Update app.py

app.py

@@ -1,38 +1,37 @@
 import streamlit as st
-from transformers import GPT2LMHeadModel, GPT2Tokenizer
+from transformers import BigBirdForQuestionAnswering, BigBirdTokenizer
 
-# Load the Mistral model for text generation
-model_name = "salesforce/gpt-3.5-mistral-small"
-tokenizer = GPT2Tokenizer.from_pretrained(model_name)
-model = GPT2LMHeadModel.from_pretrained(model_name)
+# 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("Cherry Bot")
-    st.markdown("*Your Companion for Suicide Prevention*")
-
-    # Add a sidebar to the app
-    with st.sidebar:
-        st.markdown("# Text Generation Settings")
-        # Slider for adjusting the maximum length of generated text
-        max_length = st.slider("Max Length of Generated Text:", min_value=10, max_value=200, value=50, step=10)
-
-    # Text input for user to input starting text for generation
-    starting_text = st.text_area("Enter starting text for generation:", height=100)
-
-    # Perform text generation when the user clicks the button
-    if st.button("Generate Text"):
-        if starting_text:
-            # Tokenize input text
-            input_ids = tokenizer.encode(starting_text, return_tensors="pt")
-
-            # Generate text using the loaded model
-            generated_ids = model.generate(input_ids, max_length=max_length, num_return_sequences=1)
-            generated_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
-
-            # Display the generated text
-            st.write("Generated Text:")
-            st.write(generated_text)
+    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()