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

# Load the model and tokenizer from your Hugging Face Hub repository
model_checkpoint = "abdulllah01/checkpoints"  # Replace with your actual checkpoint
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
model = AutoModelForQuestionAnswering.from_pretrained(model_checkpoint)

# Create a pipeline for question answering
qa_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)

# Streamlit UI setup
st.title("Question Answering App")
st.write("Enter a context and ask a question based on that context.")

# Text area for context input
context = st.text_area("Context:", "")

# Text input for the question
question = st.text_input("Question:", "")

if st.button("Get Answer"):
    if context and question:
        # Generate the answer using the pipeline
        result = qa_pipeline(question=question, context=context)
        answer = result['answer']
        st.write("**Answer:**", answer)
    else:
        st.write("Please enter both context and question.")