import streamlit as st
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# Load the fine-tuned model
model_name = "./tuned_model"  # Load from Hugging Face or locally
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Set up Streamlit UI
st.title("AI Coding Mentor")
st.write("Ask me any programming-related question!")

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

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

    # Generate the answer using the fine-tuned model
    output = model.generate(**inputs, max_length=200, num_return_sequences=1)

    # Decode the output
    answer = tokenizer.decode(output[0], skip_special_tokens=True)

    # Display the result
    st.write(answer)