import streamlit as st
from transformers import BartTokenizer, BartForConditionalGeneration

# Replace with your Hugging Face model repository path for QnA
model_repo_path_qna = 'ASaboor/Bart_Therapy'

# Load the model and tokenizer for QnA
model_qna = BartForConditionalGeneration.from_pretrained(model_repo_path_qna)
tokenizer_qna = BartTokenizer.from_pretrained(model_repo_path_qna)

# Streamlit app layout
st.set_page_config(page_title="QnA App", page_icon=":memo:", layout="wide")

st.title("Question and Answer App")
st.write("""
    This app uses a fine-tuned BART model to answer questions.
    Enter your question below and click "Get Answer" to see the result.
""")

# User input for QnA
question_input = st.text_input("Enter question", placeholder="Type your question here...")

# Generate the answer
if st.button("Get Answer"):
    if question_input:
        with st.spinner("Generating answer..."):
            try:
                # Tokenize input
                inputs = tokenizer_qna(question_input, return_tensors='pt', max_length=512, truncation=True)
                
                # Generate answer
                outputs = model_qna.generate(input_ids=inputs['input_ids'], attention_mask=inputs['attention_mask'], max_length=150, num_beams=5, early_stopping=True)
                
                # Decode the answer
                answer = tokenizer_qna.decode(outputs[0], skip_special_tokens=True)
                
                # Display answer
                st.subheader("Answer")
                st.write(answer)
            except Exception as e:
                st.error(f"An error occurred during QnA: {e}")
    else:
        st.warning("Please enter a question for QnA.")

# Optional: Add a footer or additional information
st.markdown("""
    ---
    Made with ❤️ using [Streamlit](https://streamlit.io) and [Hugging Face Transformers](https://huggingface.co/transformers/).
""")