Create train.py

train.py

@@ -0,0 +1,55 @@
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import LoraConfig, get_peft_model
+from datasets import load_dataset
+from transformers import TrainingArguments, Trainer
+
+# Load dataset (StackOverflow Python dataset as an example)
+dataset = load_dataset("stackoverflow", "python")
+
+# Preprocess the dataset
+def format_data(example):
+    return {
+        "text": f"### Question:\n{example['question']}\n### Answer:\n{example['answer']}"
+    }
+
+dataset = dataset.map(format_data)
+
+# Load the Mistral-7B model and tokenizer
+model_name = "mistralai/Mistral-7B-v0.1"  # or use Phi-2
+model = AutoModelForCausalLM.from_pretrained(model_name, load_in_8bit=True)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+# LoRA configuration for lightweight fine-tuning
+lora_config = LoraConfig(
+    r=8,
+    lora_alpha=32,
+    lora_dropout=0.1,
+    target_modules=["q_proj", "v_proj"]
+)
+model = get_peft_model(model, lora_config)
+
+# Training arguments
+training_args = TrainingArguments(
+    output_dir="./tuned_model",
+    per_device_train_batch_size=4,
+    num_train_epochs=3,
+    save_strategy="epoch",
+    save_total_limit=2
+)
+
+# Trainer setup
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=dataset["train"]
+)
+
+# Start the fine-tuning process
+trainer.train()
+
+# Save the model
+model.save_pretrained("./tuned_model")
+tokenizer.save_pretrained("./tuned_model")
+
+