import streamlit as st
from transformers import AutoModelForCausalLM, AutoTokenizer
from huggingface_hub import Repository
import os

# Define model directory and Hugging Face repo name
model_dir = "./tuned_model"  # Directory where the fine-tuned model is saved
repo_name = "krisha06/Python_tutor"  # Replace with your Hugging Face repo name

# Streamlit App Title
st.title("AI Coding Mentor")

# Section to upload the fine-tuned model to Hugging Face Hub
st.header("Upload Your Fine-Tuned Model to Hugging Face Hub")

# Option to upload the model
upload_model_button = st.button("Upload Model to Hugging Face Hub")

if upload_model_button:
    if os.path.exists(model_dir):
        # Initialize the Hugging Face Repository and push model to the Hub
        repo = Repository(local_dir=model_dir, clone_from=repo_name)
        repo.push_to_hub()
        st.success("Model uploaded to Hugging Face Hub successfully!")
    else:
        st.error("Model directory does not exist. Please make sure the model is fine-tuned first.")

# Section for using the model in the app
st.header("Ask Me Any Coding Question!")

# Load model and tokenizer (either from local directory or Hugging Face Hub)
model_name = repo_name  # Use the repo name if the model is on Hugging Face Hub, else use local dir
token = "<HUGGINGFACE_TOKEN>"  # Replace with your Hugging Face token if the model is private

if os.path.exists(model_dir):
    # Load the model and tokenizer from the local directory
    model = AutoModelForCausalLM.from_pretrained(model_dir, use_auth_token=token)
    tokenizer = AutoTokenizer.from_pretrained(model_dir, use_auth_token=token)
else:
    # Load the model from Hugging Face Hub
    model = AutoModelForCausalLM.from_pretrained(model_name, use_auth_token=token)
    tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=token)

# User input: Coding question
question = st.text_input("Enter your coding question:")

if question:
    input_text = f"### Question:\n{question}\n### Answer:"
    inputs = tokenizer(input_text, return_tensors="pt")

    with st.spinner("Processing..."):
        output = model.generate(**inputs, max_length=200, num_return_sequences=1)
        answer = tokenizer.decode(output[0], skip_special_tokens=True)

    st.write(answer)