Adding initial QA model

app.py

@@ -1,7 +1,56 @@
 import gradio as gr
+import numpy as np
+import torch
+from transformers import AutoModelForQuestionAnswering, AutoTokenizer, pipeline
 
-def greet(name):
-    return "Hello " + name + "!!"
+from transformers import DistilBertTokenizer, TFDistilBertForQuestionAnswering
+
+# Load model & tokenizer
+# model_name = "deepset/roberta-base-squad2"
+# model_name = "AmazonScience/qanlu"
+model_name = 'distilbert-base-cased-distilled-squad'
+# tokenizer = AutoTokenizer.from_pretrained("bert-large-uncased-whole-word-masking-finetuned-squad")
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+# tokenizer = DistilBertTokenizer.from_pretrained(model_name)
+# model = TFDistilBertForQuestionAnswering.from_pretrained(model_name)
+
+# model = AutoModelForQuestionAnswering.from_pretrained("bert-large-uncased-whole-word-masking-finetuned-squad")
+model = AutoModelForQuestionAnswering.from_pretrained(model_name)
+
+def QA_function(context, question):
+    inputs = tokenizer(question, context, add_special_tokens=True, return_tensors="pt") # pt for PyTorch tf for tensorflow
+    input_ids = inputs["input_ids"].tolist()[0]
+
+    outputs = model(**inputs)
+    answer_start_scores = outputs.start_logits
+    answer_end_scores = outputs.end_logits
+
+    ## following for torch
+    # Get the most likely beginning of answer with the argmax of the score
+    answer_start_index = torch.argmax(answer_start_scores)
+    # Get the most likely end of answer with the argmax of the score
+    answer_end_index = torch.argmax(answer_end_scores) + 1
+
+    answer = tokenizer.convert_tokens_to_string(tokenizer.convert_ids_to_tokens(input_ids[answer_start_index:answer_end_index]))
+    return answer
+
+# gradio_ui = gr.Interface(QA_function, [gr.inputs.Textbox(lines=7, label="Context"), gr.inputs.Textbox(label="Question")], gr.outputs.Textbox(label="Answer"))
+
+
+
+# a) Get predictions
+nlp = pipeline('question-answering', model=model_name, tokenizer=model_name)
+
+def get_answer(context,question):
+    QA_input = {
+        'question': question,
+        'context': context 
+    }
+    res = nlp(QA_input)
+    return res['answer']
+
+gradio_ui = gr.Interface(get_answer, [gr.inputs.Textbox(lines=7, label="Context"), gr.inputs.Textbox(label="Question")], gr.outputs.Textbox(label="Answer"))
+
+gradio_ui.launch()
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
-iface.launch()