Update app.py

app.py

@@ -1,268 +1,94 @@
 import os
-import sys
 import torch
-import json
-import gc
-import re
 from fastapi import FastAPI
 from fastapi.responses import HTMLResponse
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
-from huggingface_hub import snapshot_download
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import uvicorn
-import math
-import torch.nn as nn
-import torch.nn.functional as F
-import sentencepiece as spm
-
-if torch.cuda.is_available():
-    DEVICE = "cuda"
-    print("✅ GPU detectada")
-    torch.backends.cudnn.benchmark = True
-else:
-    DEVICE = "cpu"
-    print("⚠️ CPU mode")
-    torch.set_num_threads(4)
-
-torch.set_grad_enabled(False)
-
-MODEL_REPO = "TeszenAI/MTP-3.1.1"
-
-class LayerNorm(nn.Module):
-    def __init__(self, d_model, eps=1e-5):
-        super().__init__()
-        self.weight = nn.Parameter(torch.ones(d_model))
-        self.bias = nn.Parameter(torch.zeros(d_model))
-        self.eps = eps
-    def forward(self, x):
-        return self.weight * (x - x.mean(-1, keepdim=True)) / (x.std(-1, keepdim=True) + self.eps) + self.bias
-
-class MultiHeadAttention(nn.Module):
-    def __init__(self, d_model, n_heads, dropout=0.1):
-        super().__init__()
-        assert d_model % n_heads == 0
-        self.d_k = d_model // n_heads
-        self.n_heads = n_heads
-        self.scale = math.sqrt(self.d_k)
-        
-        self.w_q = nn.Linear(d_model, d_model)
-        self.w_k = nn.Linear(d_model, d_model)
-        self.w_v = nn.Linear(d_model, d_model)
-        self.w_o = nn.Linear(d_model, d_model)
-        self.dropout = nn.Dropout(dropout)
-        
-    def forward(self, x, mask=None):
-        B, T, C = x.shape
-        Q = self.w_q(x).view(B, T, self.n_heads, self.d_k).transpose(1, 2)
-        K = self.w_k(x).view(B, T, self.n_heads, self.d_k).transpose(1, 2)
-        V = self.w_v(x).view(B, T, self.n_heads, self.d_k).transpose(1, 2)
-        
-        attn = (Q @ K.transpose(-2, -1)) * self.scale
-        if mask is not None:
-            attn = attn.masked_fill(mask == 0, float('-inf'))
-        
-        attn = F.softmax(attn, dim=-1)
-        attn = self.dropout(attn)
-        
-        out = (attn @ V).transpose(1, 2).contiguous().view(B, T, C)
-        return self.w_o(out)
-
-class FeedForward(nn.Module):
-    def __init__(self, d_model, d_ff, dropout=0.1):
-        super().__init__()
-        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))))
-
-class TransformerBlock(nn.Module):
-    def __init__(self, d_model, n_heads, d_ff, dropout=0.1):
-        super().__init__()
-        self.attention = MultiHeadAttention(d_model, n_heads, dropout)
-        self.feed_forward = FeedForward(d_model, d_ff, dropout)
-        self.norm1 = LayerNorm(d_model)
-        self.norm2 = LayerNorm(d_model)
-        self.dropout1 = nn.Dropout(dropout)
-        self.dropout2 = nn.Dropout(dropout)
-        
-    def forward(self, x, mask=None):
-        x = x + self.dropout1(self.attention(self.norm1(x), mask))
-        x = x + self.dropout2(self.feed_forward(self.norm2(x)))
-        return x
-
-class PositionalEncoding(nn.Module):
-    def __init__(self, d_model, max_len=5000):
-        super().__init__()
-        pe = torch.zeros(max_len, d_model)
-        pos = torch.arange(0, max_len).unsqueeze(1).float()
-        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
-        pe[:, 0::2] = torch.sin(pos * div_term)
-        pe[:, 1::2] = torch.cos(pos * 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, d_model=512, n_heads=8, n_layers=6, d_ff=2048, dropout=0.1, max_len=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.norm = LayerNorm(d_model)
-        self.lm_head = nn.Linear(d_model, vocab_size)
-        
-    def forward(self, x, mask=None):
-        if mask is None:
-            mask = torch.tril(torch.ones(x.size(1), x.size(1))).unsqueeze(0).unsqueeze(0).to(x.device)
-        
-        x = self.token_embedding(x) * math.sqrt(self.d_model)
-        x = self.pos_encoding(x)
-        
-        for block in self.blocks:
-            x = block(x, mask)
-        
-        x = self.norm(x)
-        return self.lm_head(x)
-    
-    @torch.inference_mode()
-    def generate(self, input_ids, max_new_tokens=150, temperature=0.7, top_k=50):
-        generated = input_ids
-        eos_token = 3
-        
-        for _ in range(max_new_tokens):
-            logits = self(generated)
-            next_logits = logits[0, -1, :] / temperature
-            
-            if top_k > 0:
-                top_k_val = min(top_k, next_logits.size(-1))
-                top_k_values = torch.topk(next_logits, top_k_val)[0]
-                next_logits[next_logits < top_k_values[-1]] = float('-inf')
-            
-            probs = F.softmax(next_logits, dim=-1)
-            next_token = torch.multinomial(probs, 1).item()
-            
-            if next_token == eos_token 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]) > 300:
-                break
-        
-        return generated
-
-print("📦 Descargando modelo desde HuggingFace...")
-repo_path = snapshot_download(repo_id=MODEL_REPO, repo_type="model", local_dir="mtp_repo")
-
-config_path = os.path.join(repo_path, "config.json")
-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', 6)}")
-else:
-    print("⚠️ Usando configuración por defecto")
-    config = {
-        "vocab_size": 8000,
-        "d_model": 512,
-        "n_heads": 8,
-        "n_layers": 6,
-        "d_ff": 2048,
-        "dropout": 0.1,
-        "max_len": 512
-    }
-
-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)
-VOCAB_SIZE = sp.get_piece_size()
-config["vocab_size"] = VOCAB_SIZE
-
-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']}")
+import re
 
