Update train.py

train.py

@@ -3,8 +3,6 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 from peft import LoraConfig, get_peft_model
 from datasets import load_dataset
 from transformers import TrainingArguments, Trainer
-from huggingface_hub import Repository
-import os
 
 # Load dataset (StackOverflow Python dataset as an example)
 dataset = load_dataset("stackoverflow", "python")
@@ -18,7 +16,7 @@ def format_data(example):
 dataset = dataset.map(format_data)
 
 # Load the Mistral-7B model and tokenizer
-model_name = "mistralai/Mistral-7B-v0.1"  # You can also use Phi-2
+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)
 
@@ -33,11 +31,11 @@ model = get_peft_model(model, lora_config)
 
 # Training arguments
 training_args = TrainingArguments(
-    output_dir="./tuned_model",
+    output_dir="./tuned_model",  # Directory to save the fine-tuned model
     per_device_train_batch_size=4,
     num_train_epochs=3,
-    save_strategy="epoch",
-    save_total_limit=2
+    save_strategy="epoch",  # Save model after each epoch
+    save_total_limit=2      # Keep only the last 2 saved models
 )
 
 # Trainer setup
@@ -50,15 +48,6 @@ trainer = Trainer(
 # Start the fine-tuning process
 trainer.train()
 
-# Save the model locally
+# Save the model and tokenizer
 model.save_pretrained("./tuned_model")
 tokenizer.save_pretrained("./tuned_model")
-
-# Hugging Face Repo setup for uploading the model
-repo_name = "krisha06/Python_tutor"  # Replace with your repo name
-repo = Repository(local_dir="./tuned_model", clone_from=repo_name)
-
-# Upload the model to Hugging Face Hub
-repo.push_to_hub()
-
-print("Model fine-tuned and uploaded to Hugging Face Hub successfully!")