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)