app.py

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import gradio as gr

# 1. Configuración mejorada
MODEL_NAME = "HuggingFaceH4/zephyr-7b-beta"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# 2. Carga del modelo con seguridad mejorada
try:
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    
    # Ajustamos el dtype según el dispositivo
    torch_dtype = torch.float16 if DEVICE == "cuda" else torch.float32
    
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        torch_dtype=torch_dtype,
        device_map="auto"
    )
    
    # Forzamos el modelo al dispositivo correcto
    model.to(DEVICE)
    
except Exception as e:
    raise gr.Error(f"Error al cargar el modelo: {str(e)}")

# 3. Función de generación robusta
def generate_response(message, history):
    try:
        # Preparamos el input
        messages = [{"role": "user", "content": message}]
        prompt = tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True
        )
        
        inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
        
        # Generación con parámetros seguros
        outputs = model.generate(
            **inputs,
            max_new_tokens=256,
            temperature=0.7,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id
        )
        
        # Decodificamos la respuesta
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        return response.split("assistant\n")[-1].strip()
        
    except Exception as e:
        return f"⚠️ Error en la generación: {str(e)}"

# 4. Interfaz mejorada
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("## 🚀 Chatbot Gerardo - Versión Estable")
    gr.ChatInterface(
        fn=generate_response,
        examples=["Hola", "¿Cómo funciona esto?"],
        title="Chatbot de Gerardo",
        description="Asistente IA sin errores de CPU/GPU"
    )

if __name__ == "__main__":
    demo.launch()