Create train_model.py

train_model.py

@@ -0,0 +1,27 @@
+# train_model.py
+from transformers import AutoModelForCausalLM, TrainingArguments, Trainer
+from datasets import load_from_disk
+
+tokenized_dataset = load_from_disk("tokenized_dataset")
+
+model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.2")
+
+training_args = TrainingArguments(
+    output_dir="./checkpoints",
+    num_train_epochs=1,
+    per_device_train_batch_size=1,
+    gradient_accumulation_steps=8,
+    evaluation_strategy="no",
+    save_strategy="epoch",
+    fp16=True,  # if using GPU
+    logging_steps=50,
+)
+
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=tokenized_dataset,
+)
+
+trainer.train()
+model.save_pretrained("./my_ai_assistant", safe_serialization=True)  # saves .safetensors