Update app.py

app.py

@@ -4,11 +4,12 @@ import torch
 import json
 import time
 import gc
+import re
 from fastapi import FastAPI, Request
-from fastapi.responses import HTMLResponse, StreamingResponse
+from fastapi.responses import HTMLResponse
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel, Field
-from huggingface_hub import snapshot_download, hf_hub_download
+from huggingface_hub import snapshot_download
 import uvicorn
 import math
 import torch.nn as nn
@@ -33,7 +34,7 @@ torch.set_grad_enabled(False)
 MODEL_REPO = "TeszenAI/MTP-3"
 
 # ======================
-# DEFINIR ARQUITECTURA DEL MODELO (MTP-1.1)
+# ARQUITECTURA DEL MODELO (MISMA QUE EN colab.py)
 # ======================
 class LayerNorm(nn.Module):
     def __init__(self, d_model: int, eps: float = 1e-5):
@@ -41,7 +42,6 @@ class LayerNorm(nn.Module):
         self.weight = nn.Parameter(torch.ones(d_model))
         self.bias = nn.Parameter(torch.zeros(d_model))
         self.eps = eps
-    
     def forward(self, x):
         mean = x.mean(-1, keepdim=True)
         std = x.std(-1, keepdim=True)
@@ -60,7 +60,6 @@ class MultiHeadAttention(nn.Module):
         self.w_o = nn.Linear(d_model, d_model)
         self.dropout = nn.Dropout(dropout)
         self.scale = math.sqrt(self.d_k)
