Update app.py

app.py

@@ -86,20 +86,20 @@ def generate_response(user_input, chat_history):
         logger.info("Generating response for user input...")
         global total_water_consumption
 
-        # Calcula el consumo de agua para la entrada
+        # Calculate water consumption for input
         input_water_consumption = calculate_water_consumption(user_input, True)
         total_water_consumption += input_water_consumption
 
-        # Construye el historial de conversación sin etiquetas `[INST]`
+        # Create prompt with Llama 2 chat format
         conversation_history = ""
         if chat_history:
             for message in chat_history:
+                # Remove any [INST] tags from the history
                 user_msg = message[0].replace("[INST]", "").replace("[/INST]", "").strip()
                 assistant_msg = message[1].replace("[INST]", "").replace("[/INST]", "").strip()
-                conversation_history += f"{user_msg} {assistant_msg} "
-
-        # Crea el prompt sin las etiquetas `[INST]`
-        prompt = f"{system_message}\n\n{conversation_history}{user_input}"
+                conversation_history += f"[INST] {user_msg} [/INST] {assistant_msg} "
+        
+        prompt = f"<s>[INST] {system_message}\n\n{conversation_history}[INST] {user_input} [/INST]"
 
         logger.info("Generating model response...")
         outputs = model_gen(
@@ -107,20 +107,41 @@ def generate_response(user_input, chat_history):
             max_new_tokens=256,
             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")
 
-        # Limpia la salida del modelo
-        assistant_response = outputs[0]['generated_text'].replace("[INST]", "").replace("[/INST]", "").strip()
+        # Clean up the response by removing any [INST] tags and trimming
+        assistant_response = outputs[0]['generated_text'].strip()
+        assistant_response = assistant_response.replace("[INST]", "").replace("[/INST]", "").strip()
+        
+        # If the response is too short, try to generate a more detailed one
+        if len(assistant_response.split()) < 10:
+            prompt += "\nPlease provide a more detailed answer with context and explanation."
+            outputs = model_gen(
+                prompt,
+                max_new_tokens=256,
+                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
+            )
+            assistant_response = outputs[0]['generated_text'].strip()
+            assistant_response = assistant_response.replace("[INST]", "").replace("[/INST]", "").strip()
 
-        # Calcula el consumo de agua para la salida
+        # Calculate water consumption for output
         output_water_consumption = calculate_water_consumption(assistant_response, False)
         total_water_consumption += output_water_consumption
 
-        # Actualiza el historial de chat
+        # Update chat history with the cleaned messages
         chat_history.append([user_input, assistant_response])
 
-        # Prepara el mensaje de consumo de agua
+        # Prepare water consumption message
         water_message = f"""
         <div style="position: fixed; top: 20px; right: 20px;
                     background-color: white; padding: 15px;
@@ -143,7 +164,6 @@ def generate_response(user_input, chat_history):
         chat_history.append([user_input, error_message])
         return chat_history, show_water
 
-
 # Constants for water consumption calculation
 WATER_PER_TOKEN = {
     "input_training": 0.0000309,