TeszenAI/MTP-1.1
0
1# ================================================================2# MTP - app.py para Hugging Face Space (Gradio, CPU)3# Carga el checkpoint MTP_MODEL.pt desde el repo TeszenAI/MTP-14# ================================================================5import os6import math7import torch8import torch.nn as nn9import torch.nn.functional as F10import gradio as gr11from starlette.middleware import Middleware12from fastapi.middleware.cors import CORSMiddleware13from pydantic import BaseModel14from typing import Optional15from huggingface_hub import hf_hub_download16 17# ---------------- Optimización para CPU ----------------18# Limita hilos a los núcleos disponibles (evita overhead en Spaces pequeños)19torch.set_num_threads(max(1, os.cpu_count() or 1))20torch.set_grad_enabled(False) # solo inferencia, nunca necesitamos gradientes21 22DEVICE = "cpu"23 24REPO_ID = "TeszenAI/MTP-1-1"25FILENAME = "MTP_MODEL.pt"26 27# ---------------- Arquitectura (idéntica a la de entrenamiento) ----------------28class CausalSelfAttention(nn.Module):29 def __init__(self, n_embd, n_head, block_size, dropout):30 super().__init__()31 self.n_head = n_head32 self.head_dim = n_embd // n_head33 self.qkv = nn.Linear(n_embd, 3 * n_embd)34 self.proj = nn.Linear(n_embd, n_embd)35 self.attn_dropout = nn.Dropout(dropout)36 self.resid_dropout = nn.Dropout(dropout)37 mask = torch.tril(torch.ones(block_size, block_size)).view(1, 1, block_size, block_size)38 self.register_buffer("mask", mask)39 40 def forward(self, x):41 B, T, C = x.shape42 qkv = self.qkv(x)43 q, k, v = qkv.split(C, dim=2)44 q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2)45 k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2)46 v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2)47 att = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)48 att = att.masked_fill(self.mask[:, :, :T, :T] == 0, float("-inf"))49 att = F.softmax(att, dim=-1)50 att = self.attn_dropout(att)51 out = (att @ v).transpose(1, 2).contiguous().view(B, T, C)52 return self.resid_dropout(self.proj(out))53 54 55class FeedForward(nn.Module):56 def __init__(self, n_embd, dropout):57 super().__init__()58 self.net = nn.Sequential(59 nn.Linear(n_embd, 4 * n_embd), nn.GELU(),60 nn.Linear(4 * n_embd, n_embd), nn.Dropout(dropout),61 )62 63 def forward(self, x):64 return self.net(x)65 66 67class Block(nn.Module):68 def __init__(self, n_embd, n_head, block_size, dropout):69 super().__init__()70 self.ln1 = nn.LayerNorm(n_embd)71 self.attn = CausalSelfAttention(n_embd, n_head, block_size, dropout)72 self.ln2 = nn.LayerNorm(n_embd)73 self.ff = FeedForward(n_embd, dropout)74 75 def forward(self, x):76 x = x + self.attn(self.ln1(x))77 x = x + self.ff(self.ln2(x))78 return x79 80 81class MTP(nn.Module):82 def __init__(self, vocab_size, block_size, n_layer, n_head, n_embd, dropout):83 super().__init__()84 self.block_size = block_size85 self.tok_emb = nn.Embedding(vocab_size, n_embd)86 self.pos_emb = nn.Embedding(block_size, n_embd)87 self.drop = nn.Dropout(dropout)88 self.blocks = nn.ModuleList([Block(n_embd, n_head, block_size, dropout) for _ in range(n_layer)])89 self.ln_f = nn.LayerNorm(n_embd)90 self.lm_head = nn.Linear(n_embd, vocab_size, bias=False)91 self.lm_head.weight = self.tok_emb.weight92 93 def forward(self, idx):94 B, T = idx.shape95 pos = torch.arange(T, device=idx.device)96 x = self.tok_emb(idx) + self.pos_emb(pos)97 x = self.drop(x)98 for block in self.blocks:99 x = block(x)100 x = self.ln_f(x)101 return self.lm_head(x)102 103 104# ---------------- Carga del checkpoint (una sola vez, al iniciar el Space) ----------------105print("Descargando checkpoint desde el Hub...")106ckpt_path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)107checkpoint = torch.load(ckpt_path, map_location=DEVICE)108 109cfg = checkpoint["config"]110stoi = checkpoint["stoi"]111itos = {int(k): v for k, v in checkpoint["itos"].items()}112special = checkpoint["special_tokens"]113gen_defaults = checkpoint["generation_defaults"]114 115PAD_ID, BOS_ID, EOS_ID, UNK_ID = special["pad_id"], special["bos_id"], special["eos_id"], special["unk_id"]116 117model = MTP(118 vocab_size=cfg["vocab_size"], block_size=cfg["block_size"],119 n_layer=cfg["n_layer"], n_head=cfg["n_head"],120 n_embd=cfg["n_embd"], dropout=cfg["dropout"],121).to(DEVICE)122model.load_state_dict(checkpoint["model_state_dict"])123model.eval()124 125# fusiona LayerNorm/Linear estáticamente no aplica aquí, pero fija modo eval126# y evita cualquier dropout durante inferencia.127BLOCK_SIZE = cfg["block_size"]128 129print(f"MTP cargado ({checkpoint['meta']['model_name']}, "130 f"entrenado con {checkpoint['meta']['trained_examples']} ejemplos)")131 132 133def encode_text(s):134 return [stoi.get(ch, UNK_ID) for ch in s]135 136 137def decode_ids(ids):138 return "".join(itos.get(i, "") for i in ids if i not in (PAD_ID, BOS_ID, EOS_ID))139 140 141# ---------------- Generación ----------------142@torch.inference_mode()143def generate(idx, max_new_tokens, temperature, top_k, top_p, repetition_penalty):144 for _ in range(max_new_tokens):145 idx_cond = idx[:, -BLOCK_SIZE:]146 logits = model(idx_cond)147 logits = logits[:, -1, :] / max(temperature, 1e-5)148 149 if repetition_penalty and repetition_penalty != 1.0:150 for token_id in set(idx[0].tolist()):151 logits[0, token_id] /= repetition_penalty152 153 if top_k is not None and top_k > 0:154 v, _ = torch.topk(logits, min(top_k, logits.size(-1)))155 logits[logits < v[:, [-1]]] = float("-inf")156 157 probs = F.softmax(logits, dim=-1)158 159 if top_p is not None and 0 < top_p < 1:160 sorted_probs, sorted_idx = torch.sort(probs, descending=True)161 cum_probs = torch.cumsum(sorted_probs, dim=-1)162 cutoff = cum_probs > top_p163 cutoff[:, 1:] = cutoff[:, :-1].clone()164 cutoff[:, 0] = False165 sorted_probs[cutoff] = 0.0166 sorted_probs = sorted_probs / sorted_probs.sum(dim=-1, keepdim=True)167 next_id = sorted_idx.gather(-1, torch.multinomial(sorted_probs, 1))168 else:169 next_id = torch.multinomial(probs, num_samples=1)170 171 idx = torch.cat([idx, next_id], dim=1)172 if next_id.item() == EOS_ID:173 break174 return idx175 176 177def run_inference(text, max_new_tokens=None, temperature=None, top_k=None, top_p=None, repetition_penalty=None):178 """Núcleo de generación, reutilizado por la UI de Gradio y por la API /generate.179 No reduce calidad por estar en CPU: usa exactamente el mismo muestreo180 (top_k + top_p + repetition_penalty) que en la Celda 2 de entrenamiento,181 solo que tarda más en devolver el resultado."""182 max_new_tokens = int(max_new_tokens) if max_new_tokens else gen_defaults["max_new_tokens"]183 temperature = float(temperature) if temperature is not None else gen_defaults["temperature"]184 top_k = int(top_k) if top_k is not None else gen_defaults["top_k"]185 top_p = float(top_p) if top_p is not None else gen_defaults["top_p"]186 repetition_penalty = float(repetition_penalty) if repetition_penalty is not None else gen_defaults["repetition_penalty"]187 188 # Techo máximo de generación: no obliga a generar siempre esto, es solo189 # el límite superior disponible cuando la respuesta realmente lo amerite190 # (el modelo igual corta antes solo con el token <eos> en respuestas cortas).191 # 4000 caracteres ronda el tamaño de una respuesta larga tipo ChatGPT.192 MAX_TOKENS_HARD_LIMIT = 4000193 max_new_tokens = max(1, min(max_new_tokens, MAX_TOKENS_HARD_LIMIT))194 195 prefix = f"Usuario: {text}\nMTP: "196 ids = [BOS_ID] + encode_text(prefix)197 idx = torch.tensor([ids], dtype=torch.long, device=DEVICE)198 199 out = generate(idx, max_new_tokens, temperature, top_k, top_p, repetition_penalty)200 new_ids = out[0].tolist()[len(ids):]201 return decode_ids(new_ids).strip()202 203 204def chat_fn(message, history, max_new_tokens, temperature, top_k, top_p, repetition_penalty):205 return run_inference(message, max_new_tokens, temperature, top_k, top_p, repetition_penalty)206 207 208# ---------------- Interfaz Gradio (para probar el modelo desde el navegador) ----------------209with gr.Blocks(title="MTP Chat") as demo:210 gr.Markdown("# MTP\nModelo GPT entrenado desde cero (char-level). Ejecutándose en CPU.")211 212 with gr.Accordion("Parámetros de generación", open=False):213 max_new_tokens_ui = gr.Slider(16, 4000, value=gen_defaults["max_new_tokens"], step=10, label="max_new_tokens")214 temperature_ui = gr.Slider(0.1, 2.0, value=gen_defaults["temperature"], step=0.05, label="temperature")215 top_k_ui = gr.Slider(0, 100, value=gen_defaults["top_k"], step=1, label="top_k")216 top_p_ui = gr.Slider(0.1, 1.0, value=gen_defaults["top_p"], step=0.05, label="top_p")217 repetition_penalty_ui = gr.Slider(1.0, 2.0, value=gen_defaults["repetition_penalty"], step=0.05,218 label="repetition_penalty")219 220 chatbot = gr.ChatInterface(221 fn=chat_fn,222 additional_inputs=[max_new_tokens_ui, temperature_ui, top_k_ui, top_p_ui, repetition_penalty_ui],223 title=None,224 examples=[225 ["Hola, ¿cómo estás?"],226 ["¿Cuánto es 8 + 5?"],227 ["Explícame qué es un algoritmo."],228 ],229 cache_examples=False,230 )231 232demo.queue(max_size=16)233 234# ---------------- API REST /generate (la que consume el PHP) ----------------235# El PHP hace: fetch(url, { method:'POST', body: JSON.stringify({text, max_tokens, temperature}) })236# y espera de vuelta: { "reply": "..." }237#238# IMPORTANTE:239# - ssr_mode=False: Gradio 6 usa un servidor Node.js aparte para SSR, que240# intentaba levantarse en el puerto 7861 y chocaba. Lo desactivamos porque241# no lo necesitamos para servir la API.242# - El middleware CORS se pasa vía app_kwargs ANTES de llamar a launch(),243# porque una vez que la app arranca, Starlette ya no permite añadir244# middleware (por eso fallaba con app.add_middleware() después).245 246class GenerateRequest(BaseModel):247 text: str248 max_tokens: Optional[int] = None249 temperature: Optional[float] = None250 top_k: Optional[int] = None251 top_p: Optional[float] = None252 repetition_penalty: Optional[float] = None253 254 255PORT = int(os.environ.get("PORT", 7860))256demo.launch(257 server_name="0.0.0.0",258 server_port=PORT,259 prevent_thread_lock=True,260 ssr_mode=False,261 app_kwargs={262 "middleware": [263 Middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"]),264 ]265 },266)267 268app = demo.app269 270 271@app.post("/generate")272def generate_endpoint(req: GenerateRequest):273 if not req.text or not req.text.strip():274 return {"reply": "Escribe algo para que pueda responder."}275 try:276 reply = run_inference(277 req.text,278 max_new_tokens=req.max_tokens,279 temperature=req.temperature,280 top_k=req.top_k,281 top_p=req.top_p,282 repetition_penalty=req.repetition_penalty,283 )284 if not reply:285 reply = "No pude generar una respuesta."286 return {"reply": reply}287 except Exception as e:288 return {"reply": f"Error del modelo: {e}"}289 290 291@app.get("/generate")292def generate_health():293 # Solo para poder comprobar en el navegador que la ruta existe (GET no genera texto)294 return {"status": "ok", "info": "Usa POST con JSON {text, max_tokens, temperature}"}295 296 297# demo.launch(prevent_thread_lock=True) ya dejó el servidor corriendo en un298# hilo en segundo plano (un solo proceso, un solo puerto). Mantenemos vivo299# el hilo principal para que el contenedor del Space no termine.300demo.block_thread()