import torch
import streamlit as st
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer
from peft import LoraConfig, get_peft_model, TaskType
from huggingface_hub import HfApi, Repository
from datasets import load_dataset

# Hugging Face details
HF_TOKEN = "your_huggingface_token"
REPO_NAME = "tinyllama-lora-finetuned"

# Initialize Streamlit
st.title("🧑‍🏫 Python Tutor AI (Fine-tuned with LoRA)")

# Create HF repo if it doesn't exist
api = HfApi()
api.create_repo(REPO_NAME, token=HF_TOKEN, repo_type="model", exist_ok=True)
repo = Repository(local_dir=REPO_NAME, clone_from=f"hf://{REPO_NAME}", use_auth_token=HF_TOKEN)

# Load TinyLlama Model & Tokenizer
MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float16,
    device_map="auto"
)

# LoRA Configuration
lora_config = LoraConfig(
    task_type=TaskType.CAUSAL_LM,
    inference_mode=False,
    r=8,
    lora_alpha=32,
    lora_dropout=0.1
)

# Apply LoRA to Model
model = get_peft_model(model, lora_config)

# Load dataset for fine-tuning
dataset = load_dataset("Abirate/english_python_code_instructions", split="train[:2%]")

# Fine-Tuning Parameters
training_args = TrainingArguments(
    output_dir="./results",
    per_device_train_batch_size=1,
    gradient_accumulation_steps=4,
    optim="adamw_torch",
    num_train_epochs=1,
    logging_steps=10,
    save_strategy="no"
)

# Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset
)

# Fine-tune Model
st.write("🎯 Fine-tuning Model (LoRA)...")
trainer.train()
st.success("✅ Fine-tuning complete!")

# Push model to Hugging Face
model.push_to_hub(REPO_NAME, use_auth_token=HF_TOKEN)
tokenizer.push_to_hub(REPO_NAME, use_auth_token=HF_TOKEN)
st.success("🚀 Model pushed to Hugging Face!")

# User Input for Python Tutoring
user_input = st.text_area("📝 Ask me a Python question:")
if st.button("Get Answer"):
    if user_input:
        inputs = tokenizer(user_input, return_tensors="pt").to("cuda")
        outputs = model.generate(**inputs, max_new_tokens=100)
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        st.write("💡 AI Tutor:", response)
    else:
        st.warning("⚠️ Please enter a question.")