Update app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 import os
+import gc
 
 # Variables globales para el modelo y tokenizador
 model = None
@@ -17,21 +18,27 @@ def load_model_if_needed():
     global model, tokenizer
     
     if model is None:
-        print("Cargando modelo Zephyr-7B...")
+        print("Cargando modelo Mistral-7B-Instruct-v0.1...")
         
         # Cargar el tokenizador
-        model_name = "HuggingFaceH4/zephyr-7b-beta"
+        model_name = "mistralai/Mistral-7B-Instruct-v0.1"
         tokenizer = AutoTokenizer.from_pretrained(model_name)
         
-        # Cargar el modelo con configuraciones para ahorrar memoria
+        # Configuración para CPU con optimizaciones de memoria
         model = AutoModelForCausalLM.from_pretrained(
             model_name,
-            torch_dtype=torch.float16,  # Usar precisión reducida
-            device_map="auto",          # Distribuir el modelo automáticamente
-            load_in_8bit=True,          # Cuantización a 8 bits
+            torch_dtype=torch.float32,  # Usar float32 para CPU
+            low_cpu_mem_usage=True,     # Optimización para CPU con poca memoria
         )
         
-        print("Modelo cargado correctamente!")
+        # Mover el modelo a CPU explícitamente
+        model = model.to("cpu")
+        
+        # Forzar recolección de basura para liberar memoria
+        gc.collect()
+        torch.cuda.empty_cache() if torch.cuda.is_available() else None
+        
+        print("Modelo cargado correctamente en CPU!")
 
 # Función principal que procesa las preguntas del usuario
 def generate_response(user_message):
@@ -45,25 +52,22 @@ def generate_response(user_message):
         # Obtener el prompt del sistema
         system_prompt = get_system_prompt()
         
-        # Crear el formato de conversación para Zephyr
-        messages = [
-            {"role": "system", "content": system_prompt},
-            {"role": "user", "content": user_message}
-        ]
+        # Formato de prompt para Mistral-7B-Instruct-v0.1
+        prompt = f"""<s>[INST] {system_prompt}
+
+Pregunta del usuario: {user_message} [/INST]</s>"""
         
-        # Convertir mensajes al formato que espera el modelo
-        prompt = tokenizer.apply_chat_template(messages, tokenize=False)
+        # Tokenizar el prompt
+        inputs = tokenizer(prompt, return_tensors="pt")
         
-        # Generar respuesta
-        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-        
-        # Configuración de generación
+        # Configuración de generación optimizada para CPU
         generation_config = {
-            "max_new_tokens": 500,
+            "max_new_tokens": 512,      # Limitar tokens para ahorrar memoria
             "temperature": 0.7,
             "top_p": 0.9,
             "do_sample": True,
-            "pad_token_id": tokenizer.eos_token_id
+            "pad_token_id": tokenizer.eos_token_id,
+            "num_return_sequences": 1
         }
         
         # Generar respuesta
@@ -73,8 +77,12 @@ def generate_response(user_message):
         # Decodificar la respuesta
         full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
         
-        # Extraer solo la respuesta del asistente (después del último mensaje del usuario)
-        assistant_response = full_response.split("assistant:")[-1].strip()
+        # Extraer solo la respuesta del asistente (después del prompt)
+        assistant_response = full_response.replace(prompt.replace("<s>", "").replace("</s>", ""), "").strip()
+        
+        # Forzar recolección de basura para liberar memoria
+        gc.collect()
+        torch.cuda.empty_cache() if torch.cuda.is_available() else None
         
         return assistant_response
     
@@ -100,7 +108,7 @@ demo = gr.Interface(
     allow_flagging="never"
 )
 
-# Lanzar la aplicación
+# Lanzar la aplicación con configuración para ahorrar memoria
 if __name__ == "__main__":
-    demo.launch()
-
+    # Configurar menos workers para ahorrar memoria
+    demo.queue(max_size=1).launch(share=False, debug=False)