Create app.py

app.py

@@ -0,0 +1,593 @@
+# ================================================================
+# MTP-2.5 - app.py para Hugging Face Space (Gradio, CPU)
+# Generacion 2: RoPE + SwiGLU (reemplaza position embeddings aprendidos y el
+# GELU-MLP de MTP-1.x). Requiere un checkpoint entrenado con la Celda 1/2 de
+# MTP-2.5; NO carga checkpoints de MTP-2.0 (RMSNorm cambia el state_dict) ni de MTP-1.x.
+#
+# OPTIMIZACIÓN DE VELOCIDAD (sin tocar el resto de la logica de muestreo):
+#   - KV-cache en la atención: en generación autoregresiva, cada paso
+#     antes recomputaba TODO el contexto desde cero (O(n^2) en total).
+#     Ahora se reutiliza lo ya calculado y solo se procesa el token
+#     nuevo (O(n) en total). Es el mismo cálculo matemático, solo que
+#     no se repite trabajo ya hecho.
+#   - F.scaled_dot_product_attention: kernel fusionado de PyTorch,
+#     mismo resultado que el softmax manual pero más rápido en CPU.
+#     Si la versión de PyTorch no lo trae, cae automáticamente al
+#     cálculo manual (fallback), así que no se rompe en ningún entorno.
+#   - repetition_penalty vectorizado + bloqueo de n-gramas repetidos
+#     (evita que la respuesta final copie literalmente un fragmento ya
+#     generado, sin que esto sea un "modelo de n-gramas": el modelo que
+#     predice sigue siendo 100% transformer).
+# ================================================================
+import os
+import math
+import time
+import logging
+import threading
+import traceback
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import gradio as gr
+import sentencepiece as spm
+from starlette.middleware import Middleware
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+from typing import Optional
+from huggingface_hub import hf_hub_download
+
+# ---------------- Logging ----------------
+# print() se pierde facil entre el ruido de arranque de Gradio/Starlette en
+# los logs de un Space; con logging queda todo con timestamp y nivel, y es
+# mas facil de filtrar si algo sale mal en produccion.
+logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
+logger = logging.getLogger("mtp")
+
+# ---------------- Dispositivo ----------------
+# Autodetecta GPU si el Space corre en un tier con GPU (DEVICE=cuda por
+# variable de entorno tambien fuerza el valor si hace falta). Si no hay
+# GPU disponible, cae a CPU como siempre.
+DEVICE = os.environ.get("DEVICE") or ("cuda" if torch.cuda.is_available() else "cpu")
+
+if DEVICE == "cpu":
+    # Limita hilos a los núcleos disponibles (evita overhead en Spaces pequeños).
+    # No tiene sentido en GPU, donde el computo no lo hace la CPU.
+    torch.set_num_threads(max(1, os.cpu_count() or 1))
+    # set_num_interop_threads solo puede llamarse una vez y antes de cualquier
+    # operación paralela; lo protegemos por si el entorno ya lo fijó.
+    try:
+        torch.set_num_interop_threads(1)
+    except RuntimeError:
+        pass
+
+torch.set_grad_enabled(False)  # solo inferencia, nunca necesitamos gradientes
+
+# Disponibilidad de scaled_dot_product_attention (PyTorch >= 2.0).
+# Si no está disponible, usamos el softmax manual original como fallback.
+_HAS_SDPA = hasattr(F, "scaled_dot_product_attention")
+
+REPO_ID = os.environ.get("MTP_REPO_ID", "TeszenAI/MTP-2.7")  # <-- ajusta al nombre real de tu repo/Space en el Hub
+FILENAME = os.environ.get("MTP_FILENAME", "MTP2_7_MODEL.pt")
+
+# Origenes permitidos para CORS. Por defecto "*" (como antes), pero se puede
+# restringir en produccion con la variable de entorno MTP_ALLOWED_ORIGINS
+# (separados por coma), por ejemplo: "https://teszen.com,https://www.teszen.com"
+ALLOWED_ORIGINS = [
+    o.strip() for o in os.environ.get("MTP_ALLOWED_ORIGINS", "*").split(",") if o.strip()
+] or ["*"]
+
+# Techo de caracteres del input antes de tokenizar. No es por seguridad (el
+# modelo igual solo "ve" los ultimos BLOCK_SIZE tokens), es para no perder
+# tiempo tokenizando un texto absurdamente largo por error o abuso.
+MAX_INPUT_CHARS = 4000
+
+# Si es "1", los errores devueltos por /generate incluyen el detalle interno
+# de la excepcion (util mientras desarrollas). En produccion, dejar en "0"
+# para no filtrarle al cliente detalles internos del servidor.
+DEBUG_ERRORS = os.environ.get("MTP_DEBUG_ERRORS", "0") == "1"
+
+# Serializa las llamadas a generate(): sin esto, dos requests concurrentes
+# (por ejemplo la UI de Gradio y el endpoint /generate al mismo tiempo, o
+# varias visitas simultaneas al sitio) compiten por los mismos hilos de CPU
+# y todas terminan mas lentas en vez de una rapida y la otra esperando. Con
+# el lock, cada generacion corre de punta a punta antes de que empiece la
+# siguiente -- mismo comportamiento de fondo que antes bajo carga baja, pero
+# estable bajo carga alta en vez de degradarse.
+_generation_lock = threading.Lock()
+
+# ---------------- Arquitectura MTP-2.x: RoPE + SwiGLU, con KV-cache ----------------
+def rotate_half(x):
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat((-x2, x1), dim=-1)
+
+def apply_rope(q, k, cos, sin):
+    cos = cos.unsqueeze(0).unsqueeze(0)
+    sin = sin.unsqueeze(0).unsqueeze(0)
+    q_rot = (q * cos) + (rotate_half(q) * sin)
+    k_rot = (k * cos) + (rotate_half(k) * sin)
+    return q_rot, k_rot
+
+class RotaryEmbedding(nn.Module):
+    def __init__(self, head_dim, max_seq_len, base=10000):
+        super().__init__()
+        inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self._build_cache(max_seq_len)
+
+    def _build_cache(self, seq_len):
+        t = torch.arange(seq_len, dtype=self.inv_freq.dtype, device=self.inv_freq.device)
+        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.register_buffer("cos_cached", emb.cos(), persistent=False)
+        self.register_buffer("sin_cached", emb.sin(), persistent=False)
+        self.max_seq_len_cached = seq_len
+
+    def forward(self, seq_len, device, dtype, offset=0):
+        # Con KV-cache, `offset` es cuantos tokens ya estan en la cache: el
+        # token nuevo necesita el angulo correspondiente a SU posicion
+        # absoluta, no a la posicion relativa dentro de este forward.
+        if offset + seq_len > self.max_seq_len_cached:
+            self._build_cache(offset + seq_len)
+        cos = self.cos_cached[offset:offset + seq_len].to(device=device, dtype=dtype)
+        sin = self.sin_cached[offset:offset + seq_len].to(device=device, dtype=dtype)
+        return cos, sin
+
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, n_embd, n_head, block_size, dropout):
+        super().__init__()
+        self.n_head = n_head
+        self.head_dim = n_embd // n_head
+        self.qkv = nn.Linear(n_embd, 3 * n_embd)
+        self.proj = nn.Linear(n_embd, n_embd)
+        self.attn_dropout = nn.Dropout(dropout)
+        self.resid_dropout = nn.Dropout(dropout)
+        mask = torch.tril(torch.ones(block_size, block_size)).view(1, 1, block_size, block_size)
+        self.register_buffer("mask", mask)
+
+    def forward(self, x, cos, sin, past_kv=None, use_cache=False):
+        B, T, C = x.shape
+        qkv = self.qkv(x)
+        q, k, v = qkv.split(C, dim=2)
+        q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
+        k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
+        v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
+
+        # RoPE se aplica ANTES de guardar en cache, con el angulo absoluto de
+        # cada token (pasado por `cos`/`sin`, ya calculado con el offset
+        # correcto en MTP.forward). Asi el k cacheado ya trae rotada su
+        # posicion real y no hay que re-rotar nada en pasos futuros.
+        q, k = apply_rope(q, k, cos, sin)
+
+        if past_kv is not None:
+            past_k, past_v = past_kv
+            k = torch.cat([past_k, k], dim=2)
+            v = torch.cat([past_v, v], dim=2)
+
+        present_kv = (k, v) if use_cache else None
+
+        is_causal = (past_kv is None) and (T > 1)
+
+        if _HAS_SDPA:
+            out = F.scaled_dot_product_attention(
+                q, k, v, attn_mask=None, dropout_p=0.0, is_causal=is_causal,
+            )
+        else:
+            att = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+            if is_causal:
+                Tk = k.size(-2)
+                causal_mask = torch.tril(torch.ones(T, Tk, device=x.device, dtype=torch.bool))
+                att = att.masked_fill(~causal_mask, float("-inf"))
+            att = F.softmax(att, dim=-1)
+            att = self.attn_dropout(att)
+            out = att @ v
+
+        out = out.transpose(1, 2).contiguous().view(B, T, C)
+        out = self.resid_dropout(self.proj(out))
+        return out, present_kv
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, n_embd, dropout):
+        super().__init__()
+        hidden = int(2 * (4 * n_embd) / 3)
+        hidden = ((hidden + 7) // 8) * 8
+        self.w_gate = nn.Linear(n_embd, hidden, bias=False)
+        self.w_up = nn.Linear(n_embd, hidden, bias=False)
+        self.w_down = nn.Linear(hidden, n_embd, bias=False)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        return self.dropout(self.w_down(F.silu(self.w_gate(x)) * self.w_up(x)))
+
+
+class RMSNorm(nn.Module):
+    """Debe coincidir exactamente con la version de entrenamiento. No
+    necesita ningun cambio para funcionar con KV-cache: normaliza cada
+    posicion de forma independiente, igual que LayerNorm."""
+    def __init__(self, dim, eps=1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        norm = x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
+        return norm * self.weight
+
+
+class Block(nn.Module):
+    def __init__(self, n_embd, n_head, block_size, dropout):
+        super().__init__()
+        self.ln1 = RMSNorm(n_embd)
+        self.attn = CausalSelfAttention(n_embd, n_head, block_size, dropout)
+        self.ln2 = RMSNorm(n_embd)
+        self.ff = SwiGLU(n_embd, dropout)
+
+    def forward(self, x, cos, sin, past_kv=None, use_cache=False):
+        attn_out, present_kv = self.attn(self.ln1(x), cos, sin, past_kv=past_kv, use_cache=use_cache)
+        x = x + attn_out
+        x = x + self.ff(self.ln2(x))
+        return x, present_kv
+
+
+class MTP(nn.Module):
+    def __init__(self, vocab_size, block_size, n_layer, n_head, n_embd, dropout):
+        super().__init__()
+        self.block_size = block_size
+        self.head_dim = n_embd // n_head
+        self.tok_emb = nn.Embedding(vocab_size, n_embd)
+        self.rope = RotaryEmbedding(self.head_dim, max_seq_len=block_size)
+        self.drop = nn.Dropout(dropout)
+        self.blocks = nn.ModuleList([Block(n_embd, n_head, block_size, dropout) for _ in range(n_layer)])
+        self.ln_f = RMSNorm(n_embd)
+        self.lm_head = nn.Linear(n_embd, vocab_size, bias=False)
+        self.lm_head.weight = self.tok_emb.weight
+
+    def forward(self, idx, past_key_values=None, use_cache=False, pos_offset=0):
+        B, T = idx.shape
+        x = self.tok_emb(idx)
+        x = self.drop(x)
+        cos, sin = self.rope(T, idx.device, x.dtype, offset=pos_offset)
+
+        new_past = [] if use_cache else None
+        for i, block in enumerate(self.blocks):
+            past_kv = past_key_values[i] if past_key_values is not None else None
+            x, present_kv = block(x, cos, sin, past_kv=past_kv, use_cache=use_cache)
+            if use_cache:
+                new_past.append(present_kv)
+
+        x = self.ln_f(x)
+        logits = self.lm_head(x)
+        return logits, new_past
+
+
+# ---------------- Carga del checkpoint (una sola vez, al iniciar el Space) ----------------
+def _download_checkpoint_with_retries(repo_id, filename, max_retries=5):
+    """Reintenta la descarga con backoff. hf_hub_download puede fallar por un
+    corte de red momentaneo o un problema del backend Xet de HF; sin esto,
+    un solo fallo transitorio tira abajo el arranque completo del Space."""
+    last_err = None
+    for attempt in range(1, max_retries + 1):
+        try:
+            logger.info(f"Descargando checkpoint desde el Hub (intento {attempt}/{max_retries})...")
+            return hf_hub_download(repo_id=repo_id, filename=filename)
+        except Exception as e:
+            last_err = e
+            logger.warning(f"Fallo la descarga del checkpoint: {e}")
+            if attempt < max_retries:
+                time.sleep(5 * attempt)
+    raise RuntimeError(f"No se pudo descargar el checkpoint tras {max_retries} intentos") from last_err
+
+ckpt_path = _download_checkpoint_with_retries(REPO_ID, FILENAME)
+checkpoint = torch.load(ckpt_path, map_location=DEVICE)
+
+cfg = checkpoint["config"]
+special = checkpoint["special_tokens"]
+gen_defaults = checkpoint["generation_defaults"]
+
+PAD_ID, BOS_ID, EOS_ID, UNK_ID = special["pad_id"], special["bos_id"], special["eos_id"], special["unk_id"]
+
+# El tokenizer es BPE (SentencePiece) entrenado desde cero junto con el modelo.
+# No es un modelo preentrenado externo: viene embebido como bytes dentro del
+# mismo checkpoint que los pesos. Se carga directo desde memoria con
+# load_from_serialized_proto, sin necesidad de escribirlo a disco primero.
+sp = spm.SentencePieceProcessor()
+sp.load_from_serialized_proto(checkpoint["spm_model_bytes"])
+
+model = MTP(
+    vocab_size=cfg["vocab_size"], block_size=cfg["block_size"],
+    n_layer=cfg["n_layer"], n_head=cfg["n_head"],
+    n_embd=cfg["n_embd"], dropout=cfg["dropout"],
+).to(DEVICE)
+model.load_state_dict(checkpoint["model_state_dict"])
+model.eval()
+
+BLOCK_SIZE = cfg["block_size"]
+
+logger.info(
+    f"MTP cargado ({checkpoint['meta']['model_name']}, "
+    f"entrenado con {checkpoint['meta']['trained_examples']} ejemplos) "
+    f"| device={DEVICE} | SDPA={'sí' if _HAS_SDPA else 'no (fallback manual)'}"
+)
+
+
+import re as _re_indent
+
+def protect_indentation(text):
+    """Debe coincidir exactamente con la funcion usada en el entrenamiento."""
+    lines = text.split("\n")
+    new_lines = []
+    for line in lines:
+        stripped = line.lstrip(" ")
+        n_spaces = len(line) - len(stripped)
+        n_levels = n_spaces // 4
+        remainder = n_spaces % 4
+        if n_levels > 0:
+            prefix = " " + " ".join(["<tab>"] * n_levels) + " " + " " * remainder
+        else:
+            prefix = " " * remainder
+        new_lines.append(prefix + stripped)
+    text = "\n".join(new_lines)
+    def _repl(m):
+        n = len(m.group())
+        return " " + " ".join(["<nl>"] * n) + " "
+    text = _re_indent.sub(r"\n+", _repl, text)
+    return text
+
+
+def restore_indentation(text):
+    text = _re_indent.sub(r"(<nl>\s*)+", lambda m: "\n" * m.group().count("<nl>"), text)
+    text = _re_indent.sub(r"(<tab>\s*)+", lambda m: "    " * m.group().count("<tab>"), text)
+    return text
+
+
+def encode_text(s):
+    return sp.encode(protect_indentation(s), out_type=int)
+
+
+def decode_ids(ids):
+    text = sp.decode([i for i in ids if i not in (PAD_ID, BOS_ID, EOS_ID)])
+    return restore_indentation(text)
+
+
+def _block_repeated_ngrams(generated_ids, logits, ngram_size):
+    """Prohibe repetir literalmente un n-grama ya generado en esta misma
+    respuesta (tecnica de decoding tipo GPT-2/3, no un modelo de n-gramas:
+    el modelo que predice sigue siendo 100% transformer con KV-cache)."""
+    if ngram_size <= 0 or len(generated_ids) < ngram_size:
+        return logits
+    prefix = tuple(generated_ids[-(ngram_size - 1):])
+    banned = set()
+    for i in range(len(generated_ids) - ngram_size + 1):
+        if tuple(generated_ids[i:i + ngram_size - 1]) == prefix:
+            banned.add(generated_ids[i + ngram_size - 1])
+    if banned:
+        logits[0, list(banned)] = float("-inf")
+    return logits
+
+
+# ---------------- Generación (con KV-cache) ----------------
+@torch.inference_mode()
+def generate(idx, max_new_tokens, temperature, top_k, top_p, repetition_penalty, no_repeat_ngram_size=3):
+    past_key_values = None
+    cache_len = 0  # cuántos tokens del extremo derecho de `idx` ya están en la caché
+
+    for _ in range(max_new_tokens):
+        total_len = idx.shape[1]
+
+        if total_len <= BLOCK_SIZE:
+            if past_key_values is None:
+                # Primer paso: una sola pasada ("prefill") sobre todo el prompt.
+                logits, past_key_values = model(idx, use_cache=True)
+                cache_len = total_len
+            else:
+                # Pasos siguientes: solo se procesa el último token generado,
+                # reutilizando la caché de todo lo anterior.
+                last_token = idx[:, -1:]
+                logits, past_key_values = model(
+                    last_token,
+                    past_key_values=past_key_values,
+                    use_cache=True,
+                    pos_offset=cache_len,
+                )
+                cache_len += 1
+            logits = logits[:, -1, :]
+        else:
+            # Se superó block_size: mismo comportamiento que el modelo original
+            # (ventana deslizante recalculada por completo). Solo ocurre en
+            # respuestas muy largas; la caché se reinicia para esa ventana.
+            idx_cond = idx[:, -BLOCK_SIZE:]
+            logits, past_key_values = model(idx_cond, use_cache=True)
+            cache_len = BLOCK_SIZE
+            logits = logits[:, -1, :]
+
+        logits = logits / max(temperature, 1e-5)
+
+        if repetition_penalty and repetition_penalty != 1.0:
+            # Vectorizado: antes era `for token_id in set(idx[0].tolist())`,
+            # un bucle Python nuevo por cada token generado.
+            unique_ids = torch.unique(idx[0])
+            logits[0, unique_ids] /= repetition_penalty
+
+        logits = _block_repeated_ngrams(idx[0].tolist(), logits, no_repeat_ngram_size)
+
+        if top_k is not None and top_k > 0:
+            v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+            logits[logits < v[:, [-1]]] = float("-inf")
+
+        probs = F.softmax(logits, dim=-1)
+
+        if top_p is not None and 0 < top_p < 1:
+            sorted_probs, sorted_idx = torch.sort(probs, descending=True)
+            cum_probs = torch.cumsum(sorted_probs, dim=-1)
+            cutoff = cum_probs > top_p
+            cutoff[:, 1:] = cutoff[:, :-1].clone()
+            cutoff[:, 0] = False
+            sorted_probs[cutoff] = 0.0
+            sorted_probs = sorted_probs / sorted_probs.sum(dim=-1, keepdim=True)
+            next_id = sorted_idx.gather(-1, torch.multinomial(sorted_probs, 1))
+        else:
+            next_id = torch.multinomial(probs, num_samples=1)
+
+        idx = torch.cat([idx, next_id], dim=1)
+        if next_id.item() == EOS_ID:
+            break
+
+    return idx
+
+
+def run_inference(text, max_new_tokens=None, temperature=None, top_k=None, top_p=None, repetition_penalty=None, no_repeat_ngram_size=None):
+    """Núcleo de generación, reutilizado por la UI de Gradio y por la API /generate.
+    No reduce calidad por estar en CPU: usa exactamente el mismo muestreo
+    (top_k + top_p + repetition_penalty + bloqueo de n-gramas repetidos) que
+    en la Celda 2 de entrenamiento, solo que ahora con KV-cache es notablemente
+    más rápido en respuestas largas."""
+    max_new_tokens = int(max_new_tokens) if max_new_tokens else gen_defaults["max_new_tokens"]
+    temperature = float(temperature) if temperature is not None else gen_defaults["temperature"]
+    top_k = int(top_k) if top_k is not None else gen_defaults["top_k"]
+    top_p = float(top_p) if top_p is not None else gen_defaults["top_p"]
+    repetition_penalty = float(repetition_penalty) if repetition_penalty is not None else gen_defaults["repetition_penalty"]
+    no_repeat_ngram_size = int(no_repeat_ngram_size) if no_repeat_ngram_size is not None else gen_defaults.get("no_repeat_ngram_size", 3)
+
+    # Techo máximo de generación: 4000 no era realista en CPU (cada token
+    # adicional cuesta tiempo real). 700 sigue siendo una respuesta larga y
+    # mantiene el tiempo de respuesta bajo control en el peor caso.
+    MAX_TOKENS_HARD_LIMIT = 700
+    max_new_tokens = max(1, min(max_new_tokens, MAX_TOKENS_HARD_LIMIT))
+
+    if len(text) > MAX_INPUT_CHARS:
+        logger.warning(f"Input de {len(text)} caracteres recortado a {MAX_INPUT_CHARS}")
+        text = text[:MAX_INPUT_CHARS]
+
+    prefix = f"Usuario: {text}\nMTP: "
+    ids = [BOS_ID] + encode_text(prefix)
+    idx = torch.tensor([ids], dtype=torch.long, device=DEVICE)
+
+    # Con el lock, si llegan varias generaciones al mismo tiempo (UI + API,
+    # o varios usuarios a la vez) se procesan una despues de otra en vez de
+    # pisarse los hilos de CPU entre si.
+    with _generation_lock:
+        out = generate(idx, max_new_tokens, temperature, top_k, top_p, repetition_penalty, no_repeat_ngram_size)
+
+    new_ids = out[0].tolist()[len(ids):]
+    return decode_ids(new_ids).strip()
+
+
+def chat_fn(message, history, max_new_tokens, temperature, top_k, top_p, repetition_penalty):
+    return run_inference(message, max_new_tokens, temperature, top_k, top_p, repetition_penalty)
+
+
+# ---------------- Interfaz Gradio (para probar el modelo desde el navegador) ----------------
+with gr.Blocks(title="MTP-2.5 Chat") as demo:
+    gr.Markdown(f"# MTP-2.5\nModelo GPT (RoPE + SwiGLU + RMSNorm) entrenado desde cero, tokenizer BPE. Ejecutándose en {DEVICE.upper()}.")
+
+    with gr.Accordion("Parámetros de generación", open=False):
+        max_new_tokens_ui = gr.Slider(16, 4000, value=gen_defaults["max_new_tokens"], step=10, label="max_new_tokens")
+        temperature_ui = gr.Slider(0.1, 2.0, value=gen_defaults["temperature"], step=0.05, label="temperature")
+        top_k_ui = gr.Slider(0, 100, value=gen_defaults["top_k"], step=1, label="top_k")
+        top_p_ui = gr.Slider(0.1, 1.0, value=gen_defaults["top_p"], step=0.05, label="top_p")
+        repetition_penalty_ui = gr.Slider(1.0, 2.0, value=gen_defaults["repetition_penalty"], step=0.05,
+                                           label="repetition_penalty")
+
+    chatbot = gr.ChatInterface(
+        fn=chat_fn,
+        additional_inputs=[max_new_tokens_ui, temperature_ui, top_k_ui, top_p_ui, repetition_penalty_ui],
+        title=None,
+        examples=[
+            ["Hola, ¿cómo estás?"],
+            ["¿Cuánto es 8 + 5?"],
+            ["Explícame qué es un algoritmo."],
+        ],
+        cache_examples=False,
+    )
+
+demo.queue(max_size=16)
+
+# ---------------- API REST /generate (la que consume el PHP) ----------------
+# El PHP hace: fetch(url, { method:'POST', body: JSON.stringify({text, max_tokens, temperature}) })
+# y espera de vuelta: { "reply": "..." }
+#
+# IMPORTANTE:
+# - ssr_mode=False: Gradio 6 usa un servidor Node.js aparte para SSR, que
+#   intentaba levantarse en el puerto 7861 y chocaba. Lo desactivamos porque
+#   no lo necesitamos para servir la API.
+# - El middleware CORS se pasa vía app_kwargs ANTES de llamar a launch(),
+#   porque una vez que la app arranca, Starlette ya no permite añadir
+#   middleware (por eso fallaba con app.add_middleware() después).
+
+class GenerateRequest(BaseModel):
+    text: str
+    max_tokens: Optional[int] = None
+    temperature: Optional[float] = None
+    top_k: Optional[int] = None
+    top_p: Optional[float] = None
+    repetition_penalty: Optional[float] = None
+    no_repeat_ngram_size: Optional[int] = None
+
+
+PORT = int(os.environ.get("PORT", 7860))
+demo.launch(
+    server_name="0.0.0.0",
+    server_port=PORT,
+    prevent_thread_lock=True,
+    ssr_mode=False,
+    app_kwargs={
+        "middleware": [
+            Middleware(CORSMiddleware, allow_origins=ALLOWED_ORIGINS, allow_methods=["*"], allow_headers=["*"]),
+        ]
+    },
+)
+
+app = demo.app
+
+
+@app.post("/generate")
+def generate_endpoint(req: GenerateRequest):
+    if not req.text or not req.text.strip():
+        return {"reply": "Escribe algo para que pueda responder."}
+    try:
+        reply = run_inference(
+            req.text,
+            max_new_tokens=req.max_tokens,
+            temperature=req.temperature,
+            top_k=req.top_k,
+            top_p=req.top_p,
+            repetition_penalty=req.repetition_penalty,
+            no_repeat_ngram_size=req.no_repeat_ngram_size,
+        )
+        if not reply:
+            reply = "No pude generar una respuesta."
+        return {"reply": reply}
+    except Exception as e:
+        # El detalle completo va al log del Space (con traceback), no a la
+        # respuesta publica: devolver la excepcion cruda a quien llama podria
+        # filtrar rutas internas, nombres de variables, etc. Con
+        # MTP_DEBUG_ERRORS=1 se puede activar el detalle mientras desarrollas.
+        logger.error(f"Error generando respuesta: {e}\n{traceback.format_exc()}")
+        reply = f"Error del modelo: {e}" if DEBUG_ERRORS else "Ocurrió un error al generar la respuesta. Intenta de nuevo en un momento."
+        return {"reply": reply}
+
+
+@app.get("/generate")
+def generate_health():
+    # Solo para poder comprobar en el navegador que la ruta existe (GET no genera texto)
+    return {"status": "ok", "info": "Usa POST con JSON {text, max_tokens, temperature}"}
+
+
+@app.get("/health")
+def health():
+    # Health check real: confirma que el modelo esta cargado y listo, no
+    # solo que el proceso esta vivo. Util para monitoreo externo (uptime
+    # checks) o para que el PHP sepa si conviene reintentar mas tarde.
+    return {
+        "status": "ok",
+        "model": checkpoint["meta"]["model_name"],
+        "device": DEVICE,
+        "trained_examples": checkpoint["meta"]["trained_examples"],
+    }
+
+
+# demo.launch(prevent_thread_lock=True) ya dejó el servidor corriendo en un
+# hilo en segundo plano (un solo proceso, un solo puerto). Mantenemos vivo
+# el hilo principal para que el contenedor del Space no termine.
+demo.block_thread()