train model

scripts/requirements.in

@@ -3,4 +3,5 @@ datasets
 jinja2
 transformers
 jsonlines
-litgpt[all]
+litgpt[all]
+litdata

scripts/train_model.py

@@ -1,17 +1,31 @@
 import gc
+
+import torch
+from torch.optim import AdamW
+import bitsandbytes as bnb
 from datasets import load_dataset, Dataset
 
+from transformers import (
+    AutoConfig,
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    TrainingArguments,
+    Trainer,
+    DataCollatorForLanguageModeling,
+)
 
-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()
+
+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()
+    return
 
     # code
     dataset = (
@@ -166,4 +180,79 @@ def batch_iterator():
         yield f'{row["character"]}\n{row["unicode"]}\n{row["short description"]}\n{row["tags"]}\n{row["LLM description"]}'
     
     del dataset
-    gc.collect()
+    gc.collect()
+
+
+def batch_iterator():
+    for text in _batch_iterator():
+        for i in range(0, len(text), 2048):
+            chunk = text[i:i + 2048]
+            yield {'text': chunk}
+
+
+tokenizer = AutoTokenizer.from_pretrained('../')
+print(tokenizer)
+
+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 = 1792 # int(512 * 3.5)
+config.max_position_embeddings = 32768 # 32 * 1024
+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)
+print(model)
+
+dataset = Dataset.from_generator(batch_iterator)
+print(dataset)
+
+data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
+print(data_collator)
+
+optimizer = bnb.optim.AdamW8bit(
+    model.parameters(),
+    lr=1e-5,
+    betas=(0.9, 0.95),
+    weight_decay=0.1,
+)
+print(optimizer)
+
+training_args = TrainingArguments(
+    output_dir='./mistral-custom',
+    num_train_epochs=3,
+    per_device_train_batch_size=1,
+    gradient_accumulation_steps=8,
+    warmup_steps=500,
+    learning_rate=1e-5,
+    fp16=False,
+    bf16=True,
+    logging_dir='./logs',
+    logging_steps=10,
+    evaluation_strategy='no',
+    save_strategy='epoch',
+    torch_compile=True,
+    remove_unused_columns=False,
+)
+print(training_args)
+
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=dataset,
+    data_collator=data_collator,
+    optimizers=(optimizer, None)
+)
+print(trainer)
+
+trainer.train()
+trainer.save_model('./mistral-custom-final')