-        
     def forward(self, x, mask=None):
         batch_size, seq_len, _ = x.shape
         Q = self.w_q(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
@@ -81,7 +80,6 @@ class FeedForward(nn.Module):
         self.linear1 = nn.Linear(d_model, d_ff)
         self.linear2 = nn.Linear(d_ff, d_model)
         self.dropout = nn.Dropout(dropout)
-        
     def forward(self, x):
         return self.linear2(self.dropout(F.gelu(self.linear1(x))))
 
@@ -94,7 +92,6 @@ class TransformerBlock(nn.Module):
         self.norm2 = LayerNorm(d_model)
         self.dropout1 = nn.Dropout(dropout)
         self.dropout2 = nn.Dropout(dropout)
-        
     def forward(self, x, mask=None):
         attn_output = self.attention(x, mask)
         x = x + self.dropout1(attn_output)
@@ -113,22 +110,19 @@ class PositionalEncoding(nn.Module):
         pe[:, 0::2] = torch.sin(position * div_term)
         pe[:, 1::2] = torch.cos(position * div_term)
         self.register_buffer('pe', pe.unsqueeze(0))
-        
     def forward(self, x):
         return x + self.pe[:, :x.size(1), :]
 
 class MTPModel(nn.Module):
-    def __init__(self, vocab_size: int, d_model: int = 128, n_heads: int = 4, 
-                 n_layers: int = 4, d_ff: int = 512, dropout: float = 0.1, max_len: int = 256):
+    def __init__(self, vocab_size: int, d_model: int = 256, n_heads: int = 8,
+                 n_layers: int = 6, d_ff: int = 1024, dropout: float = 0.1, max_len: int = 512):
         super().__init__()
         self.vocab_size = vocab_size
         self.d_model = d_model
         self.max_len = max_len
         self.token_embedding = nn.Embedding(vocab_size, d_model)
         self.pos_encoding = PositionalEncoding(d_model, max_len)
-        self.blocks = nn.ModuleList([
-            TransformerBlock(d_model, n_heads, d_ff, dropout) for _ in range(n_layers)
-        ])
+        self.blocks = nn.ModuleList([TransformerBlock(d_model, n_heads, d_ff, dropout) for _ in range(n_layers)])
         self.norm = LayerNorm(d_model)
         self.lm_head = nn.Linear(d_model, vocab_size)
         
@@ -140,82 +134,34 @@ class MTPModel(nn.Module):
         for block in self.blocks:
             x = block(x, mask)
         x = self.norm(x)
-        logits = self.lm_head(x)
-        return logits
-    
-    def generate(self, input_ids, max_new_tokens=100, temperature=0.8, top_k=50, top_p=0.9, repetition_penalty=1.1):
-        """Método de generación compatible con la interfaz"""
-        generated = input_ids
-        
-        for _ in range(max_new_tokens):
-            with torch.no_grad():
-                logits = self(generated)
-                next_logits = logits[0, -1, :] / temperature
-            
-            if repetition_penalty != 1.0:
-                for token_id in set(generated[0].tolist()):
-                    next_logits[token_id] /= repetition_penalty
-            
-            if top_k > 0:
-                indices_to_remove = next_logits < torch.topk(next_logits, top_k)[0][..., -1, None]
-                next_logits[indices_to_remove] = float('-inf')
-            
-            if top_p < 1.0:
-                sorted_logits, sorted_indices = torch.sort(next_logits, descending=True)
-                cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
-                sorted_indices_to_remove = cumulative_probs > top_p
-                sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
-                sorted_indices_to_remove[..., 0] = 0
-                indices_to_remove = sorted_indices[sorted_indices_to_remove]
-                next_logits[indices_to_remove] = float('-inf')
-            
-            probs = F.softmax(next_logits, dim=-1)
-            next_token = torch.multinomial(probs, num_samples=1).item()
-            
-            if next_token == 3:
-                break
-            
-            generated = torch.cat([generated, torch.tensor([[next_token]], device=generated.device)], dim=1)
-        
-        return generated
+        return self.lm_head(x)
 
 # ======================
 # DESCARGA Y CARGA DEL MODELO CON REINTENTOS
 # ======================
 def download_with_retry(repo_id, local_dir, max_retries=3):
-    """Descarga el modelo con reintentos para evitar timeouts"""
-    
     for attempt in range(max_retries):
         try:
             print(f"📦 Intento {attempt + 1}/{max_retries} - Descargando modelo desde {repo_id}...")
-            
-            # Configurar timeout más largo para descargas
             repo_path = snapshot_download(
                 repo_id=repo_id,
                 repo_type="model",
                 local_dir=local_dir,
                 resume_download=True,
-                local_files_only=False,
-                ignore_patterns=["*.h5", "*.ot", "*.msgpack"]  # Ignorar archivos grandes innecesarios
+                local_files_only=False
             )
-            
             print(f"✅ Modelo descargado exitosamente en: {repo_path}")
             return repo_path
-            
         except Exception as e:
             print(f"⚠️ Error en intento {attempt + 1}: {str(e)[:200]}")
             if attempt < max_retries - 1:
-                wait_time = (attempt + 1) * 3
-                print(f"🔄 Reintentando en {wait_time} segundos...")
-                time.sleep(wait_time)
+                time.sleep(3)
             else:
-                print("❌ No se pudo descargar el modelo después de múltiples intentos")
                 raise
+    return local_dir
 
-# Intentar descargar el modelo
 print(f"🚀 Iniciando carga del modelo desde {MODEL_REPO}...")
 
-# Verificar si ya existe en caché local
 if os.path.exists("mtp_repo") and os.path.exists("mtp_repo/mtp_model.pt"):
     print("📁 Modelo encontrado en caché local")
     repo_path = "mtp_repo"
@@ -223,10 +169,8 @@ else:
     try:
         repo_path = download_with_retry(MODEL_REPO, "mtp_repo", max_retries=3)
     except Exception as e:
-        print(f"⚠️ Error crítico: {e}")
-        print("🏗️ Usando configuración por defecto...")
+        print(f"⚠️ Error: {e}")
         repo_path = "mtp_repo"
-        os.makedirs(repo_path, exist_ok=True)
 
 # Cargar configuración
 config_path = os.path.join(repo_path, "config.json")
@@ -234,35 +178,29 @@ if os.path.exists(config_path):
     with open(config_path, "r") as f:
         config = json.load(f)
 else:
+    # Configuración por defecto (MISMA que en colab.py)
     config = {
-        "vocab_size": 5000,
-        "d_model": 128,
-        "n_heads": 4,
-        "n_layers": 4,
-        "d_ff": 512,
+        "vocab_size": 2000,
+        "d_model": 256,
+        "n_heads": 8,
+        "n_layers": 6,
+        "d_ff": 1024,
         "dropout": 0.1,
-        "max_len": 256
+        "max_len": 512
     }
 
 # Cargar tokenizador
 tokenizer_path = os.path.join(repo_path, "mtp_tokenizer.model")
 if os.path.exists(tokenizer_path):
-    try:
-        sp = spm.SentencePieceProcessor()
-        sp.load(tokenizer_path)
-        VOCAB_SIZE = sp.get_piece_size()
-        print(f"✅ Tokenizador cargado: {VOCAB_SIZE} tokens")
-    except Exception as e:
-        print(f"⚠️ Error cargando tokenizador: {e}")
-        VOCAB_SIZE = config.get("vocab_size", 5000)
-        sp = None
+    sp = spm.SentencePieceProcessor()
+    sp.load(tokenizer_path)
+    VOCAB_SIZE = sp.get_piece_size()
+    config["vocab_size"] = VOCAB_SIZE
+    print(f"✅ Tokenizador cargado: {VOCAB_SIZE} tokens")
 else:
-    print("⚠️ No se encontró tokenizador, usando vocabulario por defecto")
-    VOCAB_SIZE = config.get("vocab_size", 5000)
+    print("❌ No se encontró tokenizador")
     sp = None
-
-# Actualizar vocab_size en config
-config["vocab_size"] = VOCAB_SIZE
+    VOCAB_SIZE = config.get("vocab_size", 2000)
 
 print(f"🧠 Inicializando modelo MTP...")
 print(f"   → Vocabulario: {VOCAB_SIZE}")
@@ -279,39 +217,21 @@ if os.path.exists(model_path):
     try:
         state_dict = torch.load(model_path, map_location=DEVICE)
         model.load_state_dict(state_dict)
-        print("✅ Pesos del modelo cargados")
+        print("✅ Pesos del modelo cargados correctamente")
     except Exception as e:
         print(f"⚠️ Error cargando pesos: {e}")
-        print("   Usando pesos aleatorios")
 else:
-    print("⚠️ No se encontró mtp_model.pt, usando pesos aleatorios")
+    print("⚠️ No se encontró mtp_model.pt")
 
 model.eval()
 
-# Cuantización para CPU
-if DEVICE == "cpu":
-    print("⚡ Optimizando para CPU...")
-    try:
-        model = torch.quantization.quantize_dynamic(
-            model, 
-            {nn.Linear}, 
-            dtype=torch.qint8
-        )
-        print("✅ Cuantización aplicada")
-    except Exception as e:
-        print(f"⚠️ No se pudo aplicar cuantización: {e}")
-
 param_count = sum(p.numel() for p in model.parameters())
 print(f"✅ Modelo cargado: {param_count:,} parámetros ({param_count/1e6:.1f}M)")
 
 # ======================
 # API CONFIG
 # ======================
-app = FastAPI(
-    title="MTP-1.1 API",
-    description="API para modelo de lenguaje MTP-1.1",
-    version="1.1"
-)
+app = FastAPI(title="MTP API", description="API para modelo de lenguaje MTP", version="1.0")
 
 app.add_middleware(
     CORSMiddleware,
@@ -321,58 +241,103 @@ app.add_middleware(
 )
 
 class PromptRequest(BaseModel):
-    text: str = Field(..., max_length=2000, description="Texto de entrada")
-    max_tokens: int = Field(default=150, ge=10, le=300, description="Tokens máximos a generar")
-    temperature: float = Field(default=0.7, ge=0.1, le=2.0, description="Temperatura de muestreo")
-    top_k: int = Field(default=50, ge=1, le=100, description="Top-k sampling")
-    top_p: float = Field(default=0.9, ge=0.1, le=1.0, description="Top-p (nucleus) sampling")
-    repetition_penalty: float = Field(default=1.1, ge=1.0, le=2.0, description="Penalización por repetición")
+    text: str = Field(..., max_length=2000)
+    max_tokens: int = Field(default=150, ge=10, le=300)
+    temperature: float = Field(default=0.7, ge=0.1, le=2.0)
+    top_k: int = Field(default=50, ge=1, le=100)
+    top_p: float = Field(default=0.9, ge=0.1, le=1.0)
+    repetition_penalty: float = Field(default=1.1, ge=1.0, le=2.0)
+
+# ======================
+# FUNCIÓN DE GENERACIÓN (IGUAL QUE EN colab.py)
+# ======================
+def generate_response(model, tokenizer, prompt, max_length=150, temperature=0.7, top_k=50, top_p=0.9, device='cpu'):
+    model.eval()
+    formatted_prompt = f"### Instrucción:\n{prompt}\n\n### Respuesta:\n"
+    input_ids = tokenizer.encode(formatted_prompt)
+    generated = input_ids.copy()
+    eos_id = tokenizer.eos_id()
+
+    for _ in range(max_length):
+        input_tensor = torch.tensor([generated[-model.max_len:]], dtype=torch.long).to(device)
+        with torch.no_grad():
+            logits = model(input_tensor)
+            next_logits = logits[0, -1, :] / temperature
+
+        if top_k > 0:
+            indices_to_remove = next_logits < torch.topk(next_logits, top_k)[0][..., -1, None]
+            next_logits[indices_to_remove] = float('-inf')
+
+        if top_p < 1.0:
+            sorted_logits, sorted_indices = torch.sort(next_logits, descending=True)
+            cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+            sorted_indices_to_remove = cumulative_probs > top_p
+            sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+            sorted_indices_to_remove[..., 0] = 0
+            indices_to_remove = sorted_indices[sorted_indices_to_remove]
+            next_logits[indices_to_remove] = float('-inf')
+
+        probs = F.softmax(next_logits, dim=-1)
+        next_token = torch.multinomial(probs, 1).item()
+
+        if next_token == eos_id:
+            break
+
+        if len(generated) > 20:
+            last_tokens = generated[-10:]
+            if len(set(last_tokens)) == 1:
+                break
+
+        generated.append(next_token)
+
+    response = tokenizer.decode(generated)
+    if "### Respuesta:" in response:
+        response = response.split("### Respuesta:")[-1].strip()
+    elif "Respuesta:" in response:
+        response = response.split("Respuesta:")[-1].strip()
+    elif "[/INST]" in response:
+        response = response.split("[/INST]")[-1].strip()
+
+    # Limpiar caracteres basura
+    garbage_words = ['foompañances', 'ciudadores', 'mejtedon', 'calportedon', 'rápidodcor', 'baon', 'domol']
+    for word in garbage_words:
+        response = response.replace(word, '')
+    
+    response = re.sub(r'[^\w\s\u00C0-\u00FF\u0100-\u017F.,!?¿¡()\-:;"]+', ' ', response)
+    response = re.sub(r'\s+', ' ', response).strip()
 
-def build_prompt(user_input: str) -> str:
-    """Construye el prompt en el formato del modelo"""
-    return f"### Instrucción:\n{user_input}\n\n### Respuesta:\n"
+    if len(response) < 2:
+        response = "Entendido. ¿Algo más en lo que pueda ayudarte?"
+
+    return response
 
 # ======================
-# GESTIÓN DE CARGA
+# ENDPOINTS
 # ======================
 ACTIVE_REQUESTS = 0
 
-class MTPTokenizer:
-    """Wrapper para el tokenizador de SentencePiece"""
+class TokenizerWrapper:
     def __init__(self, sp_model):
         self.sp = sp_model
-    
     def encode(self, text):
         if self.sp is None:
-            # Tokenizador simple para fallback
             return [ord(c) % 1000 for c in text[:200]]
         return self.sp.encode(text)
-    
     def decode(self, tokens):
         if self.sp is None:
             return ''.join([chr(t % 128) if 32 <= t % 128 < 127 else ' ' for t in tokens])
         return self.sp.decode(tokens)
-    
-    def bos_id(self):
-        if self.sp is None:
-            return 2
-        return self.sp.bos_id()
-    
     def eos_id(self):
-        if self.sp is None:
-            return 3
-        return self.sp.eos_id()
-    
+        return self.sp.eos_id() if self.sp else 3
+    def bos_id(self):
+        return self.sp.bos_id() if self.sp else 2
     def pad_id(self):
-        if self.sp is None:
-            return 0
-        return self.sp.pad_id()
+        return self.sp.pad_id() if self.sp else 0
 
-tokenizer_wrapper = MTPTokenizer(sp)
+tokenizer_wrapper = TokenizerWrapper(sp)
 
 @app.post("/generate")
 async def generate(req: PromptRequest):
-    """Endpoint principal de generación de texto"""
     global ACTIVE_REQUESTS
     ACTIVE_REQUESTS += 1
     
@@ -389,83 +354,51 @@ async def generate(req: PromptRequest):
         ACTIVE_REQUESTS -= 1
         return {"reply": "", "tokens_generated": 0}
 
-    full_prompt = build_prompt(user_input)
-    tokens = [tokenizer_wrapper.bos_id()] + tokenizer_wrapper.encode(full_prompt)
-    input_ids = torch.tensor([tokens], device=DEVICE)
-
     try:
-        with torch.no_grad():
-            output_ids = model.generate(
-                input_ids,
-                max_new_tokens=dyn_max_tokens,
-                temperature=dyn_temperature,
-                top_k=req.top_k,
-                top_p=req.top_p,
-                repetition_penalty=req.repetition_penalty
-            )
-
-        gen_tokens = output_ids[0, len(tokens):].tolist()
-        
-        safe_tokens = [
-            t for t in gen_tokens 
-            if 0 <= t < VOCAB_SIZE and t != tokenizer_wrapper.eos_id()
-        ]
-        
-        response = tokenizer_wrapper.decode(safe_tokens).strip()
-        
-        # Limpiar la respuesta
-        if "###" in response:
-            response = response.split("###")[0].strip()
-        
-        # Si la respuesta está vacía, devolver mensaje por defecto
-        if not response or len(response) < 2:
-            response = "Entendido. ¿En qué más puedo ayudarte?"
-
+        response = generate_response(
+            model, tokenizer_wrapper, user_input,
+            max_length=dyn_max_tokens,
+            temperature=dyn_temperature,
+            top_k=req.top_k,
+            top_p=req.top_p,
+            device=DEVICE
+        )
         return {
             "reply": response,
-            "tokens_generated": len(safe_tokens),
-            "model": "MTP-1.1"
+            "tokens_generated": len(response.split()),
+            "model": "MTP"
         }
-    
     except Exception as e:
-        print(f"❌ Error durante generación: {e}")
-        return {
-            "reply": "Lo siento, ocurrió un error al procesar tu solicitud.",
-            "error": str(e)
-        }
-    
+        print(f"❌ Error: {e}")
+        return {"reply": "Lo siento, ocurrió un error.", "error": str(e)}
     finally:
         ACTIVE_REQUESTS -= 1
         if DEVICE == "cuda":
             torch.cuda.empty_cache()
         gc.collect()
 
-# ======================
-# ENDPOINTS DE INFORMACIÓN
-# ======================
 @app.get("/health")
 def health_check():
     return {
         "status": "healthy",
-        "model": "MTP-1.1",
+        "model": "MTP",
         "device": DEVICE,
         "active_requests": ACTIVE_REQUESTS,
-        "vocab_size": VOCAB_SIZE,
-        "model_loaded": os.path.exists("mtp_repo/mtp_model.pt")
+        "vocab_size": VOCAB_SIZE
     }
 
 @app.get("/info")
 def model_info():
     return {
-        "model_name": "MTP-1.1",
-        "version": "1.1",
+        "model_name": "MTP",
+        "version": "1.0",
         "architecture": config,
         "parameters": sum(p.numel() for p in model.parameters()),
         "device": DEVICE
     }
 
 # ======================
-# INTERFAZ WEB (MODERNA DE MTP-3)
+# INTERFAZ WEB COMPLETA (CON TODAS LAS FUNCIONES ORIGINALES)
 # ======================
 @app.get("/", response_class=HTMLResponse)
 def chat_ui():
@@ -475,7 +408,7 @@ def chat_ui():
 <head>
 <meta charset="UTF-8">
 <meta name="viewport" content="width=device-width, initial-scale=1.0, maximum-scale=1.0, user-scalable=no">
-<title>MTP 1.1</title>
+<title>MTP 3</title>
 <link rel="preconnect" href="https://fonts.googleapis.com">
 <link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>
 <link href="https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600&display=swap" rel="stylesheet">
@@ -795,7 +728,7 @@ async function sendMessage(textOverride = null) {
         const response = await fetch('/generate', {
             method: 'POST',
             headers: { 'Content-Type': 'application/json' },
-            body: JSON.stringify({ text: text }),
+            body: JSON.stringify({ text: text, max_tokens: 150, temperature: 0.7 }),
             signal: abortController.signal
         });
         const data = await response.json();
@@ -883,7 +816,7 @@ window.onload = () => userInput.focus();
 
 if __name__ == "__main__":
     port = int(os.environ.get("PORT", 7860))
-    print(f"\n🚀 Iniciando servidor MTP-1.1 en puerto {port}...")
+    print(f"\n🚀 Iniciando servidor MTP en puerto {port}...")
     print(f"🌐 Interfaz web: http://0.0.0.0:{port}")
     print(f"📡 API docs: http://0.0.0.0:{port}/docs")