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

repo_name = "krisha06/Python_tutor"  # Replace with your own username and repo name
repo = Repository(local_dir="./tuned_model", clone_from=repo_name)
repo.push_to_hub()

# 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

# Check if the model is uploaded to Hugging Face Hub
if os.path.exists(model_dir):
    # Load the model and tokenizer from local directory
    model = AutoModelForCausalLM.from_pretrained(model_dir, load_in_8bit=True)
    tokenizer = AutoTokenizer.from_pretrained(model_dir)
else:
    # Load the model from Hugging Face Hub if it's not found locally
    model = AutoModelForCausalLM.from_pretrained(model_name)
    tokenizer = AutoTokenizer.from_pretrained(model_name)

# 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)