Update app.py

app.py

@@ -1,24 +1,25 @@
 import time
 import streamlit as st
 import requests
-headers = {"Authorization": "Bearer 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"}
 
+headers = {"Authorization": "Bearer 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"}
 
-st.title("question-answering-demo")
-option = st.selectbox('Select a model', ('distilbert-base-cased-distilled-squad', 'roberta-base-squad2-distilled', 'xlm-roberta-large-squad2', 'bert-large-cased-whole-word-masking-finetuned-squad'))
-if option == "distilbert-base-cased-distilled-squad":
-    API_URL = 'https://api-inference.huggingface.co/models/distilbert-base-cased-distilled-squad'
-elif option == 'roberta-base-squad2-distilled':
-    API_URL = 'https://api-inference.huggingface.co/models/deepset/roberta-base-squad2-distilled'
-elif option == 'xlm-roberta-large-squad2':
-    API_URL = "https://api-inference.huggingface.co/models/deepset/xlm-roberta-large-squad2"
-if option == 'bert-large-cased-whole-word-masking-finetuned-squad':
-    API_URL = 'https://api-inference.huggingface.co/models/bert-large-cased-whole-word-masking-finetuned-squad'
-st.write(API_URL)
+st.title("extractive-qa")
+option = st.selectbox(
+    'Select a model👇',
+    ('distilbert-base-cased-distilled-squad', 'roberta-base-squad2-distilled', 'xlm-roberta-large-squad2', 'bert-large-cased-whole-word-masking-finetuned-squad'))
 text_input = st.text_area("Enter some context👇")
 text_question = st.text_input("Enter a question regarding that context👇")
 
-
+if option == 'distilbert-base-cased-distilled-squad':
+        API_URL = "https://api-inference.huggingface.co/models/distilbert-base-cased-distilled-squad"
+elif option == "roberta-base-squad2-distilled":
+    API_URL="https://api-inference.huggingface.co/models/deepset/roberta-base-squad2-distilled"
+elif option == 'xlm-roberta-large-squad2':
+    API_URL="https://api-inference.huggingface.co/models/deepset/xlm-roberta-large-squad2"
+if option == "bert-large-cased-whole-word-masking-finetuned-squad":
+    API_URL = "https://api-inference.huggingface.co/models/bert-large-cased-whole-word-masking-finetuned-squad"
+    
 def query(payload):
     retries = 0
     while True:
@@ -31,10 +32,13 @@ def query(payload):
             print(f"Too many requests. Retrying in {wait_time} seconds...")
             time.sleep(wait_time)
         else:
-            print(f"Request failed with status code {response.status_code}.")
+            st.error(f"Request failed with status code {response.status_code}. Try Again!", icon="🚨")
             return None
 
 
 if st.button("Send"):
     output = query({"inputs": {"question": text_question, "context": text_input}})
-    st.write(output)
+    if output:
+        answer = output["answer"]
+        st.write(output)
+