"""Chat web con Minerva v0.1 (https://huggingface.co/AxoneSystem/Minerva-v0.1)."""
import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

REPO = "AxoneSystem/Minerva-v0.1"
tok = AutoTokenizer.from_pretrained(REPO)
modelo = AutoModelForCausalLM.from_pretrained(REPO).eval()
CONTEXTO = 512      # tokens que Minerva puede ver a la vez
MAX_RESPUESTA = 200


def texto(contenido):
    """El contenido de un mensaje puede venir como texto, dict o lista de partes (según la versión de Gradio)."""
    if isinstance(contenido, str):
        return contenido
    if isinstance(contenido, dict):
        return contenido.get("text") or ""
    if isinstance(contenido, (list, tuple)):
        return " ".join(texto(p) for p in contenido).strip()
    return ""


def a_mensajes(historial):
    """Acepta el historial de Gradio en formato de mensajes o de pares (según la versión)."""
    mensajes = []
    for item in historial or []:
        if isinstance(item, dict):
            t = texto(item.get("content"))
            if item.get("role") in ("user", "assistant") and t:
                mensajes.append({"role": item["role"], "content": t})
        else:
            usuario, asistente = texto(item[0]), texto(item[1])
            if usuario:
                mensajes.append({"role": "user", "content": usuario})
            if asistente:
                mensajes.append({"role": "assistant", "content": asistente})
    return mensajes


def responder(mensaje, historial, temperatura):
    mensajes = a_mensajes(historial) + [{"role": "user", "content": texto(mensaje)}]
    entrada = tok.apply_chat_template(mensajes, add_generation_prompt=True, return_tensors="pt", return_dict=True)
    ids = entrada["input_ids"][:, -(CONTEXTO - MAX_RESPUESTA):]  # solo recuerda lo último
    with torch.no_grad():
        salida = modelo.generate(input_ids=ids, attention_mask=torch.ones_like(ids), max_new_tokens=MAX_RESPUESTA,
                                 do_sample=True, temperature=temperatura, top_k=30)
    return tok.decode(salida[0, ids.shape[1]:], skip_special_tokens=True).strip()


demo = gr.ChatInterface(
    responder,
    title="Minerva v0.1 🇪🇸",
    description=("Un modelo de lenguaje en español diminuto (30M de parámetros), entrenado desde cero en un portátil. "
                 "Habla bien pero **se lo inventa casi todo**: es un modelo de juguete. "
                 "[Ver el modelo](https://huggingface.co/AxoneSystem/Minerva-v0.1)"),
    examples=[["¿Quién eres?", 0.7], ["¿Por qué te llamas Minerva?", 0.7], ["¿Qué es la paella?", 0.7],
              ["Cuéntame un cuento corto.", 0.7]],
    additional_inputs=[gr.Slider(0.1, 1.5, value=0.7, step=0.1, label="Temperatura (más alta = más loca)")],
)

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