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

# Load the fine-tuned model and tokenizer
@st.cache_resource  # Cache model to avoid reloading
def load_model():
    model_directory = "C:/Users/DELL/Desktop/QC_streamlit/QC_fine_tuned_smollm2_360m_instruct_3_epoch"
    model = AutoModelForCausalLM.from_pretrained(model_directory)
    tokenizer = AutoTokenizer.from_pretrained(model_directory)
    return model, tokenizer

# Load model and tokenizer
model, tokenizer = load_model()

# Create a pipeline
question_completion_pipeline = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    device=-1   
)

# Streamlit UI
st.title("Question Completion Model")
st.write("Provide a partial question, and the model will complete it.")

partial_question = st.text_input("Enter a partial question:", "")

if st.button("Complete Question"):
    if partial_question.strip():
        output = question_completion_pipeline(
            partial_question,
            max_length=60,
            num_return_sequences=1,
            do_sample=True
        )
       
        completed_question = output[0]["generated_text"]
        st.success(f"Completed Question: {completed_question}")    
        
    else:
        st.warning("Please enter a partial question.")