Update app.py

app.py

@@ -31,20 +31,19 @@ css = """
 </style>
 """
 
-# Prompt template
+# Updated prompt template for Mixtral
 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.
+<s>[INST] You are a professional therapist who speaks Moroccan Arabic (Darija).
+Act as a compassionate therapist and provide empathetic responses using therapeutic techniques.
+Always respond in Darija unless specifically asked otherwise.
 
 Previous conversation:
 {chat_history}
 
 User message: {question}
 
-Additional context: {context}
-
-Therapeutic response in Darija:
+Context: {context}
+[/INST]
 """
 
 class DarijaTherapist:
@@ -67,31 +66,35 @@ class DarijaTherapist:
             
             # Configure retry strategy
             retry_strategy = Retry(
-                total=3,  # number of retries
-                backoff_factor=1,  # wait 1, 2, 4 seconds between retries
-                status_forcelist=[429, 500, 502, 503, 504]  # HTTP status codes to retry on
+                total=3,
+                backoff_factor=1,
+                status_forcelist=[429, 500, 502, 503, 504]
             )
             
             # Create session with retry strategy
             session = requests.Session()
             session.mount("https://", HTTPAdapter(max_retries=retry_strategy))
             
-            # LLM setup with increased timeout and retry session
+            # Updated LLM setup for Mixtral
             self.llm = HuggingFaceEndpoint(
-                endpoint_url="https://api-inference.huggingface.co/models/MBZUAI-Paris/Atlas-Chat-2B",
+                endpoint_url="https://api-inference.huggingface.co/models/mistralai/Mixtral-8x7B-Instruct-v0.1",
                 task="text-generation",
                 temperature=0.7,
                 do_sample=True,
                 return_full_text=False,
-                timeout = 300,  # increase timeout to 5 minutes
+                timeout=300,
                 model_kwargs={
-                    "max_length": 512
-                    },
+                    "max_new_tokens": 2048,
+                    "top_p": 0.9,
+                    "repetition_penalty": 1.2,
+                    "return_text": True,
+                    "stop": ["</s>"]
+                },
                 huggingfacehub_api_token=os.getenv("HUGGINGFACE_API_TOKEN"),
-                client=session  # use configured session
+                client=session
             )
             
-            # Rest of the setup remains the same
+            # Embeddings setup
             self.embeddings = HuggingFaceBgeEmbeddings(
                 model_name="BAAI/bge-large-en"
             )
@@ -176,27 +179,35 @@ class DarijaTherapist:
                     "question": user_input,
                     "chat_history": self.memory.chat_memory.messages[-5:]  # Limit context window
                 })
+                
+                if not response or 'answer' not in response:
+                    if attempt < max_retries - 1:
+                        time.sleep(2 ** attempt)
+                        continue
+                    return "عذراً، كاين مشكل. حاول مرة أخرى."
+                
                 return response['answer']
 
             except requests.exceptions.HTTPError as e:
                 if e.response.status_code == 424:
                     if attempt < max_retries - 1:
                         st.warning("Model error, retrying with simplified input...")
-                        # Try with shorter context
-                        user_input = user_input[:256]
                         time.sleep(2 ** attempt)
                         continue
-                    return "عذراً، كاين مشكل مع النموذج. جرب سؤال أقصر."
-                    
+                return "عذراً، كاين مشكل مع النموذج. جرب سؤال أقصر."
+                
             except requests.exceptions.ReadTimeout:
                 if attempt < max_retries - 1:
                     st.warning(f"Attempt {attempt + 1} timed out, retrying...")
                     time.sleep(2 ** attempt)
                     continue
                 return "عذراً، الخادم بطيء حالياً. حاول مرة أخرى."
-                
+            
             except Exception as e:
                 st.error(f"Error: {str(e)}")
+                if attempt < max_retries - 1:
+                    time.sleep(2 ** attempt)
+                    continue
                 return "عذراً، كاين شي مشكل. حاول مرة أخرى."
 
     def run(self):
@@ -241,7 +252,8 @@ class DarijaTherapist:
         
         with st.spinner("جاري التفكير..."):
             ai_response = self.get_ai_response(user_input)
-            st.session_state.messages.append({"role": "assistant", "content": ai_response})
+            if ai_response:
+                st.session_state.messages.append({"role": "assistant", "content": ai_response})
 
 if __name__ == "__main__":
     app = DarijaTherapist()