Update train.py

train.py

@@ -3,6 +3,8 @@ 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")
@@ -16,7 +18,7 @@ def format_data(example):
 dataset = dataset.map(format_data)
 
 # Load the Mistral-7B model and tokenizer
-model_name = "mistralai/Mistral-7B-v0.1"  # or use Phi-2
+model_name = "mistralai/Mistral-7B-v0.1"  # You can also use Phi-2
 model = AutoModelForCausalLM.from_pretrained(model_name, load_in_8bit=True)
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
@@ -48,6 +50,15 @@ trainer = Trainer(
 # Start the fine-tuning process
 trainer.train()
 
-# Save the model
+# Save the model locally
 model.save_pretrained("./tuned_model")
 tokenizer.save_pretrained("./tuned_model")
+
+# Hugging Face Repo setup for uploading the model
+repo_name = "user06/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!")