Update app.py

app.py

@@ -3,7 +3,7 @@ import sys
 import torch
 import json
 import gc
-import time
+import re
 from fastapi import FastAPI
 from fastapi.responses import HTMLResponse
 from fastapi.middleware.cors import CORSMiddleware
@@ -17,11 +17,11 @@ import sentencepiece as spm
 
 if torch.cuda.is_available():
     DEVICE = "cuda"
-    print("✅ GPU")
+    print("✅ GPU detectada")
     torch.backends.cudnn.benchmark = True
 else:
     DEVICE = "cpu"
-    print("⚠️ CPU")
+    print("⚠️ CPU mode")
     torch.set_num_threads(4)
 
 torch.set_grad_enabled(False)
@@ -114,7 +114,9 @@ class MTPModel(nn.Module):
         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)
         
@@ -129,7 +131,7 @@ class MTPModel(nn.Module):
         return self.lm_head(x)
     
     @torch.inference_mode()
-    def generate(self, input_ids, max_new_tokens=100, temperature=0.8, top_k=40):
+    def generate(self, input_ids, max_new_tokens=150, temperature=0.7, top_k=50):
         generated = input_ids
         for _ in range(max_new_tokens):
             logits = self(generated)
@@ -143,7 +145,7 @@ class MTPModel(nn.Module):
             if next_token == 3 or next_token == 0 or next_token == 1:
                 break
             generated = torch.cat([generated, torch.tensor([[next_token]], device=generated.device)], dim=1)
-            if len(generated[0]) > 180:
+            if len(generated[0]) > 200:
                 break
         return generated
 
@@ -151,25 +153,52 @@ print("📦 Descargando modelo...")
 repo_path = snapshot_download(repo_id=MODEL_REPO, repo_type="model", local_dir="mtp_repo")
 
 config_path = os.path.join(repo_path, "config.json")
-with open(config_path, "r") as f:
-    config = json.load(f)
+if os.path.exists(config_path):
+    with open(config_path, "r") as f:
+        config = json.load(f)
+    print(f"✅ Configuración cargada: d_model={config.get('d_model', 512)}, layers={config.get('n_layers', 8)}")
+else:
+    print("⚠️ Usando configuración por defecto (igual que colab.py)")
+    config = {
+        "vocab_size": 8000,
+        "d_model": 512,
+        "n_heads": 8,
+        "n_layers": 8,
+        "d_ff": 2048,
+        "dropout": 0.1,
+        "max_len": 1024
+    }
 
 tokenizer_path = os.path.join(repo_path, "mtp_tokenizer.model")
+if not os.path.exists(tokenizer_path):
+    print(f"❌ Tokenizador no encontrado")
+    sys.exit(1)
+
 sp = spm.SentencePieceProcessor()
 sp.load(tokenizer_path)
-config["vocab_size"] = sp.get_piece_size()
+VOCAB_SIZE = sp.get_piece_size()
+config["vocab_size"] = VOCAB_SIZE
 
-print(f"✅ Vocab: {config['vocab_size']}")
-print(f"✅ Dim: {config['d_model']}, Layers: {config['n_layers']}")
+print(f"🧠 Inicializando modelo MTP...")
+print(f"   → Vocabulario: {VOCAB_SIZE}")
+print(f"   → Dimensión: {config['d_model']}")
+print(f"   → Capas: {config['n_layers']}")
+print(f"   → Heads: {config['n_heads']}")
 
 model = MTPModel(**config)
 model.to(DEVICE)
 
 model_path = os.path.join(repo_path, "mtp_model.pt")
-state_dict = torch.load(model_path, map_location=DEVICE)
-model.load_state_dict(state_dict, strict=False)
+if os.path.exists(model_path):
+    state_dict = torch.load(model_path, map_location=DEVICE)
+    model.load_state_dict(state_dict, strict=False)
+    print("✅ Pesos cargados correctamente")
+else:
+    print(f"❌ Modelo no encontrado")
+    sys.exit(1)
+
 model.eval()
