import streamlit as st
from transformers import BigBirdForQuestionAnswering, BigBirdTokenizer

# Load the BigBird model and tokenizer
model_name = "google/bigbird-base-trivia-itc"
model = BigBirdForQuestionAnswering.from_pretrained(model_name)
tokenizer = BigBirdTokenizer.from_pretrained(model_name)

# Streamlit app
def main():
    st.title("CherryBot")

    # Text input for user to input the question
    question = st.text_input("Enter your question:")

    # Text area for user to input the context
    context = st.text_area("Enter the context for answering:", height=200)

    # Perform question answering when the user clicks the button
    if st.button("Get Answer"):
        if question and context:
            # Tokenize input question and context
            encoded_input = tokenizer(question, context, return_tensors='pt')

            # Perform question answering using the loaded model
            output = model(**encoded_input)

            # Extract and display the answer
            answer_start = output.start_logits.argmax()
            answer_end = output.end_logits.argmax() + 1
            answer = tokenizer.convert_tokens_to_string(tokenizer.convert_ids_to_tokens(encoded_input.input_ids[0][answer_start:answer_end]))
            
            st.write("Answer:")
            st.write(answer)

if __name__ == "__main__":
    main()