Update chatbot.py

chatbot.py

@@ -1,164 +1,201 @@
-import streamlit as st
-from langchain_core.prompts import PromptTemplate
-from langchain_core.output_parsers import StrOutputParser
-from transformers import pipeline
-from langchain_huggingface import HuggingFaceEndpoint
-import numpy as np
-from pydub import AudioSegment
-import os
-
-# Model IDs
-model_id = "unsloth/Llama-3.2-1B-Instruct"
-model2_id = "mistralai/Mistral-7B-Instruct-v0.3"
-whisper_model = "openai/whisper-small"  # Using Whisper model for audio transcription
-
-def get_llm_hf_inference(model_id, max_new_tokens=128, temperature=0.1):
-    """Returns a language model for HuggingFace inference."""
-    try:
-        llm = HuggingFaceEndpoint(
-            repo_id=model_id,
-            max_new_tokens=max_new_tokens,
-            temperature=temperature,
-            token=os.getenv("HF_TOKEN")
-        )
-        return llm
-    except Exception as e:
-        st.error(f"Error initializing model: {e}")
-        return None
- 
-# Initialize Whisper transcription model
-def load_transcription_model():
-    try:
-        transcriber = pipeline("automatic-speech-recognition", model=whisper_model)
-        return transcriber
-    except Exception as e:
-        st.error(f"Error loading Whisper model: {e}")
-        return None
- 
-# Preprocess audio to 16kHz mono
-def preprocess_audio(file):
-    audio = AudioSegment.from_file(file).set_frame_rate(16000).set_channels(1)
-    audio_samples = np.array(audio.get_array_of_samples()).astype(np.float32) / (2**15)
-    return audio_samples
- 
-# Function to transcribe audio with preprocessing
-def transcribe_audio(file, transcriber):
-    audio = preprocess_audio(file)
-    transcription = transcriber(audio)["text"]
-    return transcription
-
-# Chatbot page content
-def display_chatbot():
-    st.title("Personal Psychologist Chatbot")
-    st.markdown(f"*This is a simple chatbot that acts as a psychologist and gives solutions to your psychological problems. It uses the {model_id}.*")
- 
-    # Sidebar for settings
-    with st.sidebar:
-        # Reset Chat History
-        reset_history = st.button("Reset Chat History")
-        go_home = st.button("Back to Home")
-        if go_home:
-            st.session_state.page = "home"
-            st.experimental_rerun()  # Go back to homepage
- 
-    # Initialize or reset chat history
-    if reset_history:
-        st.session_state.chat_history = [{"role": "assistant", "content": st.session_state.starter_message}]
- 
-    def get_response(system_message, chat_history, user_text, model_id, max_new_tokens=256):
-        """Generates a response from the chatbot model."""
-        hf = get_llm_hf_inference(model_id=model_id, max_new_tokens=max_new_tokens)
-        if hf is None:
-            return "Error: Model not initialized.", chat_history
- 
-        # Create the prompt template
-        prompt = PromptTemplate.from_template(
-            (
-                "[INST] {system_message}"
-                "\nCurrent Conversation:\n{chat_history}\n\n"
-                "\nUser: {user_text}.\n [/INST]"
-                "\nAI:"
-            )
-        )
-        chat = prompt | hf.bind(skip_prompt=True) | StrOutputParser(output_key='content')
- 
-        # Generate the response
-        response = chat.invoke(input=dict(system_message=system_message, user_text=user_text, chat_history=chat_history))
-        response = response.split("AI:")[-1].strip()
- 
-        # Check for "thank you" in the user's input and prompt for report
-        if "thank you" in user_text.lower():
-            follow_up_question = "Would you like to have a report of your current health?"
-            response += f"\n\n{follow_up_question}"
- 
-        # Update the chat history
-        chat_history.append({'role': 'user', 'content': user_text})
-        chat_history.append({'role': 'assistant', 'content': response})
-        return response, chat_history
- 
-    def get_summary_of_chat_history(chat_history, model2_id):
-        """Generates a summary of the entire chat history using GPT-2."""
-        hf = get_llm_hf_inference(model_id=model2_id, max_new_tokens=256)
-        if hf is None:
-            return "Error: Model not initialized."
- 
-        # Create the summary prompt
-        chat_content = "\n".join([f"{message['role']}: {message['content']}" for message in chat_history])
-        prompt = PromptTemplate.from_template(
-            f"Please summarize the following conversation and generate a report of the user's current health:\n\n{chat_content}"
-        )
- 
-        summary = prompt | hf.bind(skip_prompt=True) | StrOutputParser(output_key='content')
-        summary_response = summary.invoke(input={})
- 
-        return summary_response
- 
-    # Load Whisper model for transcription
-    transcriber = load_transcription_model()
- 
-    # User input for audio and text
-    st.markdown("### Choose your input:")
-    audio_file = st.file_uploader("Upload an audio file for transcription", type=["mp3", "wav", "m4a"])
-    st.session_state.user_text = st.chat_input(placeholder="Or enter your text here.")
- 
-    # Check if audio file is uploaded and transcribe if available
-    if audio_file is not None and transcriber:
-        with st.spinner("Transcribing audio..."):
-            try:
-                st.session_state.user_text = transcribe_audio(audio_file, transcriber)
-                st.success("Audio transcribed successfully!")
-            except Exception as e:
-                st.error(f"Error transcribing audio: {e}")
- 
-    # Chat interface
-    output_container = st.container()
- 
-    # Display chat messages
-    with output_container:
-        for message in st.session_state.chat_history:
-            if message['role'] == 'system':
-                continue
-            with st.chat_message(message['role'], avatar=st.session_state['avatars'][message['role']]):
-                st.markdown(message['content'])
- 
-    # Process text input for chatbot response
-    if st.session_state.user_text:
-        with st.chat_message("user", avatar=st.session_state.avatars['user']):
-            st.markdown(st.session_state.user_text)
- 
-        with st.chat_message("assistant", avatar=st.session_state.avatars['assistant']):
-            with st.spinner("Addressing your concerns..."):
-                response, st.session_state.chat_history = get_response(
-                    system_message=st.session_state.system_message,
-                    user_text=st.session_state.user_text,
-                    chat_history=st.session_state.chat_history,
-                    model_id=model_id,
-                    max_new_tokens=st.session_state.max_response_length,
-                )
-                st.markdown(response)
- 
-                # Check if the user has agreed to the report
-                if "yes" in st.session_state.user_text.lower() and "Would you like to have a report of your current health?" in response:
-                    with st.spinner("Generating your health report..."):
-                        report = get_summary_of_chat_history(st.session_state.chat_history, model2_id)
+import streamlit as st
+from langchain_core.prompts import PromptTemplate
+from langchain_core.output_parsers import StrOutputParser
+from transformers import pipeline
+from langchain_huggingface import HuggingFaceEndpoint
+import numpy as np
+from pydub import AudioSegment
+import os
+
+# Model IDs
+model_id = "meta-llama/Llama-3.2-1B-Instruct"
+model2_id = "meta-llama/Llama-3.2-1B-Instruct"
+whisper_model = "openai/whisper-small"  # Using Whisper model for audio transcription
+
+def get_llm_hf_inference(model_id, max_new_tokens=128, temperature=0.1):
+    """Returns a language model for HuggingFace inference."""
+    try:
+        llm = HuggingFaceEndpoint(
+            repo_id=model_id,
+            max_new_tokens=max_new_tokens,
+            temperature=temperature,
+            token=os.getenv("HF_TOKEN")
+        )
+        return llm
+    except Exception as e:
+        st.error(f"Error initializing model: {e}")
+        return None
+ 
+# Initialize Whisper transcription model
+def load_transcription_model():
+    try:
+        transcriber = pipeline("automatic-speech-recognition", model=whisper_model)
+        return transcriber
+    except Exception as e:
+        st.error(f"Error loading Whisper model: {e}")
+        return None
+ 
+# Preprocess audio to 16kHz mono
+def preprocess_audio(file):
+    audio = AudioSegment.from_file(file).set_frame_rate(16000).set_channels(1)
+    audio_samples = np.array(audio.get_array_of_samples()).astype(np.float32) / (2**15)
+    return audio_samples
+ 
+# Function to transcribe audio with preprocessing
+def transcribe_audio(file, transcriber):
+    audio = preprocess_audio(file)
+    transcription = transcriber(audio)["text"]
+    return transcription
+
+# Chatbot page content
+def display_chatbot():
+    st.title("Personal Psychologist Chatbot")
+    st.markdown(f"*This is a simple chatbot that acts as a psychologist and gives solutions to your psychological problems. It uses the {model_id}.*")
+ 
+    # Sidebar for settings
+    with st.sidebar:
+        # Reset Chat History
+        reset_history = st.button("Reset Chat History")
+        go_home = st.button("Back to Home")
+        if go_home:
+            st.session_state.chat_started = False
+            st.experimental_rerun()  # This will reload the app to show the homepage
+ 
+    # Initialize or reset chat history
+    if reset_history:
+        st.session_state.chat_history = [{"role": "assistant", "content": st.session_state.starter_message}]
+ 
+    def get_response(system_message, chat_history, user_text, model_id, max_new_tokens=256):
+        """Generates a response from the chatbot model."""
+        hf = get_llm_hf_inference(model_id=model_id, max_new_tokens=max_new_tokens)
+        if hf is None:
+            return "Error: Model not initialized.", chat_history
+ 
+        # Create the prompt template
+        prompt = PromptTemplate.from_template(
+            (
+                "[INST] {system_message}"
+                "\nCurrent Conversation:\n{chat_history}\n\n"
+                "\nUser: {user_text}.\n [/INST]"
+                "\nAI:"
+            )
+        )
+        chat = prompt | hf.bind(skip_prompt=True) | StrOutputParser(output_key='content')
+ 
+        # Generate the response
+        response = chat.invoke(input=dict(system_message=system_message, user_text=user_text, chat_history=chat_history))
+        response = response.split("AI:")[-1].strip()
+ 
+        # Enhanced end-of-conversation detection:
+        # Check if response contains low engagement or end-of-conversation patterns
+        low_engagement_threshold = 3  # Threshold for short responses
+        end_keywords = ["thank you", "thanks", "goodbye", "bye", "that's all", "done"]
+        
+        # Check for short responses over multiple turns
+        short_responses = len(user_text.split()) <= low_engagement_threshold
+        end_pattern_match = any(keyword in user_text.lower() for keyword in end_keywords)
+        
+        # Recent short responses pattern
+        recent_short_responses = all(len(msg["content"].split()) <= low_engagement_threshold for msg in chat_history[-2:])
+        response_is_acknowledgment = user_text.lower() in ["yes", "okay", "alright"]
+
+        # Trigger health report prompt based on combination of patterns
+        if (end_pattern_match or (short_responses and recent_short_responses)) and not response_is_acknowledgment:
+            follow_up_question = "Would you like to have a report of your current health? Yes/No"
+            response += f"\n\n{follow_up_question}"
+
+        # Update the chat history
+        chat_history.append({'role': 'user', 'content': user_text})
+        chat_history.append({'role': 'assistant', 'content': response})
+        return response, chat_history
+ 
+    def get_summary_of_chat_history(chat_history, model2_id):
+        """Generates a comprehensive summary of the chat history and a health report."""
+        hf = get_llm_hf_inference(model_id=model2_id, max_new_tokens=256)
+        if hf is None:
+            return "Error: Model not initialized."
+
+        # Format the chat content
+        chat_content = "\n".join([f"{message['role']}: {message['content']}" for message in chat_history])
+
+        # Improved summary prompt
+        prompt = PromptTemplate.from_template(
+            (
+                """
+                Generate a detailed report based on the following conversation between a therapist and patient. The report should include:
+
+                1. **Patient Information:** 
+                - Include placeholders for Name, Age, Gender, Date of Session.
+                
+                2. **Conversation Summary:** 
+                - Summarize the main points of the conversation, focusing on the patient’s primary concerns and emotional state. Note any specific causes of stress or distress, how these issues affect the patient's personal life, and their expressed desires or goals.
+
+                3. **Preliminary Diagnosis:** 
+                - Identify the main symptoms observed in the conversation, such as mood, energy levels, motivation, etc.
+                - Suggest a potential preliminary diagnosis based on the symptoms described, e.g., stress-induced burnout or other relevant concerns. Mention the need for further assessment if applicable.
+
+                4. **Recommendations & Strategies:** 
+                - Provide practical, achievable strategies tailored to the patient’s needs. 
+
+
+                Format the report neatly with headings and subheadings as shown in the example. Aim to keep the language supportive and professional.
+                """
+            )
+        )
+
+        summary = prompt | hf.bind(skip_prompt=True) | StrOutputParser(output_key='content')
+        summary_response = summary.invoke(input={"chat_content": chat_content})
+
+        return summary_response
+
+
+ 
+    # Load Whisper model for transcription
+    transcriber = load_transcription_model()
+ 
+    # User input for audio and text
+    st.markdown("### Choose your input:")
+    audio_file = st.file_uploader("Upload an audio file for transcription", type=["mp3", "wav", "m4a"])
+    st.session_state.user_text = st.chat_input(placeholder="Or enter your text here.")
+ 
+    # Check if audio file is uploaded and transcribe if available
+    if audio_file is not None and transcriber:
+        with st.spinner("Transcribing audio..."):
+            try:
+                st.session_state.user_text = transcribe_audio(audio_file, transcriber)
+                st.success("Audio transcribed successfully!")
+            except Exception as e:
+                st.error(f"Error transcribing audio: {e}")
+ 
+    # Chat interface
+    output_container = st.container()
+ 
+    # Display chat messages
+    with output_container:
+        for message in st.session_state.chat_history:
+            if message['role'] == 'system':
+                continue
+            with st.chat_message(message['role'], avatar=st.session_state['avatars'][message['role']]):
+                st.markdown(message['content'])
+ 
+    # Process text input for chatbot response
+    if st.session_state.user_text:
+        with st.chat_message("user", avatar=st.session_state.avatars['user']):
+            st.markdown(st.session_state.user_text)
+ 
+        with st.chat_message("assistant", avatar=st.session_state.avatars['assistant']):
+            with st.spinner("Addressing your concerns..."):
+                response, st.session_state.chat_history = get_response(
+                    system_message=st.session_state.system_message,
+                    user_text=st.session_state.user_text,
+                    chat_history=st.session_state.chat_history,
+                    model_id=model_id,
+                    max_new_tokens=st.session_state.max_response_length,
+                )
+                st.markdown(response)
+ 
+                # Check if the user has agreed to the report
+                if "yes" in st.session_state.user_text.lower() and "Would you like to have a report of your current health?" in response:
+                    with st.spinner("Generating your health report..."):
+                        report = get_summary_of_chat_history(st.session_state.chat_history, model2_id)
                         st.markdown(report)