Update app.py

app.py

@@ -36,7 +36,7 @@ try:
     )
     logger.info("Tokenizer loaded successfully")
 
-    # Load model - Modificada la carga del modelo
+    # Load model
     logger.info("Loading model...")
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
@@ -47,7 +47,7 @@ try:
         model = model.to(device)
     logger.info("Model loaded successfully")
 
-    # Create pipeline - Modificada la creación del pipeline
+    # Create pipeline
     logger.info("Creating generation pipeline...")
     model_gen = pipeline(
         "text-generation",
@@ -67,13 +67,10 @@ except Exception as e:
     raise
 
 # Configure system message
-system_message = {
-    "role": "system",
-    "content": """You are AQuaBot, an AI assistant aware of environmental impact.
-    You help users with any topic while raising awareness about water consumption
-    in AI. Did you know that training GPT-3 consumed 5.4 million liters of water,
-    equivalent to the daily consumption of a city of 10,000 people?"""
-}
+system_message = """You are AQuaBot, an AI assistant aware of environmental impact.
+You help users with any topic while raising awareness about water consumption
+in AI. Did you know that training GPT-3 consumed 5.4 million liters of water,
+equivalent to the daily consumption of a city of 10,000 people?"""
 
 # Constants for water consumption calculation
 WATER_PER_TOKEN = {
@@ -84,7 +81,6 @@ WATER_PER_TOKEN = {
 }
 
 # Initialize variables
-messages = [system_message]
 total_water_consumption = 0
 
 def calculate_tokens(text):
@@ -100,30 +96,27 @@ def calculate_water_consumption(text, is_input=True):
         return tokens * (WATER_PER_TOKEN["input_training"] + WATER_PER_TOKEN["input_inference"])
     return tokens * (WATER_PER_TOKEN["output_training"] + WATER_PER_TOKEN["output_inference"])
 
+def format_message(role, content):
+    return {"role": role, "content": content}
+
 @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, messages
+        global total_water_consumption
 
         # Calculate water consumption for input
         input_water_consumption = calculate_water_consumption(user_input, True)
         total_water_consumption += input_water_consumption
 
-        # Add user input to messages
-        messages.append({"role": "user", "content": user_input})
-
         # Create prompt
-        prompt = ""
-        for m in messages:
-            if m["role"] == "system":
-                prompt += f"<START SYSTEM MESSAGE>\n{m['content']}\n<END SYSTEM MESSAGE>\n\n"
-            elif m["role"] == "user":
-                prompt += f"User: {m['content']}\n"
-            else:
-                prompt += f"Assistant: {m['content']}\n"
-        prompt += "Assistant:"
+        conversation_history = ""
+        if chat_history:
+            for message in chat_history:
+                conversation_history += f"User: {message[0]}\nAssistant: {message[1]}\n"
+        
+        prompt = f"{system_message}\n\n{conversation_history}User: {user_input}\nAssistant:"
 
         logger.info("Generating model response...")
         outputs = model_gen(
@@ -140,11 +133,8 @@ def generate_response(user_input, chat_history):
         output_water_consumption = calculate_water_consumption(assistant_response, False)
         total_water_consumption += output_water_consumption
 
-        # Add assistant's response to messages
-        messages.append({"role": "assistant", "content": assistant_response})
-
-        # Update chat history
-        chat_history.append((user_input, assistant_response))
+        # Update chat history with the new formatted messages
+        chat_history.append([user_input, assistant_response])
 
         # Prepare water consumption message
         water_message = f"""
@@ -166,7 +156,7 @@ def generate_response(user_input, chat_history):
     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))
+        chat_history.append([user_input, error_message])
         return chat_history, show_water
 
 # Create Gradio interface
@@ -183,7 +173,7 @@ try:
             </div>
         """)
 
-        chatbot = gr.Chatbot(type="messages")
+        chatbot = gr.Chatbot()
         message = gr.Textbox(
             placeholder="Type your message here...",
             show_label=False