Update app.py

app.py

@@ -1,42 +1,64 @@
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForQuestionAnswering
+from transformers import AutoTokenizer, AutoModelForQuestionAnswering, AutoModelForCausalLM
 import torch
 
-# Load Bio_ClinicalBERT model and tokenizer
-model_name = "emilyalsentzer/Bio_ClinicalBERT"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForQuestionAnswering.from_pretrained(model_name)
+# Load BioBERT for medical question answering
+medical_model_name = "dmis-lab/biobert-base-cased-v1.1"
+medical_tokenizer = AutoTokenizer.from_pretrained(medical_model_name)
+medical_model = AutoModelForQuestionAnswering.from_pretrained(medical_model_name)
+
+# Load DialoGPT for conversation
+conversation_model_name = "microsoft/DialoGPT-small"
+conversation_tokenizer = AutoTokenizer.from_pretrained(conversation_model_name)
+conversation_model = AutoModelForCausalLM.from_pretrained(conversation_model_name)
 
 # Streamlit app layout
 st.title("Medical Chatbot")
-st.write("Enter your medical query below:")
+st.write("Ask your medical-related question below:")
+
+# Conversation history tracker
+if 'conversation_history' not in st.session_state:
+    st.session_state.conversation_history = []
 
 # Text input
-user_input = st.text_area("Input Text")
+user_input = st.text_input("You:")
 
-if st.button("Get Response"):
+if st.button("Send"):
     if user_input:
-        # Example context for medical queries
-        context = """
-        A fever is a temporary increase in body temperature, often due to an illness.
-        It's a sign that something unusual is going on in your body.
-        For adults, a fever may be uncomfortable, but usually isn’t a cause for concern unless it reaches 103 F (39.4 C) or higher.
-        """  # This context should be based on your use case
-        
-        # Tokenize input
-        inputs = tokenizer(user_input, context, return_tensors="pt")
+        # Append user input to the conversation history
+        st.session_state.conversation_history.append(f"User: {user_input}")
 
-        # Get model outputs
+        # BioBERT for Medical Q&A
+        context = "Medical knowledge base or a long text about medical conditions..."  # Insert medical reference or context
+        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 = model(**inputs)
+            outputs = medical_model(**inputs)
+            answer_start_scores = outputs.start_logits
+            answer_end_scores = outputs.end_logits
 
-        # Get the most likely beginning and end of the answer
-        answer_start = torch.argmax(outputs.start_logits)
-        answer_end = torch.argmax(outputs.end_logits) + 1
+        answer_start = torch.argmax(answer_start_scores)
+        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]))
+
+        # 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)
 
-        # Convert tokens to string
-        answer = tokenizer.convert_tokens_to_string(tokenizer.convert_ids_to_tokens(inputs["input_ids"][0][answer_start:answer_end]))
+        # Append conversational response to conversation history
+        st.session_state.conversation_history.append(f"Bot (Conversational): {conversation_response}")
 
-        st.write(f"Bot: {answer}")
+        # Display conversation history
+        for message in st.session_state.conversation_history:
+            st.write(message)
     else:
-        st.write("Please enter a medical query.")
+        st.write("Please enter a medical question.")