Update app.py

app.py

@@ -1,109 +1,188 @@
 import streamlit as st
 import torch
-from audio_handler import AudioHandler
-from chat_handler import ChatHandler
+import torchaudio
+import soundfile as sf
+from pathlib import Path
 from transformers import pipeline
-import base64
-import time
-
-# Page configuration
-st.set_page_config(
-    page_title="Darija Mental Health Assistant",
-    page_icon="🤗",
-    layout="centered"
-)
-
-# Custom CSS
-st.markdown("""
+from langchain.memory import ConversationBufferMemory
+from langchain.chains import ConversationalRetrievalChain
+from langchain.llms import HuggingFaceHub
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.vectorstores import FAISS
+from langchain import PromptTemplate
+import os
+from dotenv import load_dotenv
+
+# Load environment variables
+load_dotenv()
+
+# CSS Styling
+css = """
 <style>
-    .stButton > button {
-        background-color: #4CAF50;
-        color: white;
-        padding: 10px 20px;
-        border-radius: 5px;
-    }
-    .main-header {
-        text-align: center;
-        color: #2e7d32;
-    }
-    .chat-container {
-        padding: 20px;
-        border-radius: 10px;
-        background-color: #f5f5f5;
-        margin: 10px 0;
-    }
+.chat-message { padding: 1.5rem; border-radius: 0.5rem; margin-bottom: 1rem; display: flex; }
+.chat-message.user { background-color: #2b313e; }
+.chat-message.bot { background-color: #475063; }
+.avatar { margin-right: 1rem; }
+.message { color: white; }
 </style>
-""", unsafe_allow_html=True)
-
-# Initialize session state
-if 'messages' not in st.session_state:
-    st.session_state.messages = []
-if 'recording' not in st.session_state:
-    st.session_state.recording = False
-if 'audio_handler' not in st.session_state:
-    st.session_state.audio_handler = AudioHandler()
-if 'chat_handler' not in st.session_state:
-    st.session_state.chat_handler = ChatHandler()
-
-def main():
-    # Header
-    st.markdown("<h1 class='main-header'>Darija Mental Health Assistant 🤗</h1>", unsafe_allow_html=True)
-    st.markdown("<h3 style='text-align: center;'>تحدث معي بالدارجة عن مشاعرك</h3>", unsafe_allow_html=True)
-
-    # Sidebar for settings
-    with st.sidebar:
-        st.header("Settings ⚙️")
-        voice_enabled = st.toggle("Enable Voice Input 🎤", True)
-        st.divider()
-        st.markdown("### About")
-        st.info("This assistant provides mental health support in Moroccan Arabic (Darija). Feel free to speak or type your thoughts.")
-
-    # Main chat interface
-    st.markdown("<div class='chat-container'>", unsafe_allow_html=True)
-    
-    # Voice input section
-    if voice_enabled:
-        cols = st.columns([1, 1])
-        with cols[0]:
+"""
+
+# Prompt template
+PROMPT_TEMPLATE = """
+You are a professional therapist who speaks Moroccan Arabic (Darija). 
+Respond with empathy and use therapeutic techniques.
+Always respond in Darija unless specifically asked to use another language.
+
+Previous conversation context:
+{chat_history}
+
+Current message: {question}
+
+Therapeutic response:
+"""
+
+class DarijaTherapist:
+    def __init__(self):
+        self.setup_models()
+        self.initialize_session_state()
+        self.setup_memory()
+        
+    def setup_models(self):
+        # Speech recognition setup
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.asr_pipe = pipeline(
+            "automatic-speech-recognition",
+            model="facebook/seamless-m4t-v2-large",
+            device=self.device
+        )
+        
+        # LLM and conversation chain setup
+        self.llm = HuggingFaceHub(
+            repo_id="MBZUAI-Paris/Atlas-Chat-27B",
+            model_kwargs={"temperature": 0.7, "max_length": 512},
+        )
+        
+        # Create embeddings and vector store
+        self.embeddings = HuggingFaceEmbeddings()
+        self.vectorstore = FAISS.from_texts(
+            ["Initial therapeutic context"], 
+            self.embeddings
+        )
+        
+    def setup_memory(self):
+        self.memory = ConversationBufferMemory(
+            memory_key="chat_history",
+            return_messages=True
+        )
+        
+        self.conversation_chain = ConversationalRetrievalChain.from_llm(
+            llm=self.llm,
+            retriever=self.vectorstore.as_retriever(),
+            memory=self.memory,
+            combine_docs_chain_kwargs={"prompt": PromptTemplate.from_template(PROMPT_TEMPLATE)}
+        )
+
+    def initialize_session_state(self):
+        if "messages" not in st.session_state:
+            st.session_state.messages = []
+        if "recording" not in st.session_state:
+            st.session_state.recording = False
+        if "audio_buffer" not in st.session_state:
+            st.session_state.audio_buffer = []
+
+    def handle_audio_input(self):
+        if not st.session_state.recording:
+            return
+        
+        try:
+            # Record audio using torchaudio
+            waveform, sample_rate = torchaudio.load("temp_audio.wav")
+            st.session_state.audio_buffer.append(waveform)
+        except Exception as e:
+            st.error(f"Error recording audio: {str(e)}")
+
+    def process_audio(self):
+        if not st.session_state.audio_buffer:
+            return None
+            
+        try:
+            # Concatenate audio buffer
+            audio_data = torch.cat(st.session_state.audio_buffer, dim=1)
+            # Save temporary file
+            torchaudio.save("temp_audio.wav", audio_data, 16000)
+            # Transcribe
+            audio, rate = sf.read("temp_audio.wav", dtype='float32')
+            result = self.asr_pipe(
+                audio,
+                generate_kwargs={"task": "transcribe", "language": "ara"}
+            )
+            return result["text"]
+        except Exception as e:
+            st.error(f"Error processing audio: {str(e)}")
+            return None
+        finally:
+            # Clear buffer
+            st.session_state.audio_buffer = []
+
+    def get_ai_response(self, user_input):
+        try:
+            response = self.conversation_chain({"question": user_input})
+            return response['answer']
+        except Exception as e:
+            st.error(f"Error getting AI response: {str(e)}")
+            return "عذراً، كاين شي مشكل. حاول مرة أخرى."
+
+    def run(self):
+        st.set_page_config(page_title="Darija AI Therapist", page_icon="🧠")
+        st.markdown(css, unsafe_allow_html=True)
+        
+        st.title("Darija AI Therapist 🧠")
+        st.subheader("تكلم معايا بالدارجة على اللي كيجول فبالك")
+
+        # Sidebar
+        with st.sidebar:
+            st.header("Settings ⚙️")
+            if st.button("Clear Chat History"):
+                st.session_state.messages = []
+                self.memory.clear()
+            
+            st.markdown("### About")
+            st.info("This AI therapist speaks Darija and is here to help. "
+                   "You can either type or speak your messages.")
+
+        # Audio input
+        col1, col2 = st.columns(2)
+        with col1:
             if st.button("🎤 Start Recording", disabled=st.session_state.recording):
                 st.session_state.recording = True
-                st.session_state.audio_handler.start_recording()
-                st.experimental_rerun()
+                st.session_state.audio_buffer = []
         
-        with cols[1]:
+        with col2:
             if st.button("⏹️ Stop Recording", disabled=not st.session_state.recording):
-                if st.session_state.recording:
-                    audio_file = st.session_state.audio_handler.stop_recording()
-                    with st.spinner("Processing your message..."):
-                        transcription = st.session_state.audio_handler.transcribe_audio(audio_file)
-                        process_input(transcription)
-                    st.session_state.recording = False
-                    st.experimental_rerun()
-
-    # Text input
-    user_input = st.text_input("اكتب رسالتك هنا:", key="text_input")
-    if user_input:
-        process_input(user_input)
-
-    # Display chat history
-    display_chat_history()
-
-    st.markdown("</div>", unsafe_allow_html=True)
-
-def process_input(user_input):
-    if user_input:
+                st.session_state.recording = False
+                transcription = self.process_audio()
+                if transcription:
+                    self.process_message(transcription)
+
+        # Text input
+        user_input = st.text_input("اكتب رسالتك هنا:", key="text_input")
+        if user_input:
+            self.process_message(user_input)
+
+        # Display chat history
+        for message in st.session_state.messages:
+            with st.chat_message(message["role"]):
+                st.write(message["content"])
+
+    def process_message(self, user_input):
         # Add user message
         st.session_state.messages.append({"role": "user", "content": user_input})
         
-        # Get bot response
+        # Get and add AI response
         with st.spinner("جاري التفكير..."):
-            response = st.session_state.chat_handler.get_response(user_input)
-            st.session_state.messages.append({"role": "assistant", "content": response})
-
-def display_chat_history():
-    for message in st.session_state.messages:
-        with st.chat_message(message["role"]):
-            st.write(message["content"])
+            ai_response = self.get_ai_response(user_input)
+            st.session_state.messages.append({"role": "assistant", "content": ai_response})
 
 if __name__ == "__main__":
-    main()
+    app = DarijaTherapist()
+    app.run()