import streamlit as st
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
from dotenv import load_dotenv
load_dotenv()
from huggingface_hub import login
import os

access_token = os.environ["HUGGING_FACE_HUB_TOKEN"]
login(token=access_token)

# Load the base model and adapter
@st.cache_resource
def load_model():
    base_model_name = "Qwen/Qwen2.5-Coder-3B-Instruct"  # Ensure this is the correct base model
    adapter_model_name = "mohamedyd/Natural-Coder-3B-Instruct-V1"

    # Load the base model
    base_model = AutoModelForCausalLM.from_pretrained(base_model_name, trust_remote_code=True)
    
    # Load the PEFT adapter on top of the base model
    model = PeftModel.from_pretrained(base_model, adapter_model_name)

    # Load tokenizer from base model
    tokenizer = AutoTokenizer.from_pretrained(base_model_name, trust_remote_code=True)

    return model, tokenizer

model, tokenizer = load_model()

# Streamlit App UI
st.title("Natural-Coder-3B-Instruct-V1 Model Interaction")

user_input = st.text_area("Enter your prompt here:", height=150)

if st.button("Generate Response"):
    if user_input:
        inputs = tokenizer(user_input, return_tensors="pt")
        outputs = model.generate(**inputs, max_length=512, num_return_sequences=1)
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        st.write("Model Response:")
        st.write(response)
    else:
        st.write("Please enter a prompt to generate a response.")