-print(f"✅ Modelo: {sum(p.numel() for p in model.parameters()):,} params")
+print(f"✅ Modelo listo: {sum(p.numel() for p in model.parameters()):,} params")
 
 app = FastAPI()
 app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
@@ -180,24 +209,29 @@ class PromptRequest(BaseModel):
 def build_prompt(user_input):
     return f"### Instrucción:\n{user_input}\n\n### Respuesta:\n"
 
+def clean_response(text):
+    if not text:
+        return ""
+    text = re.sub(r'<unk>|<pad>|<s>|</s>', '', text)
+    text = re.sub(r'\s+', ' ', text).strip()
+    return text
+
 @app.post("/generate")
 async def generate(req: PromptRequest):
-    start = time.time()
     user_input = req.text.strip()
-    
     if not user_input:
-        return {"reply": "Escribe un mensaje", "time": 0}
+        return {"reply": ""}
     
     prompt = build_prompt(user_input)
     tokens = sp.encode(prompt)
     
-    if len(tokens) > 800:
-        tokens = tokens[-800:]
+    if len(tokens) > 900:
+        tokens = tokens[-900:]
     
     input_ids = torch.tensor([tokens], device=DEVICE)
     
     try:
-        output_ids = model.generate(input_ids, max_new_tokens=80, temperature=0.75, top_k=40)
+        output_ids = model.generate(input_ids, max_new_tokens=120, temperature=0.7, top_k=50)
         
         gen_tokens = output_ids[0, len(tokens):].tolist()
         
@@ -208,26 +242,15 @@ async def generate(req: PromptRequest):
             clean_tokens.append(t)
         
         response = sp.decode(clean_tokens).strip() if clean_tokens else ""
+        response = clean_response(response)
         
-        markers = ["### Respuesta:", "Respuesta:", "[/INST]", "Asistente:", "Usuario:"]
-        for marker in markers:
-            if marker in response:
-                response = response.split(marker)[-1]
-        
-        response = response.replace('<unk>', '').replace('<pad>', '').replace('<s>', '').replace('</s>', '')
-        response = ' '.join(response.split())
-        
-        if not response or len(response) < 2:
-            response = "Entiendo tu pregunta. ¿Podrías darme más detalles?"
+        print(f"📝 {user_input[:40]} -> {len(clean_tokens)} tokens")
         
-        elapsed = time.time() - start
-        print(f"✅ {user_input[:25]}... -> {elapsed:.1f}s ({len(clean_tokens)} tokens)")
-        
-        return {"reply": response[:350], "time": elapsed}
+        return {"reply": response[:500]}
     
     except Exception as e:
         print(f"❌ Error: {e}")
-        return {"reply": "Error, intenta de nuevo", "time": 0}
+        return {"reply": ""}
 
 @app.get("/health")
 def health():
@@ -360,11 +383,11 @@ body {
 <body>
 <div class="header">
     <h1><span class="dot"></span> MTP Assistant</h1>
-    <p>Modelo Transformer | Generación 100% por IA</p>
+    <p>Modelo Transformer 512-dim / 8-capas</p>
 </div>
 <div class="chat" id="chat">
     <div class="message bot">
-        <div class="message-content">¡Hola! Soy MTP. ¿En qué puedo ayudarte?</div>
+        <div class="message-content">Hola, soy MTP. ¿En qué puedo ayudarte?</div>
     </div>
 </div>
 <div class="input-area">
@@ -415,8 +438,6 @@ async function send() {
     sendBtn.disabled = true;
     addTyping();
     
-    const startTime = Date.now();
-    
     try {
         const res = await fetch('/generate', {
             method: 'POST',
@@ -424,13 +445,11 @@ async function send() {
             body: JSON.stringify({ text: text })
         });
         const data = await res.json();
-        const elapsed = ((Date.now() - startTime) / 1000).toFixed(1);
         removeTyping();
         addMessage(data.reply || "No pude generar respuesta.", false);
-        console.log(`Respuesta en ${elapsed}s`);
     } catch (err) {
         removeTyping();
-        addMessage("Error de conexión. Intenta de nuevo.", false);
+        addMessage("Error de conexión.", false);
     } finally {
         loading = false;
         sendBtn.disabled = false;