scripts/train_model.py

import gc
import sys

from datasets import load_dataset, Dataset
from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer
from transformers import AutoConfig
from transformers import DataCollatorForLanguageModeling

import torch
from torch.utils.data import DataLoader
import torch.multiprocessing as mp


x = input('Are you sure? [y/N] ')

if x not in ('y', 'Y', 'yes'):
    sys.exit(0)


device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
mp.set_start_method('spawn', force=True)


def _batch_iterator():
    ## code
    # dataset = load_dataset('bigcode/programming-languages-keywords', split='train')

    # for row in dataset:
    #     for n in row['keywords']:
    #         yield n

    # del dataset
    # gc.collect()

    # code
    dataset = (
        load_dataset('bigcode/the-stack-smol-xs', lang, split='train', trust_remote_code=True)
        for lang in [
            'ada', 'agda', 'alloy', 'antlr', 'applescript', 'assembly', 'augeas', 'awk', 'batchfile', 'bison', 'bluespec', 'c',
            'c++', 'c-sharp', 'clojure', 'cmake', 'coffeescript', 'common-lisp', 'css', 'cuda', 'dart', 'dockerfile', 'elixir',
            'elm', 'emacs-lisp','erlang', 'f-sharp', 'fortran', 'glsl', 'go', 'groovy', 'haskell','html', 'idris', 'isabelle', 'java', 
            'java-server-pages', 'javascript', 'julia', 'kotlin', 'lean', 'literate-agda', 'literate-coffeescript', 'literate-haskell',
            'lua', 'makefile', 'maple', 'markdown', 'mathematica', 'matlab', 'ocaml', 'pascal', 'perl', 'php', 'powershell', 'prolog',
            'protocol-buffer', 'python', 'r', 'racket', 'restructuredtext', 'rmarkdown', 'ruby', 'rust', 'sas', 'scala', 'scheme', 
            'shell', 'smalltalk', 'solidity', 'sparql', 'sql', 'stan', 'standard-ml', 'stata', 'systemverilog', 'tcl', 'tcsh', 'tex', 
            'thrift', 'typescript', 'verilog', 'vhdl', 'visual-basic', 'xslt', 'yacc', 'zig'
        ]
    )

    for d in dataset:
        for row in d:
            yield row['content']

    del dataset
    gc.collect()
    return

    # text
    dataset = load_dataset('nampdn-ai/tiny-textbooks', split='train')

    for row in dataset:
        yield row['text']

    del dataset
    gc.collect()

    ## text
    # dataset = (
    #     load_dataset('wikimedia/wikisource', lang, split='train')
    #     for lang in ['20231201.ar', '20231201.as', '20231201.az', '20231201.ban', '20231201.be', '20231201.bg', '20231201.bn', '20231201.br', '20231201.bs', '20231201.ca', '20231201.cs', '20231201.cy', '20231201.da', '20231201.de', '20231201.el', '20231201.en', '20231201.eo', '20231201.es', '20231201.et', '20231201.eu', '20231201.fa', '20231201.fi', '20231201.fo', '20231201.fr', '20231201.gl', '20231201.gu', '20231201.he', '20231201.hi', '20231201.hr', '20231201.hu', '20231201.hy', '20231201.id', '20231201.is', '20231201.it', '20231201.ja', '20231201.jv', '20231201.kn', '20231201.ko', '20231201.la', '20231201.li', '20231201.lij', '20231201.lt', '20231201.mk', '20231201.ml', '20231201.mr', '20231201.nap', '20231201.nl', '20231201.no', '20231201.or', '20231201.pa', '20231201.pl', '20231201.pms', '20231201.pt', '20231201.ro', '20231201.ru', '20231201.sa', '20231201.sah', '20231201.sk', '20231201.sl', '20231201.sr', '20231201.su', '20231201.sv', '20231201.ta', '20231201.te', '20231201.th', '20231201.tr', '20231201.uk', '20231201.vec', '20231201.vi', '20231201.wa', '20231201.yi', '20231201.zh', '20231201.zh-min-nan']
    # )
    #
    # for d in dataset:
    #     for row in d['text']:
    #         yield row
    #
    # del dataset
    # gc.collect()

    # text
    dataset = (
        load_dataset('xu-song/cc100-samples', lang, split='train')
        for lang in ['am', 'ar', 'as', 'az', 'be', 'bg', 'bn', 'bn_rom', 'br', 'bs', 'ca', 'cs', 'cy', 'da', 'de', 'el', 'en', 'eo', 'es', 'et', 'eu', 'fa', 'ff', 'fi', 'fr', 'fy', 'ga', 'gd', 'gl', 'gn', 'gu', 'ha', 'he', 'hi', 'hi_rom', 'hr', 'ht', 'hu', 'hy', 'id', 'ig', 'is', 'it', 'ja', 'jv', 'ka', 'kk', 'km', 'kn', 'ko', 'ku', 'ky', 'la', 'lg', 'li', 'ln', 'lo', 'lt', 'lv', 'mg', 'mk', 'ml', 'mn', 'mr', 'ms', 'my', 'my_zaw', 'ne', 'nl', 'no', 'ns', 'om', 'or', 'pa', 'pl', 'ps', 'pt', 'qu', 'rm', 'ro', 'ru', 'sa', 'si', 'sc', 'sd', 'sk', 'sl', 'so', 'sq', 'sr', 'ss', 'su', 'sv', 'sw', 'ta', 'ta_rom', 'te', 'te_rom', 'th', 'tl', 'tn', 'tr', 'ug', 'uk', 'ur', 'ur_rom', 'uz', 'vi', 'wo', 'xh', 'yi', 'yo', 'zh-Hans', 'zh-Hant', 'zu']
    )

    for d in dataset:
        for row in d['text']:
            yield row

    del dataset
    gc.collect()

    ## text
    # dataset = (
    #     load_dataset('csebuetnlp/xlsum', lang, split='train')
    #     for lang in ['amharic', 'arabic', 'azerbaijani', 'bengali', 'burmese', 'chinese_simplified', 'chinese_traditional', 'english', 'french', 'gujarati', 'hausa', 'hindi', 'igbo', 'indonesian', 'japanese', 'kirundi', 'korean', 'kyrgyz', 'marathi', 'nepali', 'oromo', 'pashto', 'persian', 'pidgin', 'portuguese', 'punjabi', 'russian', 'scottish_gaelic', 'serbian_cyrillic', 'serbian_latin', 'sinhala', 'somali', 'spanish', 'swahili', 'tamil', 'telugu', 'thai', 'tigrinya', 'turkish', 'ukrainian', 'urdu', 'uzbek', 'vietnamese', 'welsh', 'yoruba']
    # )
    #
    # for d in dataset:
    #     for row in d['text']:
    #         yield row
    #
    # del dataset
    # gc.collect()

    ## text
    # dataset = load_dataset('recursal/SuperWikiNEXT-32B', split='train')
    #
    # for row in dataset['text']:
    #     yield row
    #
    # del dataset
    # gc.collect()

    # code
    dataset = load_dataset('m-a-p/CodeFeedback-Filtered-Instruction', split='train')

    for row in dataset:
        yield row['query'] + '\n' + row['answer']

    del dataset
    gc.collect()

    # code
    dataset = load_dataset('nampdn-ai/tiny-codes', split='train')
    
    for row in dataset:
        yield row['prompt'] + '\n' + row['response']
    
    del dataset
    gc.collect()

    # math
    dataset = load_dataset('ajibawa-2023/Maths-College', split='train')
    
    for row in dataset:
        yield row['instruction'] + '\n' + row['output']
    
    del dataset
    gc.collect()

    # math
    dataset = load_dataset('microsoft/orca-math-word-problems-200k', split='train')

    for row in dataset:
        yield row['question'] + '\n' + row['answer']

    del dataset
    gc.collect()

    # text
    dataset = load_dataset('mlabonne/FineTome-100k', split='train')

    for row in dataset['conversations']:
        yield '\n'.join(n['value'] for n in row)

    del dataset
    gc.collect()

    # instruction
    dataset = load_dataset('arcee-ai/agent-data', split='train')
    
    for row in dataset['conversations']:
        yield '\n'.join(n['value'] for n in row)
    
    del dataset
    gc.collect()

    # instruction
    dataset = (
        load_dataset('cognitivecomputations/SystemChat-2.0', data_files='SystemChat_filtered.jsonl', split='train'),
        load_dataset('cognitivecomputations/SystemChat-2.0', data_files='SystemChat_multilingual.jsonl', split='train'),
    )
    
    for d in dataset:
        for row in d['messages']:
            yield '\n'.join(n['content'] for n in row)
    
    del dataset
    gc.collect()

    # emoji
    dataset = load_dataset('badrex/llm-emoji-dataset', split='train')
    
    for row in dataset:
        yield f'{row["character"]}\n{row["unicode"]}\n{row["short description"]}\n{row["tags"]}\n{row["LLM description"]}'
    
    del dataset
    gc.collect()


