Update app.py

app.py

@@ -71,13 +71,24 @@ except Exception as e:
     logger.error(f"Error during initialization: {str(e)}")
     raise
 
-# Improved system message
-system_message = """You are AQuaBot, an AI assistant focused on providing accurate and environmentally conscious information. Your responses should be:
-1. Clear and concise yet informative
-2. Based on verified information when discussing economic and financial topics
-3. Balanced and well-reasoned
-4. Mindful of environmental impact
-5. Professional but conversational in tone"""
+
+
+@spaces.GPU(duration=60)
+@torch.inference_mode()
+def clean_response(text):
+    """Limpia la respuesta del modelo eliminando etiquetas y texto no deseado"""
+    # Eliminar etiquetas INST y wikipedia references
+    text = text.replace('[INST]', '').replace('[/INST]', '')
+    text = text.replace('(You can find more about it at wikipedia)', '')
+    
+    # Eliminar cualquier texto que comience con "User:" o "Assistant:"
+    lines = text.split('\n')
+    cleaned_lines = []
+    for line in lines:
+        if not line.strip().startswith(('User:', 'Assistant:', 'Human:', 'AI:')):
+            cleaned_lines.append(line)
+    
+    return '\n'.join(cleaned_lines).strip()
 
 @spaces.GPU(duration=60)
 @torch.inference_mode()
@@ -90,14 +101,21 @@ def generate_response(user_input, chat_history):
         input_water_consumption = calculate_water_consumption(user_input, True)
         total_water_consumption += input_water_consumption
 
-        # Create a clean conversation history
-        conversation_history = ""
+        # Format conversation history without using INST tags
+        formatted_history = ""
         if chat_history:
-            for user_msg, assistant_msg in chat_history:
-                conversation_history += f"User: {user_msg}\nAssistant: {assistant_msg}\n\n"
+            for prev_input, prev_response in chat_history:
+                formatted_history += f"Question: {prev_input}\nAnswer: {prev_response}\n\n"
+
+        # Create prompt using a más natural format
+        prompt = f"""
+{system_message}
+
+Previous conversation:
+{formatted_history}
 
-        # Create a clean prompt format
-        prompt = f"{system_message}\n\nConversation History:\n{conversation_history}\nUser: {user_input}\nAssistant:"
+Question: {user_input}
+Answer:"""
 
         logger.info("Generating model response...")
         outputs = model_gen(
@@ -105,12 +123,23 @@ def generate_response(user_input, chat_history):
             max_new_tokens=512,
             return_full_text=False,
             pad_token_id=tokenizer.eos_token_id,
+            do_sample=True,
+            temperature=0.7,
+            top_p=0.9,
+            repetition_penalty=1.1
         )
-        logger.info("Model response generated successfully")
-
-        # Clean up response
-        assistant_response = outputs[0]['generated_text'].strip()
-        assistant_response = assistant_response.split('User:')[0].split('Assistant:')[-1].strip()
+        
+        # Limpiar y procesar la respuesta
+        assistant_response = outputs[0]['generated_text']
+        assistant_response = clean_response(assistant_response)
+        
+        # Si la respuesta sigue conteniendo texto no deseado, intentar extraer solo la parte relevante
+        if 'Question:' in assistant_response or 'Answer:' in assistant_response:
+            parts = assistant_response.split('Answer:')
+            if len(parts) > 1:
+                assistant_response = parts[1].split('Question:')[0].strip()
+        
+        logger.info("Response cleaned and processed")
 
         # Calculate water consumption for output
         output_water_consumption = calculate_water_consumption(assistant_response, False)
@@ -119,7 +148,7 @@ def generate_response(user_input, chat_history):
         # Update chat history
         chat_history.append([user_input, assistant_response])
 
-        # Water consumption message
+        # Update water consumption display
         water_message = f"""
         <div style="position: fixed; top: 20px; right: 20px;
                     background-color: white; padding: 15px;
@@ -142,6 +171,19 @@ def generate_response(user_input, chat_history):
         chat_history.append([user_input, error_message])
         return chat_history, show_water
 
+# Actualizar el system message para ser más específico sobre el formato
+system_message = """You are AQuaBot, an AI assistant focused on providing accurate and environmentally conscious information. 
+
+Guidelines for your responses:
+1. Provide direct, clear answers without any special tags or markers
+2. Do not reference external sources like Wikipedia in your responses
+3. Stay focused on the question asked
+4. Be concise but informative
+5. Be mindful of environmental impact
+6. Use a natural, conversational tone
+
+Remember: Never include [INST] tags or other technical markers in your responses."""
+
 
 # Constants for water consumption calculation
 WATER_PER_TOKEN = {