train model

scripts/train_model.py

@@ -1,7 +1,16 @@
 import gc
+import sys
 
 from datasets import load_dataset, Dataset
-from transformers import AutoTokenizer
+from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer
+from transformers import AutoConfig
+from transformers import DataCollatorForLanguageModeling
+
+
+x = input('Are you sure? [y/N]')
+
+if x not in ('y', 'Y', 'yes'):
+    sys.exit(0)
 
 
 def _batch_iterator():
@@ -36,6 +45,7 @@ def _batch_iterator():
 
     del dataset
     gc.collect()
+    return
 
     # text
     dataset = load_dataset('nampdn-ai/tiny-textbooks', split='train')
@@ -180,5 +190,60 @@ tokenizer = AutoTokenizer.from_pretrained('../')
 
 dataset = Dataset.from_generator(batch_iterator)
 print(dataset)
-print(dir(dataset))
-input()
+
+
+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
+print(config)
+
+model = AutoModelForCausalLM.from_config(config)
+print(model)
+
+training_args = TrainingArguments(
+    output_dir='./results',
+    num_train_epochs=3,
+    per_device_train_batch_size=2,  # Adjust based on your GPU memory
+    per_device_eval_batch_size=2,
+    warmup_steps=500,
+    weight_decay=0.01,
+    logging_dir='./logs',
+    logging_steps=10,
+    fp16=False,
+    bf16=True,
+)
+print(training_args)
+
+data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
+print(data_collator)
+
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=tokenized_datasets['train'],
+    eval_dataset=tokenized_datasets['test'],
+    tokenizer=tokenizer,
+    data_collator=data_collator,
+)
+
+print(trainer)
+trainer.train()

scripts/train_tokenizer.py

@@ -1,4 +1,5 @@
 import gc
+import sys
 import string
 
 from datasets import load_dataset
@@ -9,6 +10,12 @@ from tokenizers.trainers import BpeTrainer
 from tokenizers.processors import TemplateProcessing
 
 
+x = input('Are you sure?')
+
+if x not in ('y', 'Y', 'yes'):
+    sys.exit(0)
+
+
 def batch_iterator():
     # code
     dataset = load_dataset('bigcode/programming-languages-keywords', split='train')