-model = MTPModel(**config)
-model.to(DEVICE)
+# ==================== CONFIGURACIÓN ====================
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"📱 Dispositivo: {DEVICE}")
 
-model_path = os.path.join(repo_path, "mtp_model.pt")
-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)
+# Usar un modelo pequeño pero FUNCIONAL de HuggingFace
+# Opciones: "microsoft/DialoGPT-small" (mejor para conversación)
+#          "TinyLlama/TinyLlama-1.1B-Chat-v1.0" (más potente pero más lento)
+MODEL_NAME = "microsoft/DialoGPT-small"  # ~60MB, rápido y funcional
 
+print(f"📦 Cargando modelo {MODEL_NAME}...")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME).to(DEVICE)
 model.eval()
-print(f"✅ Modelo listo: {sum(p.numel() for p in model.parameters()):,} params")
+print(f"✅ Modelo cargado: {sum(p.numel() for p in model.parameters()):,} parámetros")
 
+# ==================== API ====================
 app = FastAPI()
 app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
 
 class PromptRequest(BaseModel):
     text: str
 
-def build_prompt(user_input):
-    return f"### Instrucción:\n{user_input}\n\n### Respuesta:\n"
-
-def clean_response(text):
+def clean_response(text: str) -> str:
+    """Limpia la respuesta del modelo"""
     if not text:
         return ""
-    text = re.sub(r'<unk>|<pad>|<s>|</s>', '', text)
+    
+    # Eliminar caracteres especiales
+    text = re.sub(r'<\|.*?\|>', '', text)
+    text = re.sub(r'\[.*?\]', '', text)
     text = re.sub(r'\s+', ' ', text).strip()
-    return text
+    
+    # Limitar longitud
+    if len(text) > 400:
+        text = text[:400]
+        last_dot = text.rfind('.')
+        if last_dot > 200:
+            text = text[:last_dot + 1]
+    
+    return text if text else "Lo siento, no pude generar una respuesta."
 
 @app.post("/generate")
 async def generate(req: PromptRequest):
     user_input = req.text.strip()
     if not user_input:
-        return {"reply": ""}
+        return {"reply": "Escribe un mensaje"}
+    
+    # Formatear entrada para el modelo
+    formatted_input = f"User: {user_input}\nBot:"
+    
+    # Tokenizar
+    inputs = tokenizer.encode(formatted_input, return_tensors="pt").to(DEVICE)
+    
+    # Generar
+    with torch.no_grad():
+        outputs = model.generate(
+            inputs,
+            max_new_tokens=100,
+            temperature=0.7,
+            top_k=50,
+            top_p=0.9,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id
+        )
     
-    prompt = build_prompt(user_input)
-    tokens = sp.encode(prompt)
+    # Decodificar
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
     
-    if len(tokens) > 450:
-        tokens = tokens[-450:]
+    # Extraer solo la respuesta del bot
+    if "Bot:" in response:
+        response = response.split("Bot:")[-1].strip()
+    elif "User:" in response:
+        parts = response.split("User:")
+        response = parts[-1].strip() if len(parts) > 1 else response
     
-    input_ids = torch.tensor([tokens], device=DEVICE)
+    response = clean_response(response)
     
-    try:
-        output_ids = model.generate(input_ids, max_new_tokens=120, temperature=0.7, top_k=50)
-        
-        gen_tokens = output_ids[0, len(tokens):].tolist()
-        
-        clean_tokens = []
-        for t in gen_tokens:
-            if t == 3 or t == 0 or t == 1:
-                break
-            clean_tokens.append(t)
-        
-        response = sp.decode(clean_tokens).strip() if clean_tokens else ""
-        response = clean_response(response)
-        
-        if not response:
-            response = "Lo siento, no pude generar una respuesta."
-        
-        print(f"📝 {user_input[:40]} -> {len(clean_tokens)} tokens")
-        
-        return {"reply": response[:500]}
+    print(f"📝 Usuario: {user_input[:50]}")
+    print(f"🤖 Respuesta: {response[:100]}")
     
-    except Exception as e:
-        print(f"❌ Error: {e}")
-        return {"reply": "Error al generar respuesta"}
+    return {"reply": response}
 
 @app.get("/health")
 def health():
@@ -400,11 +226,11 @@ body {
 <body>
 <div class="header">
     <h1><span class="dot"></span> MTP Assistant</h1>
-    <p>Modelo Transformer 512-dim | 6 capas</p>
+    <p>DialoGPT - Modelo conversacional real</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, tu asistente. ¿En qué puedo ayudarte hoy?</div>
     </div>
 </div>
 <div class="input-area">
@@ -466,7 +292,7 @@ async function send() {
         addMessage(data.reply || "No pude generar respuesta.", false);
     } catch (err) {
         removeTyping();
-        addMessage("Error de conexión.", false);
+        addMessage("Error de conexión. Intenta de nuevo.", false);
     } finally {
         loading = false;
         sendBtn.disabled = false;