def batch_iterator():
    for text in _batch_iterator():
        row = {'text': text}
        yield row

tokenizer = AutoTokenizer.from_pretrained('../')

dataset = Dataset.from_generator(batch_iterator)
print(dataset)


def tokenize_function(examples):
    outputs = tokenizer(examples['text'], truncation=True, padding='max_length', max_length=32 * 1024)
    outputs['labels'] = outputs['input_ids'].copy()
    return outputs


tokenized_datasets = dataset.map(tokenize_function, batched=True)
tokenized_datasets = tokenized_datasets.train_test_split(test_size=0.01)

config = AutoConfig.from_pretrained('mistralai/Mistral-7B-Instruct-v0.3')
config.bos_token_id = tokenizer.bos_token_id
config.eos_token_id = tokenizer.eos_token_id
config.unk_token_id = tokenizer.unk_token_id
config.pad_token_id = tokenizer.pad_token_id
config.hidden_size = 512
config.intermediate_size = int(512 * 3.5) # 1792
config.max_position_embeddings = 32 * 1024 # 32768
config.num_attention_heads = 12
config.num_hidden_layers = 10
config.num_key_value_heads = 4
config.rope_theta = 1_000_000.0
config.sliding_window = 4096
config.torch_dtype = torch.bfloat16
config.use_cache = False
print(config)

