Create train.py

train.py

@@ -0,0 +1,88 @@
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer
+from datasets import load_dataset
+import os
+
+# Model and tokenizer setup
+MODEL_NAME = "mistralai/Mixtral-8x7B-Instruct-v0.1"  # Real Mixtral model
+OUTPUT_DIR = "./mixtral_finetuned"
+
+# Load tokenizer
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+tokenizer.pad_token = tokenizer.eos_token  # Set pad token if missing
+
+# Load model with memory optimizations
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME,
+    torch_dtype=torch.bfloat16,        # Efficient precision
+    device_map="auto",                 # Auto-distribute across GPU/CPU
+    low_cpu_mem_usage=True             # Minimize RAM usage
+)
+
+# Load dataset (local or predefined)
+# Example: local text files; replace with your paths
+dataset = load_dataset("text", data_files={"train": "train.txt", "validation": "val.txt"})
+# Or use a Hugging Face dataset locally: dataset = load_dataset("wikitext", "wikitext-2-raw-v1")
+
+# Tokenize dataset
+def tokenize_function(examples):
+    tokenized = tokenizer(
+        examples["text"],
+        padding="max_length",
+        truncation=True,
+        max_length=512,                # Adjustable; matches earlier intent
+        return_tensors="pt"
+    )
+    tokenized["labels"] = tokenized["input_ids"].clone()  # Causal LM needs labels
+    return tokenized
+
+tokenized_datasets = dataset.map(
+    tokenize_function,
+    batched=True,
+    remove_columns=["text"]            # Save memory
+)
+
+# Split dataset
+train_dataset = tokenized_datasets["train"]
+eval_dataset = tokenized_datasets["validation"]
+
+# Define training arguments
+training_args = TrainingArguments(
+    output_dir=OUTPUT_DIR,
+    evaluation_strategy="epoch",       # Eval each epoch
+    per_device_train_batch_size=2,     # Adjust for your GPU
+    per_device_eval_batch_size=2,
+    num_train_epochs=3,                # Default; tweak as needed
+    learning_rate=2e-5,                # Safe for fine-tuning
+    weight_decay=0.01,                 # Regularization
+    gradient_accumulation_steps=4,     # Effective batch size = 8
+    bf16=True,                         # Matches bfloat16 dtype
+    fp16=False,                        # Avoid if using bf16
+    save_strategy="epoch",             # Save each epoch
+    save_total_limit=2,                # Keep 2 latest checkpoints
+    logging_dir="./logs",
+    logging_steps=10,
+    load_best_model_at_end=True,       # Load best based on eval loss
+    metric_for_best_model="loss",
+    report_to="none"                   # No external logging
+)
+
+# Initialize Trainer
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=train_dataset,
+    eval_dataset=eval_dataset,
+)
+
+# Train the model
+trainer.train()
+
+# Save locally
+trainer.save_model(OUTPUT_DIR)
+tokenizer.save_pretrained(OUTPUT_DIR)
+
+# Clean up memory
+del model
+torch.cuda.empty_cache()
+print(f"Model and tokenizer saved to {OUTPUT_DIR}")