Update app.py

app.py

@@ -76,7 +76,93 @@ except Exception as e:
     raise
 
 # Configure system message
-system_message = """You are a helpful AI assistant called AQuaBot. You provide direct, clear answers to questions while being aware of environmental impact. Keep your responses natural and concise. Respond directly to questions without using any special tags or markers."""
+
+system_message = """You are a helpful AI assistant called AQuaBot. You provide direct, clear, and detailed answers to questions while being aware of environmental impact. Keep your responses natural and informative, but concise. Always provide context and explanations with your answers. Respond directly to questions without using any special tags or markers."""
+
+@spaces.GPU(duration=60)
+@torch.inference_mode()
+def generate_response(user_input, chat_history):
+    try:
+        logger.info("Generating response for user input...")
+        global total_water_consumption
+
+        # Calculate water consumption for input
+        input_water_consumption = calculate_water_consumption(user_input, True)
+        total_water_consumption += input_water_consumption
+
+        # 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"[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(
+            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
+        )
+        logger.info("Model response generated successfully")
+
+        # 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()
+
+        # Calculate water consumption for output
+        output_water_consumption = calculate_water_consumption(assistant_response, False)
+        total_water_consumption += output_water_consumption
+
+        # Update chat history with the cleaned messages
+        chat_history.append([user_input, assistant_response])
+
+        # Prepare water consumption message
+        water_message = f"""
+        <div style="position: fixed; top: 20px; right: 20px;
+                    background-color: white; padding: 15px;
+                    border: 2px solid #ff0000; border-radius: 10px;
+                    box-shadow: 0 2px 4px rgba(0,0,0,0.1);">
+            <div style="color: #ff0000; font-size: 24px; font-weight: bold;">
+                💧 {total_water_consumption:.4f} ml
+            </div>
+            <div style="color: #666; font-size: 14px;">
+                Water Consumed
+            </div>
+        </div>
+        """
+
+        return chat_history, water_message
+
+    except Exception as e:
+        logger.error(f"Error in generate_response: {str(e)}")
+        error_message = f"An error occurred: {str(e)}"
+        chat_history.append([user_input, error_message])
+        return chat_history, show_water
 
 # Constants for water consumption calculation
 WATER_PER_TOKEN = {