import gradio as gr
import torch
from transformers import pipeline

# 1. Configuración del modelo
model_id = "Fifthoply/AyudaAlan-0.1" 

print(f"Cargando {model_id} en CPU...")
pipe = pipeline(
    "text-generation",
    model=model_id,
    dtype=torch.float32,
    device_map="cpu"
)

def chat_responder(message, history):
    prompt = "<|im_start|>system\nEres un asistente breve.<|im_end|>\n"
    
    # Procesar historial con detección de formato automática
    for msg in history:
        # Caso 1: Formato de diccionario {"role": "...", "content": "..."}
        if isinstance(msg, dict):
            role = msg.get("role", "user")
            content = msg.get("content", "")
            prompt += f"<|im_start|>{role}\n{content}<|im_end|>\n"
        # Caso 2: Formato de lista clásica [user, assistant]
        elif isinstance(msg, (list, tuple)):
            u, a = msg
            prompt += f"<|im_start|>user\n{u}<|im_end|>\n<|im_start|>assistant\n{a}<|im_end|>\n"
    
    prompt += f"<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n"
    
    # Generación
    outputs = pipe(
        prompt, 
        max_new_tokens=25, 
        do_sample=False,
        pad_token_id=pipe.tokenizer.pad_token_id,
        eos_token_id=pipe.tokenizer.eos_token_id
    )
    
    # Extraer la respuesta (usando el índice [0] para asegurar el acceso al texto)
    generated_text = outputs[0]['generated_text']
    respuesta = generated_text.split("<|im_start|>assistant\n")[-1].split("<|im_end|>")[0].strip()
    
    return respuesta

# 2. Interfaz de Usuario (Sin argumentos problemáticos)
demo = gr.ChatInterface(
    fn=chat_responder,
    title="Ayuda Alan 🦥",
    description="Hazme una pregunta y te diré cómo se hace.",
    examples=["¿Cómo se toma agua?", "¿Cómo se salta?"],
)

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