Spaces:
Sleeping
Sleeping
| # ================================================================ | |
| # 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 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 | |
| # ---------------- 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)) | |
| # 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 | |
| DEVICE = "cpu" | |
| # 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 = "TeszenAI/MTP-2" # <-- ajusta al nombre real de tu repo/Space en el Hub | |
| FILENAME = "MTP2_5_MODEL.pt" | |
| # ---------------- 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) ---------------- | |
| 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"] | |
| 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"] | |
| print(f"MTP cargado ({checkpoint['meta']['model_name']}, " | |
| f"entrenado con {checkpoint['meta']['trained_examples']} ejemplos)" | |
| f" | 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) ---------------- | |
| 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)) | |
| 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, 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("# MTP-2.5\nModelo GPT (RoPE + SwiGLU + RMSNorm) entrenado desde cero, tokenizer BPE. 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 | |
| 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=["*"], allow_methods=["*"], allow_headers=["*"]), | |
| ] | |
| }, | |
| ) | |
| app = demo.app | |
| 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: | |
| return {"reply": f"Error del modelo: {e}"} | |
| 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() |