model = AutoModelForCausalLM.from_config(config)
model = model.to(torch.bfloat16)
model = torch.compile(model)
model.to(device)
print(model)

training_args = TrainingArguments(
    output_dir='./results',
    num_train_epochs=3,
    per_device_train_batch_size=1,  # Adjust based on your GPU memory
    per_device_eval_batch_size=1,
    optim='adamw_bnb_8bit',
    gradient_accumulation_steps=8,
    gradient_checkpointing=True,
    warmup_steps=500,
    weight_decay=0.01,
    logging_dir='./logs',
    logging_steps=10,
    fp16=False,
    bf16=True,
    torch_compile=True,
)
print(training_args)

data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
print(data_collator)

def collate_fn(examples):
    texts = [ex['text'] for ex in examples]
    batch = tokenizer(texts, padding=True, truncation=True, return_tensors='pt', max_length=32*1024, return_token_type_ids=False)
    batch = {k: v.to(device) for k, v in batch.items()}  # Move tensors to GPU
    batch['labels'] = batch['input_ids'].clone()
    return batch

train_dataloader = DataLoader(
    tokenized_datasets["train"],
    shuffle=True,
    collate_fn=collate_fn,
    batch_size=training_args.per_device_train_batch_size,
    pin_memory=True,
    num_workers=4
)

eval_dataloader = DataLoader(
    tokenized_datasets["test"],
    collate_fn=collate_fn,
    batch_size=training_args.per_device_eval_batch_size,
    pin_memory=True,
    num_workers=4
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_datasets['train'],
    eval_dataset=tokenized_datasets['test'],
    tokenizer=tokenizer,
    data_collator=data_collator,
)

trainer.get_train_dataloader = lambda: train_dataloader
trainer.get_eval_dataloader = lambda: eval_dataloader

print(trainer)
trainer.train()