Upload 2 files

app.py

@@ -0,0 +1,300 @@
+# ================================================================
+# MTP - app.py para Hugging Face Space (Gradio, CPU)
+# Carga el checkpoint MTP_MODEL.pt desde el repo TeszenAI/MTP-1
+# ================================================================
+import os
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import gradio as gr
+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
+
+# ---------------- Optimización para CPU ----------------
+# Limita hilos a los núcleos disponibles (evita overhead en Spaces pequeños)
+torch.set_num_threads(max(1, os.cpu_count() or 1))
+torch.set_grad_enabled(False)  # solo inferencia, nunca necesitamos gradientes
+
+DEVICE = "cpu"
+
+REPO_ID = "TeszenAI/MTP-1.2"
+FILENAME = "MTP_MODEL.pt"
+
+# ---------------- Arquitectura (idéntica a la de entrenamiento) ----------------
+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):
+        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)
+        att = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+        att = att.masked_fill(self.mask[:, :, :T, :T] == 0, float("-inf"))
+        att = F.softmax(att, dim=-1)
+        att = self.attn_dropout(att)
+        out = (att @ v).transpose(1, 2).contiguous().view(B, T, C)
+        return self.resid_dropout(self.proj(out))
+
+
+class FeedForward(nn.Module):
+    def __init__(self, n_embd, dropout):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(n_embd, 4 * n_embd), nn.GELU(),
+            nn.Linear(4 * n_embd, n_embd), nn.Dropout(dropout),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class Block(nn.Module):
+    def __init__(self, n_embd, n_head, block_size, dropout):
+        super().__init__()
+        self.ln1 = nn.LayerNorm(n_embd)
+        self.attn = CausalSelfAttention(n_embd, n_head, block_size, dropout)
+        self.ln2 = nn.LayerNorm(n_embd)
+        self.ff = FeedForward(n_embd, dropout)
+
+    def forward(self, x):
+        x = x + self.attn(self.ln1(x))
+        x = x + self.ff(self.ln2(x))
+        return x
+
+
+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.tok_emb = nn.Embedding(vocab_size, n_embd)
+        self.pos_emb = nn.Embedding(block_size, n_embd)
+        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 = nn.LayerNorm(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):
+        B, T = idx.shape
+        pos = torch.arange(T, device=idx.device)
+        x = self.tok_emb(idx) + self.pos_emb(pos)
+        x = self.drop(x)
+        for block in self.blocks:
+            x = block(x)
+        x = self.ln_f(x)
+        return self.lm_head(x)
+
+
+# ---------------- Carga del checkpoint (una sola vez, al iniciar el Space) ----------------
+print("Descargando checkpoint desde el Hub...")
+ckpt_path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
+checkpoint = torch.load(ckpt_path, map_location=DEVICE)
+
+cfg = checkpoint["config"]
+stoi = checkpoint["stoi"]
+itos = {int(k): v for k, v in checkpoint["itos"].items()}
+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"]
+
+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()
+
+# fusiona LayerNorm/Linear estáticamente no aplica aquí, pero fija modo eval
+# y evita cualquier dropout durante inferencia.
+BLOCK_SIZE = cfg["block_size"]
+
+print(f"MTP cargado ({checkpoint['meta']['model_name']}, "
+      f"entrenado con {checkpoint['meta']['trained_examples']} ejemplos)")
+
+
+def encode_text(s):
+    return [stoi.get(ch, UNK_ID) for ch in s]
+
+
+def decode_ids(ids):
+    return "".join(itos.get(i, "") for i in ids if i not in (PAD_ID, BOS_ID, EOS_ID))
+
+
+# ---------------- Generación ----------------
+@torch.inference_mode()
+def generate(idx, max_new_tokens, temperature, top_k, top_p, repetition_penalty):
+    for _ in range(max_new_tokens):
+        idx_cond = idx[:, -BLOCK_SIZE:]
+        logits = model(idx_cond)
+        logits = logits[:, -1, :] / max(temperature, 1e-5)
+
+        if repetition_penalty and repetition_penalty != 1.0:
+            for token_id in set(idx[0].tolist()):
+                logits[0, token_id] /= repetition_penalty
+
+        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):
+    """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) que en la Celda 2 de entrenamiento,
+    solo que tarda más en devolver el resultado."""
+    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"]
+
+    # Techo máximo de generación: no obliga a generar siempre esto, es solo
+    # el límite superior disponible cuando la respuesta realmente lo amerite
+    # (el modelo igual corta antes solo con el token <eos> en respuestas cortas).
+    # 4000 caracteres ronda el tamaño de una respuesta larga tipo ChatGPT.
+    MAX_TOKENS_HARD_LIMIT = 4000
+    max_new_tokens = max(1, min(max_new_tokens, MAX_TOKENS_HARD_LIMIT))
+
+    prefix = f"Usuario: {text}\nMTP: "
+    ids = [BOS_ID] + encode_text(prefix)
+    idx = torch.tensor([ids], dtype=torch.long, device=DEVICE)
+
+    out = generate(idx, max_new_tokens, temperature, top_k, top_p, repetition_penalty)
+    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 Chat") as demo:
+    gr.Markdown("# MTP\nModelo GPT entrenado desde cero (char-level). Ejecutándose en CPU.")
+
+    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
+
+
+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=["*"], 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,
+        )
+        if not reply:
+            reply = "No pude generar una respuesta."
+        return {"reply": reply}
+    except Exception as e:
+        return {"reply": f"Error del modelo: {e}"}
+
+
+@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}"}
+
+
+# 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()

requirements.txt

@@ -0,0 +1,5 @@
+torch --index-url https://download.pytorch.org/whl/cpu
+gradio
+fastapi
+uvicorn
+huggingface_hub