Update app.py

app.py

@@ -1,3 +1,4 @@
+import torch
 import streamlit as st
 from transformers import pipeline, AutoTokenizer, AutoModelForQuestionAnswering
 
@@ -12,10 +13,20 @@ st.sidebar.header('What will you like to do?')
 option = st.sidebar.radio('', ['Extractive question answering', 'Text summarization', 'Text generation', 'Sentiment analysis'])
 
 @st.cache(show_spinner=False, allow_output_mutation=True)
-def question_model():
+def question_answerer(context, question):
     tokenizer = AutoTokenizer.from_pretrained("deepset/roberta-base-squad2")
-    question_answerer = AutoModelForQuestionAnswering.from_pretrained("deepset/roberta-base-squad2")
-    return question_answerer
+    model = AutoModelForQuestionAnswering.from_pretrained("deepset/roberta-base-squad2")
+    inputs = tokenizer(question, context, add_special_tokens=True, return_tensors="pt")
+    input_ids = inputs["input_ids"].tolist()[0]
+    outputs = model(**inputs)
+    answer_start_scores = outputs.start_logits
+    answer_end_scores = outputs.end_logits
+    # Get the most likely beginning of answer with the argmax of the score
+    answer_start = torch.argmax(answer_start_scores)
+    # Get the most likely end of answer with the argmax of the score 
+    answer_end = torch.argmax(answer_end_scores) + 1
+    answer = tokenizer.convert_tokens_to_string(tokenizer.convert_ids_to_tokens(input_ids[answer_start:answer_end]))
+    return answer
 
 @st.cache(show_spinner=False, allow_output_mutation=True)
 def summarization_model():
@@ -41,10 +52,9 @@ if option == 'Extractive question answering':
         question = st.text_input(label='Enter your question')
         button = st.button('Get answer')
         if button:
-            question_answerer = question_model()
             with st.spinner(text="Getting answer..."):
-                answer = question_answerer(context=context, question=question)
-                st.write(answer["answer"])
+                answer = question_answerer(context, question)
+                st.write(answer)
     elif source == "I want to upload a file":
         uploaded_file = st.file_uploader("Choose a .txt file to upload", type=["txt"])
         question = st.text_input(label='Enter your question')