Update app.py

app.py

@@ -18,21 +18,25 @@ def generate_response(user_message):
         # Obtener el prompt del sistema
         system_prompt = get_system_prompt()
         
-        # Crear el prompt completo
-        prompt = f"{system_prompt}\n\nPregunta: {user_message}\n\nRespuesta:"
+        # Crear el prompt completo para TinyLlama
+        prompt = f"""<|system|>
+{system_prompt}
+<|user|>
+{user_message}
+<|assistant|>"""
         
         # Generar respuesta usando el pipeline
         response = generator(
             prompt,
-            max_new_tokens=512,
+            max_new_tokens=256,  # Reducido para mayor velocidad
             temperature=0.7,
             top_p=0.9,
             do_sample=True,
             num_return_sequences=1
         )[0]["generated_text"]
         
-        # Extraer solo la respuesta (después de "Respuesta:")
-        assistant_response = response.split("Respuesta:")[-1].strip()
+        # Extraer solo la respuesta del asistente (después del último <|assistant|>)
+        assistant_response = response.split("<|assistant|>")[-1].strip()
         
         # Forzar recolección de basura
         gc.collect()
@@ -45,18 +49,19 @@ def generate_response(user_message):
         return f"Lo siento, ha ocurrido un error: {str(e)}"
 
 # Cargar el modelo y crear el pipeline
-print("Cargando modelo ligero...")
+print("Cargando modelo ultra ligero...")
 
-# Usar un modelo pequeño optimizado para español
-model_name = "microsoft/phi-2"  # Modelo de 2.7B parámetros, mucho más ligero que 7B
+# Usar TinyLlama, un modelo muy pequeño (1.1B parámetros) pero potente
+model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
 
-# Configuración para reducir el uso de memoria
-# IMPORTANTE: Eliminamos device_map="auto" que no es compatible con Phi-2
+# Configuración para reducir el uso de memoria y aumentar velocidad
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
     torch_dtype=torch.float32,
-    low_cpu_mem_usage=True
+    low_cpu_mem_usage=True,
+    # Cargamos en 8-bit para reducir memoria y aumentar velocidad
+    load_in_8bit=True
 )
 
 # Crear el pipeline de generación de texto
@@ -88,6 +93,6 @@ demo = gr.Interface(
 
 # Lanzar la aplicación con configuración para ahorrar memoria
 if __name__ == "__main__":
-    # Configurar menos workers para ahorrar memoria
+    # Configurar menos workers para ahorrar memoria y aumentar velocidad
     demo.queue(max_size=1).launch(share=False, debug=False)