Spaces:
Sleeping
Sleeping
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() | |