Update app.py

app.py

@@ -1,65 +1,69 @@
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
-from threading import Thread
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import gradio as gr
+import warnings
 
-# 1. Configuración robusta del modelo (usamos uno más liviano)
-MODEL_NAME = "microsoft/phi-2"  # Modelo eficiente para Spaces
+# 1. Configuración a prueba de errores
+MODEL_NAME = "HuggingFaceH4/zephyr-7b-beta"  # Modelo optimizado para Spaces
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
-# 2. Carga con manejo de errores
+# 2. Carga segura del modelo
 try:
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_NAME,
         torch_dtype=torch.float16,
-        device_map="auto"
+        device_map="auto",
+        low_cpu_mem_usage=True
     )
+    if DEVICE == "cuda":
+        model = model.to(DEVICE)
 except Exception as e:
-    raise gr.Error(f"❌ Error al cargar el modelo: {str(e)}")
+    raise gr.Error(f"Error inicialización: {str(e)}")
 
-# 3. Función de generación mejorada
-def chat_with_gerardo(message, history):
+# 3. Función de chat mejorada
+def generate_response(message, history):
     try:
-        # Construimos el prompt manualmente
-        prompt = "Eres Gerardo, un asistente IA útil. Responde preguntas claramente.\n\n"
-        for user_msg, bot_msg in history:
-            prompt += f"Usuario: {user_msg}\nGerardo: {bot_msg}\n"
-        prompt += f"Usuario: {message}\nGerardo:"
+        # Limpieza de memoria
+        if DEVICE == "cuda":
+            torch.cuda.empty_cache()
         
-        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-        streamer = TextIteratorStreamer(tokenizer)
+        # Formateo del prompt
+        messages = [{"role": "user", "content": message}]
+        prompt = tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True
+        )
         
-        generation_kwargs = dict(
-            inputs,
-            streamer=streamer,
-            max_new_tokens=300,  # Reducido para evitar OOM
+        # Generación con parámetros seguros
+        inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=256,
             temperature=0.7,
+            do_sample=True,
             pad_token_id=tokenizer.eos_token_id
         )
         
-        thread = Thread(target=model.generate, kwargs=generation_kwargs)
-        thread.start()
+        # Decodificación segura
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return response.split("assistant\n")[-1].strip()
         
-        partial_message = ""
-        for new_token in streamer:
-            partial_message += new_token
-            yield partial_message
-            
-    except torch.cuda.OutOfMemoryError:
-        yield "⚠️ Error: Memoria de GPU agotada. Intenta con una consulta más corta."
     except Exception as e:
-        yield f"❌ Error: {str(e)}"
+        warnings.warn(str(e))
+        return f"Error: {str(e)}"
 
-# 4. Interfaz con configuración optimizada
-with gr.Blocks(title="Chatbot Gerardo v2.0") as demo:
-    gr.Markdown("## 🤖 Chatbot creado por Gerardo")
-    gr.ChatInterface(
-        fn=chat_with_gerardo,
-        examples=["Hola Gerardo", "¿Qué puedes hacer?"],
-        cache_examples=False  # Importante para Spaces
+# 4. Interfaz a prueba de fallos
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("## 🤖 Chatbot Gerardo - Versión Estable")
+    chatbot = gr.ChatInterface(
+        fn=generate_response,
+        examples=["Hola", "¿Cómo estás?"],
+        title="Chatbot Personalizado",
+        description="Asistente IA creado por Gerardo",
+        cache_examples=False
     )
 
-# 5. Lanzamiento específico para Hugging Face
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
-  
+    demo.launch(server_name="0.0.0.0", server_port=7860)