Update app.py

app.py

@@ -14,11 +14,11 @@ from scipy.io.wavfile import write as write_wav
 import os
 import re
 from huggingface_hub import login
-import threading # <-- FIX: Added threading import
+import threading
 
 # --- Login to Hugging Face using secret ---
 # Make sure HF_TOKEN is set in your Hugging Face Space > Settings > Repository secrets
-hf_token = os.environ.get("hugface") # 
+hf_token = os.environ.get("hugface") # Using "HF_TOKEN" is the standard on Spaces
 if not hf_token:
     raise ValueError("HF_TOKEN not found. Please set it in Hugging Face Space repository secrets.")
 login(token=hf_token)
@@ -63,13 +63,12 @@ class WeeboAssistant:
 
         # LLM
         print(f"Loading LLM: {LLM_MODEL_ID}")
-        # <-- FIX: Initialize tokenizer separately to use it with the streamer
         self.llm_tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL_ID)
         self.llm_pipeline = pipeline(
             "text-generation",
             model=LLM_MODEL_ID,
             model_kwargs={"torch_dtype": self.torch_dtype},
-            tokenizer=self.llm_tokenizer, # Pass the tokenizer here
+            tokenizer=self.llm_tokenizer,
             device=self.device,
         )
         print("LLM pipeline loaded successfully.")
@@ -117,11 +116,20 @@ class WeeboAssistant:
         return output_path
 
     def get_llm_response(self, chat_history):
-        messages = [{'role': 'system', 'content': self.SYSTEM_PROMPT}]
-        for turn in chat_history:
-            messages.append({'role': 'user', 'content': turn[0]})
-            if turn[1] is not None:
-                messages.append({'role': 'assistant', 'content': turn[1]})
+        # <-- START OF FIX: Rebuild message history without a 'system' role -->
+        messages = []
+        for user_msg, assistant_msg in chat_history:
+            # Add the user's message
+            messages.append({"role": "user", "content": user_msg})
+            # Add the assistant's message if it exists
+            if assistant_msg:
+                messages.append({"role": "assistant", "content": assistant_msg})
+
+        # Prepend the system prompt to the content of the very first user message.
+        # This is the correct way to use a system prompt with Gemma models.
+        if messages:
+            messages[0]["content"] = f"{self.SYSTEM_PROMPT}\n\n{messages[0]['content']}"
+        # <-- END OF FIX -->
         
         prompt = self.llm_pipeline.tokenizer.apply_chat_template(
             messages, tokenize=False, add_generation_prompt=True
@@ -131,7 +139,6 @@ class WeeboAssistant:
             self.llm_pipeline.tokenizer.convert_tokens_to_ids("<|eot_id|>")
         ]
         
-        # <-- START OF FIX: Use TextIteratorStreamer instead of gr.TextIterator -->
         streamer = TextIteratorStreamer(
             self.llm_pipeline.tokenizer, skip_prompt=True, skip_special_tokens=True
         )
@@ -145,12 +152,10 @@ class WeeboAssistant:
             top_p=0.9,
         )
         
-        # Run the pipeline in a separate thread to enable streaming
         thread = threading.Thread(target=self.llm_pipeline, args=[prompt], kwargs=generation_kwargs)
         thread.start()
         
         return streamer
-        # <-- END OF FIX -->
 
 assistant = WeeboAssistant()
 
@@ -163,12 +168,12 @@ def s2s_pipeline(audio_input, chat_history):
         return
         
     chat_history.append((user_text, ""))
-    yield chat_history, None, "..." # Show thinking indicator
+    yield chat_history, None, "..."
     
     response_stream = assistant.get_llm_response(chat_history)
     llm_response_text = ""
     for text_chunk in response_stream:
-        llm_response_text += text_chunk # <-- FIX: Append chunk to full response
+        llm_response_text += text_chunk
         chat_history[-1] = (user_text, llm_response_text)
         yield chat_history, None, llm_response_text
         
@@ -183,7 +188,7 @@ def t2t_pipeline(text_input, chat_history):
     response_stream = assistant.get_llm_response(chat_history)
     llm_response_text = ""
     for text_chunk in response_stream:
-        llm_response_text += text_chunk # <-- FIX: Append chunk to full response
+        llm_response_text += text_chunk
         chat_history[-1] = (text_input, llm_response_text)
         yield chat_history
 
@@ -232,7 +237,7 @@ with gr.Blocks(theme=gr.themes.Soft(), title="Msaidizi wa Kiswahili") as demo:
         outputs=[s2s_chatbot, s2s_audio_out, s2s_text_out],
         queue=True
     ).then(
-        fn=lambda: gr.Audio(value=None), # Clear audio input after submit
+        fn=lambda: gr.Audio(value=None),
         inputs=None,
         outputs=s2s_audio_in
     )
@@ -240,7 +245,7 @@ with gr.Blocks(theme=gr.themes.Soft(), title="Msaidizi wa Kiswahili") as demo:
     t2t_submit_btn.click(
         fn=t2t_pipeline,
         inputs=[t2t_text_in, t2t_chatbot],
-        outputs=[t2t_chatbot], # <-- FIX: Only output to the chatbot
+        outputs=[t2t_chatbot],
         queue=True
     ).then(
         fn=clear_textbox,
@@ -248,7 +253,6 @@ with gr.Blocks(theme=gr.themes.Soft(), title="Msaidizi wa Kiswahili") as demo:
         outputs=t2t_text_in
     )
     
-    # Also allow Enter key to submit text
     t2t_text_in.submit(
         fn=t2t_pipeline,
         inputs=[t2t_text_in, t2t_chatbot],
@@ -260,7 +264,6 @@ with gr.Blocks(theme=gr.themes.Soft(), title="Msaidizi wa Kiswahili") as demo:
         outputs=t2t_text_in
     )
 
-
     tool_s2t_btn.click(
         fn=assistant.transcribe_audio,
         inputs=tool_s2t_audio_in,