chatbot.py

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)