Update app.py

app.py

@@ -1,6 +1,40 @@
 import gradio as gr
+from gradio.mix import Parallel
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
-iface = gr.Interface.load("huggingface/csarron/bert-base-uncased-squad-v1",  
-)
+import torch
 
-iface.launch()
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+tokenizer = AutoTokenizer.from_pretrained('jaimin/QA')
+model = AutoModelForSeq2SeqLM.from_pretrained('jaimin/QA').to('cuda:0')
+
+def get_answer(question,context,device="cpu"):
+
+    question_doc = "question: {} context: {}".format(question, context)
+    
+    encoding = tokenizer.encode_plus(question_doc, pad_to_max_length=True, return_tensors="pt")
+    input_ids, attention_masks = encoding["input_ids"].to(device), encoding["attention_mask"].to(device)
+
+    model_output = model.generate(
+        input_ids=input_ids,
+        attention_mask=attention_masks,
+        num_beams=10,
+        do_sample=True,
+        max_length=30,
+        top_k=50,
+        top_p=0.95,
+        early_stopping=True,
+        no_repeat_ngram_size=1
+    )
+    generated_sent = tokenizer.decode(
+            model_output, skip_special_tokens=True, clean_up_tokenization_spaces=True)
+        
+    return generated_sent
+    
+demo = gr.Interface(get_answer, inputs=[gr.inputs.Textbox(label="question", optional=False),
+                                        gr.inputs.Textbox(label="context", optional=False],
+                                outputs=[gr.outputs.Textbox(label="Answer")])
+                                                 
+if __name__ == "__main__":
+    demo.launch()