import torch
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")

# 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"  # You can also 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 locally
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!")