Update app.py

app.py

@@ -29,10 +29,13 @@ if st.button("Send"):
         st.session_state.conversation_history.append(f"User: {user_input}")
 
         # BioBERT for Medical Q&A
-        context = "Medical knowledge base or a long text about medical conditions..."  # Insert medical reference or context
+        context = """
+        A headache can have many causes, ranging from stress, dehydration, or fatigue to more severe conditions like migraines, infections, or neurological problems. Common remedies include over-the-counter pain relievers, hydration, and rest. 
+        If headaches are persistent or severe, it may indicate an underlying condition such as tension headaches, cluster headaches, or even infections like sinusitis. If the headache is accompanied by other symptoms such as nausea, vision changes, or confusion, it is recommended to seek medical attention.
+        """
         inputs = medical_tokenizer.encode_plus(user_input, context, add_special_tokens=True, return_tensors="pt")
         input_ids = inputs["input_ids"].tolist()[0]
-        
+
         # Perform Question Answering using BioBERT
         with torch.no_grad():
             outputs = medical_model(**inputs)
@@ -43,13 +46,17 @@ if st.button("Send"):
         answer_end = torch.argmax(answer_end_scores) + 1
         medical_answer = medical_tokenizer.convert_tokens_to_string(medical_tokenizer.convert_ids_to_tokens(input_ids[answer_start:answer_end]))
 
+        # If no medical answer is found, provide a fallback response
+        if medical_answer.strip() == "":
+            medical_answer = "I'm not sure about that. You may want to consult a medical professional."
+
         # Append medical response to the conversation history
         st.session_state.conversation_history.append(f"Bot (Medical): {medical_answer}")
 
         # DialoGPT for conversational response
         conversation_input_ids = conversation_tokenizer.encode(user_input + conversation_tokenizer.eos_token, return_tensors='pt')
         conversation_bot_input_ids = torch.cat([conversation_tokenizer.encode(convo + conversation_tokenizer.eos_token, return_tensors='pt') for convo in st.session_state.conversation_history], dim=-1)
-        
+
         # Generate conversational response
         chat_history_ids = conversation_model.generate(conversation_bot_input_ids, max_length=1000, pad_token_id=conversation_tokenizer.eos_token_id)
         conversation_response = conversation_tokenizer.decode(chat_history_ids[:, conversation_bot_input_ids.shape[-1]:][0], skip_special_tokens=True)