Update app.py

app.py

@@ -1,6 +1,6 @@
 import torch
 import streamlit as st
-from transformers import pipeline, AutoTokenizer, AutoModelForQuestionAnswering
+from transformers import pipeline
 
 st.set_page_config(page_title="Common NLP Tasks")
 st.title("Common NLP Tasks")
@@ -14,19 +14,10 @@ option = st.sidebar.radio('', ['Extractive question answering', 'Text summarizat
 
 @st.cache(show_spinner=False, allow_output_mutation=True)
 def question_answerer(context, question):
-    tokenizer = AutoTokenizer.from_pretrained("deepset/roberta-base-squad2")
-    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
+    model_name = "deepset/roberta-base-squad2"
+    pipeline(model=model_name, tokenizer=model_name, revision="v1.0", task="question-answering")
+    result = question_answerer(question="What is extractive question answering?", context=context)
+    return result['answer']
 
 @st.cache(show_spinner=False, allow_output_mutation=True)
 def summarization_model():