TeszenAI/MTP-3
0
1import os2import sys3import torch4import pickle5import time6import gc7from fastapi import FastAPI, Request8from fastapi.responses import HTMLResponse, StreamingResponse9from fastapi.middleware.cors import CORSMiddleware10from pydantic import BaseModel, Field11from huggingface_hub import snapshot_download12import uvicorn13 14# ======================15# CONFIGURACIÓN DE DISPOSITIVO16# ======================17if torch.cuda.is_available():18 DEVICE = "cuda"19 print("✅ GPU NVIDIA detectada. Usando CUDA.")20else:21 DEVICE = "cpu"22 print("⚠️ GPU no detectada. Usando CPU (puede ser más lento).")23 24# Optimización de hilos para CPU25if DEVICE == "cpu":26 torch.set_num_threads(max(1, os.cpu_count() // 2))27 28torch.set_grad_enabled(False)29 30MODEL_REPO = "TeszenAI/mtp-3.1"31 32# ======================33# DESCARGA Y CARGA DEL MODELO34# ======================35print(f"📦 Descargando modelo desde {MODEL_REPO}...")36repo_path = snapshot_download(37 repo_id=MODEL_REPO,38 repo_type="model",39 local_dir="mtptz_repo"40)41 42sys.path.insert(0, repo_path)43 44# Importar modelo mejorado compatible45from model import MTPMiniModel46from tokenizer import MTPTokenizer47 48print("🔧 Cargando tensores y configuración...")49with open(os.path.join(repo_path, "mtp_mini.pkl"), "rb") as f:50 model_data = pickle.load(f)51 52tokenizer = MTPTokenizer(os.path.join(repo_path, "mtp_tokenizer.model"))53VOCAB_SIZE = tokenizer.sp.get_piece_size()54config = model_data["config"]55 56# Detectar si el modelo usa SwiGLU57use_swiglu = config["model"].get("use_swiglu", False)58 59print(f"🧠 Inicializando modelo...")60print(f" → Vocabulario: {VOCAB_SIZE}")61print(f" → Dimensión: {config['model']['d_model']}")62print(f" → Capas: {config['model']['n_layers']}")63print(f" → SwiGLU: {'✓' if use_swiglu else '✗'}")64 65model = MTPMiniModel(66 vocab_size=VOCAB_SIZE,67 d_model=config["model"]["d_model"],68 n_layers=config["model"]["n_layers"],69 n_heads=config["model"]["n_heads"],70 d_ff=config["model"]["d_ff"],71 max_seq_len=config["model"]["max_seq_len"],72 dropout=0.0,73 use_swiglu=use_swiglu # NUEVO: soporte para SwiGLU74)75 76model.load_state_dict(model_data["model_state_dict"])77model.eval()78 79# Cuantización para CPU80if DEVICE == "cpu":81 print("⚡ Aplicando cuantización dinámica para CPU...")82 model = torch.quantization.quantize_dynamic(83 model, 84 {torch.nn.Linear}, 85 dtype=torch.qint886 )87 88model.to(DEVICE)89 90param_count = sum(p.numel() for p in model.parameters())91print(f"✅ Modelo cargado: {param_count:,} parámetros ({param_count/1e6:.1f}M)")92 93# ======================94# API CONFIG95# ======================96app = FastAPI(97 title="MTP-3 API",98 description="API para modelo de lenguaje MTP-3 mejorado",99 version="3.0"100)101 102app.add_middleware(103 CORSMiddleware,104 allow_origins=["*"],105 allow_methods=["*"],106 allow_headers=["*"],107)108 109class PromptRequest(BaseModel):110 text: str = Field(..., max_length=2000, description="Texto de entrada")111 max_tokens: int = Field(default=150, ge=10, le=300, description="Tokens máximos a generar")112 temperature: float = Field(default=0.7, ge=0.1, le=2.0, description="Temperatura de muestreo")113 top_k: int = Field(default=50, ge=1, le=100, description="Top-k sampling")114 top_p: float = Field(default=0.9, ge=0.1, le=1.0, description="Top-p (nucleus) sampling")115 repetition_penalty: float = Field(default=1.1, ge=1.0, le=2.0, description="Penalización por repetición")116 117def build_prompt(user_input: str) -> str:118 """Construye el prompt en el formato del modelo"""119 return f"### Instrucción:\n{user_input}\n\n### Respuesta:\n"120 121# ======================122# ⚡ GESTIÓN DE CARGA123# ======================124ACTIVE_REQUESTS = 0125 126@app.post("/generate")127async def generate(req: PromptRequest):128 """Endpoint principal de generación de texto"""129 global ACTIVE_REQUESTS130 ACTIVE_REQUESTS += 1131 132 # Ajuste dinámico bajo carga133 dyn_max_tokens = req.max_tokens134 dyn_temperature = req.temperature135 136 if ACTIVE_REQUESTS > 2:137 print(f"⚠️ Carga alta ({ACTIVE_REQUESTS} requests). Ajustando parámetros.")138 dyn_max_tokens = min(dyn_max_tokens, 120)139 dyn_temperature = max(0.5, dyn_temperature * 0.9)140 141 user_input = req.text.strip()142 if not user_input:143 ACTIVE_REQUESTS -= 1144 return {"reply": "", "tokens_generated": 0}145 146 full_prompt = build_prompt(user_input)147 tokens = [tokenizer.bos_id()] + tokenizer.encode(full_prompt)148 input_ids = torch.tensor([tokens], device=DEVICE)149 150 try:151 with torch.no_grad():152 output_ids = model.generate(153 input_ids,154 max_new_tokens=dyn_max_tokens,155 temperature=dyn_temperature,156 top_k=req.top_k,157 top_p=req.top_p,158 repetition_penalty=req.repetition_penalty159 )160 161 gen_tokens = output_ids[0, len(tokens):].tolist()162 163 # Filtro de seguridad164 safe_tokens = [165 t for t in gen_tokens 166 if 0 <= t < VOCAB_SIZE and t != tokenizer.eos_id()167 ]168 169 response = tokenizer.decode(safe_tokens).strip()170 171 # Limpiar marcadores de sección172 if "###" in response:173 response = response.split("###")[0].strip()174 175 return {176 "reply": response,177 "tokens_generated": len(safe_tokens),178 "model": "MTP-3"179 }180 181 except Exception as e:182 print(f"❌ Error durante generación: {e}")183 return {184 "reply": "Lo siento, ocurrió un error al procesar tu solicitud.",185 "error": str(e)186 }187 188 finally:189 ACTIVE_REQUESTS -= 1190 if DEVICE == "cuda":191 torch.cuda.empty_cache()192 gc.collect()193 194# ======================195# 📡 STREAMING SSE196# ======================197@app.get("/generate_sse")198def generate_sse(199 text: str,200 max_tokens: int = 150,201 temperature: float = 0.7202):203 """Endpoint de streaming con Server-Sent Events"""204 global ACTIVE_REQUESTS205 ACTIVE_REQUESTS += 1206 207 def event_stream():208 try:209 full_prompt = build_prompt(text)210 tokens = [tokenizer.bos_id()] + tokenizer.encode(full_prompt)211 input_ids = torch.tensor([tokens], device=DEVICE)212 213 # Ajuste dinámico214 limit = 100 if ACTIVE_REQUESTS > 2 else max_tokens215 temp = max(0.5, temperature * 0.9) if ACTIVE_REQUESTS > 2 else temperature216 217 for step in range(limit):218 with torch.no_grad():219 logits, _ = model(input_ids)220 logits = logits[:, -1, :VOCAB_SIZE]221 222 # Sampling con temperatura223 probs = torch.softmax(logits / temp, dim=-1)224 next_id = torch.multinomial(probs, num_samples=1).item()225 226 if next_id == tokenizer.eos_id():227 break228 229 if 0 <= next_id < VOCAB_SIZE:230 token_text = tokenizer.decode([next_id])231 232 # Limpiar salida233 if "###" in token_text:234 break235 236 yield f"data:{token_text}\n\n"237 238 input_ids = torch.cat(239 [input_ids, torch.tensor([[next_id]], device=DEVICE)],240 dim=1241 )242 time.sleep(0.01)243 244 yield "data:[DONE]\n\n"245 246 except Exception as e:247 yield f"data:[ERROR: {str(e)}]\n\n"248 249 finally:250 ACTIVE_REQUESTS -= 1251 if DEVICE == "cuda":252 torch.cuda.empty_cache()253 254 return StreamingResponse(event_stream(), media_type="text/event-stream")255 256# ======================257# 📊 ENDPOINTS DE INFORMACIÓN258# ======================259@app.get("/health")260def health_check():261 """Check del estado del servicio"""262 return {263 "status": "healthy",264 "model": "MTP-3",265 "device": DEVICE,266 "active_requests": ACTIVE_REQUESTS,267 "vocab_size": VOCAB_SIZE,268 "parameters": sum(p.numel() for p in model.parameters())269 }270 271@app.get("/info")272def model_info():273 """Información detallada del modelo"""274 return {275 "model_name": "MTP-3",276 "version": "3.0",277 "architecture": {278 "d_model": config["model"]["d_model"],279 "n_layers": config["model"]["n_layers"],280 "n_heads": config["model"]["n_heads"],281 "d_ff": config["model"]["d_ff"],282 "max_seq_len": config["model"]["max_seq_len"],283 "vocab_size": VOCAB_SIZE,284 "use_swiglu": config["model"].get("use_swiglu", False)285 },286 "parameters": sum(p.numel() for p in model.parameters()),287 "device": DEVICE,288 "improvements": [289 "RoPE (Rotary Position Embedding)",290 "RMSNorm",291 "Label Smoothing",292 "Repetition Penalty",293 "SwiGLU (opcional)" if config["model"].get("use_swiglu") else None294 ]295 }296 297# ======================298# 🎨 INTERFAZ WEB299# ======================300@app.get("/", response_class=HTMLResponse)301def chat_ui():302 return """303<!DOCTYPE html>304<html lang="es">305<head>306<meta charset="UTF-8">307<meta name="viewport" content="width=device-width, initial-scale=1.0, maximum-scale=1.0, user-scalable=no">308<title>MTP 3</title>309<link rel="preconnect" href="https://fonts.googleapis.com">310<link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>311<link href="https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600&display=swap" rel="stylesheet">312<style>313:root {314 --bg-color: #131314;315 --surface-color: #1E1F20;316 --accent-color: #4a9eff;317 --text-primary: #e3e3e3;318 --text-secondary: #9aa0a6;319 --user-bubble: #282a2c;320 --bot-actions-color: #c4c7c5;321 --logo-url: url('https://i.postimg.cc/yxS54PF3/IMG-3082.jpg');322}323* { box-sizing: border-box; outline: none; -webkit-tap-highlight-color: transparent; }324body {325 margin: 0;326 background-color: var(--bg-color);327 font-family: 'Inter', sans-serif;328 color: var(--text-primary);329 height: 100dvh;330 display: flex;331 flex-direction: column;332 overflow: hidden;333}334header {335 padding: 12px 20px;336 display: flex;337 align-items: center;338 justify-content: space-between;339 background: rgba(19, 19, 20, 0.85);340 backdrop-filter: blur(12px);341 position: fixed;342 top: 0;343 width: 100%;344 z-index: 50;345 border-bottom: 1px solid rgba(255,255,255,0.05);346}347.brand-wrapper {348 display: flex;349 align-items: center;350 gap: 12px;351 cursor: pointer;352}353.brand-logo {354 width: 32px;355 height: 32px;356 border-radius: 50%;357 background-image: var(--logo-url);358 background-size: cover;359 background-position: center;360 border: 1px solid rgba(255,255,255,0.1);361}362.brand-text {363 font-weight: 500;364 font-size: 1.05rem;365 display: flex;366 align-items: center;367 gap: 8px;368}369.version-badge {370 font-size: 0.75rem;371 background: rgba(74, 158, 255, 0.15);372 color: #8ab4f8;373 padding: 2px 8px;374 border-radius: 12px;375 font-weight: 600;376}377.chat-scroll {378 flex: 1;379 overflow-y: auto;380 padding: 80px 20px 40px 20px;381 display: flex;382 flex-direction: column;383 gap: 30px;384 max-width: 850px;385 margin: 0 auto;386 width: 100%;387 scroll-behavior: smooth;388}389.msg-row {390 display: flex;391 gap: 16px;392 width: 100%;393 opacity: 0;394 transform: translateY(10px);395 animation: slideUpFade 0.4s cubic-bezier(0.2, 0.8, 0.2, 1) forwards;396}397.msg-row.user { justify-content: flex-end; }398.msg-row.bot { justify-content: flex-start; align-items: flex-start; }399.msg-content {400 line-height: 1.6;401 font-size: 1rem;402 word-wrap: break-word;403 max-width: 85%;404}405.user .msg-content {406 background-color: var(--user-bubble);407 padding: 10px 18px;408 border-radius: 18px;409 border-top-right-radius: 4px;410 color: #fff;411}412.bot .msg-content-wrapper {413 display: flex;414 flex-direction: column;415 gap: 8px;416 width: 100%;417}418.bot .msg-text {419 padding-top: 6px;420 color: var(--text-primary);421}422.bot-avatar {423 width: 34px;424 height: 34px;425 min-width: 34px;426 border-radius: 50%;427 background-image: var(--logo-url);428 background-size: cover;429 box-shadow: 0 2px 6px rgba(0,0,0,0.2);430}431.bot-actions {432 display: flex;433 gap: 10px;434 opacity: 0;435 transition: opacity 0.3s;436 margin-top: 5px;437}438.action-btn {439 background: transparent;440 border: none;441 color: var(--text-secondary);442 cursor: pointer;443 padding: 4px;444 border-radius: 4px;445 display: flex;446 align-items: center;447 transition: color 0.2s, background 0.2s;448}449.action-btn:hover {450 color: var(--text-primary);451 background: rgba(255,255,255,0.08);452}453.action-btn svg { width: 16px; height: 16px; fill: currentColor; }454.typing-cursor::after {455 content: '';456 display: inline-block;457 width: 10px;458 height: 10px;459 background: var(--accent-color);460 border-radius: 50%;461 margin-left: 5px;462 vertical-align: middle;463 animation: blink 1s infinite;464}465.footer-container {466 padding: 0 20px 20px 20px;467 background: linear-gradient(to top, var(--bg-color) 85%, transparent);468 position: relative;469 z-index: 60;470}471.input-box {472 max-width: 850px;473 margin: 0 auto;474 background: var(--surface-color);475 border-radius: 28px;476 padding: 8px 10px 8px 20px;477 display: flex;478 align-items: center;479 border: 1px solid rgba(255,255,255,0.1);480 transition: border-color 0.2s, box-shadow 0.2s;481}482.input-box:focus-within {483 border-color: rgba(74, 158, 255, 0.5);484 box-shadow: 0 0 0 2px rgba(74, 158, 255, 0.1);485}486#userInput {487 flex: 1;488 background: transparent;489 border: none;490 color: white;491 font-size: 1rem;492 font-family: inherit;493 padding: 10px 0;494}495#mainBtn {496 background: white;497 color: black;498 border: none;499 width: 36px;500 height: 36px;501 border-radius: 50%;502 display: flex;503 align-items: center;504 justify-content: center;505 cursor: pointer;506 margin-left: 8px;507 transition: transform 0.2s;508}509#mainBtn:hover { transform: scale(1.05); }510.disclaimer {511 text-align: center;512 font-size: 0.75rem;513 color: #666;514 margin-top: 12px;515}516@keyframes slideUpFade {517 from { opacity: 0; transform: translateY(15px); }518 to { opacity: 1; transform: translateY(0); }519}520@keyframes blink { 0%, 100% { opacity: 1; } 50% { opacity: 0; } }521@keyframes pulseAvatar {522 0% { box-shadow: 0 0 0 0 rgba(74, 158, 255, 0.4); }523 70% { box-shadow: 0 0 0 8px rgba(74, 158, 255, 0); }524 100% { box-shadow: 0 0 0 0 rgba(74, 158, 255, 0); }525}526.pulsing { animation: pulseAvatar 1.5s infinite; }527::-webkit-scrollbar { width: 8px; }528::-webkit-scrollbar-track { background: transparent; }529::-webkit-scrollbar-thumb { background: #333; border-radius: 4px; }530</style>531</head>532<body>533<header>534 <div class="brand-wrapper" onclick="location.reload()">535 <div class="brand-logo"></div>536 <div class="brand-text">537 MTP <span class="version-badge">3</span>538 </div>539 </div>540</header>541<div id="chatScroll" class="chat-scroll">542 <div class="msg-row bot" style="animation-delay: 0.1s;">543 <div class="bot-avatar"></div>544 <div class="msg-content-wrapper">545 <div class="msg-text">546 ¡Hola! Soy MTP 3. ¿En qué puedo ayudarte hoy?547 </div>548 </div>549 </div>550</div>551<div class="footer-container">552 <div class="input-box">553 <input type="text" id="userInput" placeholder="Escribe un mensaje..." autocomplete="off">554 <button id="mainBtn" onclick="handleBtnClick()"></button>555 </div>556 <div class="disclaimer">557 MTP puede cometer errores. Considera verificar la información importante.558 </div>559</div>560<script>561const chatScroll = document.getElementById('chatScroll');562const userInput = document.getElementById('userInput');563const mainBtn = document.getElementById('mainBtn');564let isGenerating = false;565let abortController = null;566let typingTimeout = null;567let lastUserPrompt = "";568const ICON_SEND = `<svg width="20" height="20" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2"><path d="M22 2L11 13M22 2l-7 20-4-9-9-4 20-7z"></path></svg>`;569const ICON_STOP = `<svg width="14" height="14" viewBox="0 0 24 24" fill="currentColor" stroke="currentColor" stroke-width="0"><rect x="2" y="2" width="20" height="20" rx="4" ry="4"></rect></svg>`;570mainBtn.innerHTML = ICON_SEND;571function scrollToBottom() {572 chatScroll.scrollTop = chatScroll.scrollHeight;573}574function setBtnState(state) {575 if (state === 'sending') {576 mainBtn.innerHTML = ICON_STOP;577 isGenerating = true;578 } else {579 mainBtn.innerHTML = ICON_SEND;580 isGenerating = false;581 abortController = null;582 }583}584function handleBtnClick() {585 if (isGenerating) {586 stopGeneration();587 } else {588 sendMessage();589 }590}591function stopGeneration() {592 if (abortController) abortController.abort();593 if (typingTimeout) clearTimeout(typingTimeout);594 const activeCursor = document.querySelector('.typing-cursor');595 if (activeCursor) activeCursor.classList.remove('typing-cursor');596 const activeAvatar = document.querySelector('.pulsing');597 if (activeAvatar) activeAvatar.classList.remove('pulsing');598 setBtnState('idle');599 userInput.focus();600}601async function sendMessage(textOverride = null) {602 const text = textOverride || userInput.value.trim();603 if (!text) return;604 lastUserPrompt = text;605 if (!textOverride) {606 userInput.value = '';607 addMessage(text, 'user');608 }609 setBtnState('sending');610 abortController = new AbortController();611 const botRow = document.createElement('div');612 botRow.className = 'msg-row bot';613 const avatar = document.createElement('div');614 avatar.className = 'bot-avatar pulsing'; 615 const wrapper = document.createElement('div');616 wrapper.className = 'msg-content-wrapper';617 const msgText = document.createElement('div');618 msgText.className = 'msg-text'; 619 wrapper.appendChild(msgText);620 botRow.appendChild(avatar);621 botRow.appendChild(wrapper);622 chatScroll.appendChild(botRow);623 scrollToBottom();624 try {625 const response = await fetch('/generate', {626 method: 'POST',627 headers: { 'Content-Type': 'application/json' },628 body: JSON.stringify({ text: text }),629 signal: abortController.signal630 });631 const data = await response.json();632 if (!isGenerating) return; 633 avatar.classList.remove('pulsing');634 const reply = data.reply || "No entendí eso.";635 await typeWriter(msgText, reply);636 if (isGenerating) {637 addActions(wrapper, reply);638 setBtnState('idle');639 }640 } catch (error) {641 if (error.name === 'AbortError') {642 msgText.textContent += " [Detenido]";643 } else {644 avatar.classList.remove('pulsing');645 msgText.textContent = "Error de conexión.";646 msgText.style.color = "#ff8b8b";647 setBtnState('idle');648 }649 }650}651function addMessage(text, sender) {652 const row = document.createElement('div');653 row.className = `msg-row ${sender}`;654 const content = document.createElement('div');655 content.className = 'msg-content';656 content.textContent = text;657 row.appendChild(content);658 chatScroll.appendChild(row);659 scrollToBottom();660}661function typeWriter(element, text, speed = 12) {662 return new Promise(resolve => {663 let i = 0;664 element.classList.add('typing-cursor');665 function type() {666 if (!isGenerating) {667 element.classList.remove('typing-cursor');668 resolve();669 return;670 }671 if (i < text.length) {672 element.textContent += text.charAt(i);673 i++;674 scrollToBottom();675 typingTimeout = setTimeout(type, speed + Math.random() * 5);676 } else {677 element.classList.remove('typing-cursor');678 resolve();679 }680 }681 type();682 });683}684function addActions(wrapperElement, textToCopy) {685 const actionsDiv = document.createElement('div');686 actionsDiv.className = 'bot-actions';687 const copyBtn = document.createElement('button');688 copyBtn.className = 'action-btn';689 copyBtn.innerHTML = `<svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><rect x="9" y="9" width="13" height="13" rx="2" ry="2"></rect><path d="M5 15H4a2 2 0 0 1-2-2V4a2 2 0 0 1 2-2h9a2 2 0 0 1 2 2v1"></path></svg>`;690 copyBtn.onclick = () => {691 navigator.clipboard.writeText(textToCopy);692 };693 const regenBtn = document.createElement('button');694 regenBtn.className = 'action-btn';695 regenBtn.innerHTML = `<svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><path d="M23 4v6h-6"></path><path d="M1 20v-6h6"></path><path d="M3.51 9a9 9 0 0 1 14.85-3.36L23 10M1 14l4.64 4.36A9 9 0 0 0 20.49 15"></path></svg>`;696 regenBtn.onclick = () => {697 sendMessage(lastUserPrompt);698 };699 actionsDiv.appendChild(copyBtn);700 actionsDiv.appendChild(regenBtn);701 wrapperElement.appendChild(actionsDiv);702 requestAnimationFrame(() => actionsDiv.style.opacity = "1");703 scrollToBottom();704}705userInput.addEventListener('keydown', (e) => {706 if (e.key === 'Enter') handleBtnClick();707});708window.onload = () => userInput.focus();709</script>710</body>711</html>712"""713 714if __name__ == "__main__":715 port = int(os.environ.get("PORT", 7860))716 print(f"\n🚀 Iniciando servidor en puerto {port}...")717 print(f"🌐 Interfaz web: http://0.0.0.0:{port}")718 print(f"📡 API docs: http://0.0.0.0:{port}/docs")719 720 uvicorn.run(721 app,722 host="0.0.0.0",723 port=port,724 log_level="info"725 )