Update app.py

app.py

@@ -1,28 +1,36 @@
 
 from transformers import AutoTokenizer, AutoModelForQuestionAnswering
 import torch
+import gradio as gr
 
 tokenizer = AutoTokenizer.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")
 model = AutoModelForQuestionAnswering.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")
 
 def answer_question(context, question):
     inputs = tokenizer.encode_plus(question, context, return_tensors="pt")
-    answer_start_scores, answer_end_scores = model(**inputs).values()
+    outputs = model(**inputs)
+    start_scores = outputs.start_logits
+    end_scores = outputs.end_logits
+
+    start = torch.argmax(start_scores)
+    end = torch.argmax(end_scores) + 1
 
-    start = torch.argmax(answer_start_scores)
-    end = torch.argmax(answer_end_scores) + 1
+    if start >= end:
+        return "I couldn't find an answer."
 
     answer = tokenizer.convert_tokens_to_string(
         tokenizer.convert_ids_to_tokens(inputs["input_ids"][0][start:end])
     )
     return answer
-import gradio as gr
+
 
 def chatbot_response(question):
-    context = "Include a reliable medical knowledge base here like PubMed abstracts or simplified texts."
+    context = (
+        "COVID-19 is a respiratory illness caused by the SARS-CoV-2 virus. "
+        "Common symptoms include fever, cough, fatigue, and loss of taste or smell. "
+        "Fever usually lasts for 3-5 days. Treatment is mostly supportive, and vaccination reduces severity."
+    )
     return answer_question(context, question)
 
 iface = gr.Interface(fn=chatbot_response, inputs="text", outputs="text", title="Medical Chatbot")
 iface.launch()
-
-