import streamlit as st
import torch
from transformers import BartForConditionalGeneration, BartTokenizer

# Load the model and tokenizer
model_repo_path = 'AbdurRehman313/hotpotQA_BART_Finetuned_E5'
model = BartForConditionalGeneration.from_pretrained(model_repo_path)
tokenizer = BartTokenizer.from_pretrained(model_repo_path)

# Ensure the model is in evaluation mode
model.eval()
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

# Streamlit app layout
st.title("Multi-Hop Question Answering Application")

# User input for context and question
context_input = st.text_area("Enter context", height=200)
question_input = st.text_area("Enter question")

# Generate the answer
if st.button("Get Answer"):
    if context_input and question_input:
        with st.spinner("Generating answer..."):
            try:
                # Prepare the input for the model
                input_text = f"context: {context_input} question: {question_input}"
                inputs = tokenizer(input_text, return_tensors='pt')
                inputs = {key: value.to(device) for key, value in inputs.items()}

                # Perform inference
                with torch.no_grad():
                    outputs = model.generate(inputs['input_ids'], max_length=50)

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

                st.subheader("Answer")
                st.write(answer)
            except Exception as e:
                st.error(f"Error during question answering: {e}")
    else:
        st.warning("Please enter both context and question.")