Update app.py

app.py

@@ -2,8 +2,9 @@ import os
 import sys
 import torch
 import pickle
+import time
 from fastapi import FastAPI
-from fastapi.responses import HTMLResponse
+from fastapi.responses import HTMLResponse, StreamingResponse
 from pydantic import BaseModel
 from huggingface_hub import snapshot_download
 import uvicorn
@@ -11,7 +12,6 @@ import uvicorn
 # ======================
 # CONFIGURACIÓN DE DISPOSITIVO (GPU/CPU)
 # ======================
-# Detectar automáticamente si hay una GPU NVIDIA disponible
 if torch.cuda.is_available():
     DEVICE = "cuda"
     print("✅ GPU NVIDIA detectada. Usando CUDA.")
@@ -19,12 +19,18 @@ else:
     DEVICE = "cpu"
     print("⚠️ GPU no detectada. Usando CPU (puede ser más lento).")
 
+# ======================
+# OPTIMIZACIÓN CPU
+# ======================
+torch.set_grad_enabled(False)
+torch.set_num_threads(max(1, os.cpu_count() // 2))
+
 MODEL_REPO = "teszenofficial/mtptz"
 
 # ======================
 # DESCARGA DEL MODELO
 # ======================
-print(f"--- SISTEMA MTP 1.1 ---")
+print(f"--- SISTEMA MTP 2 ---")
 print(f"Descargando/Verificando modelo desde {MODEL_REPO}...")
 repo_path = snapshot_download(
     repo_id=MODEL_REPO,
@@ -34,12 +40,8 @@ repo_path = snapshot_download(
 
 sys.path.insert(0, repo_path)
 
-try:
-    from model import MTPMiniModel
-    from tokenizer import MTPTokenizer
-except ImportError:
-    print("Advertencia: Verifica la estructura de archivos del modelo.")
-    pass
+from model import MTPMiniModel
+from tokenizer import MTPTokenizer
 
 # ======================
 # CARGA DEL MODELO
@@ -52,10 +54,11 @@ tokenizer = MTPTokenizer(
     os.path.join(repo_path, "mtp_tokenizer.model")
 )
 
+VOCAB_SIZE = tokenizer.sp.get_piece_size()
 config = model_data["config"]
 
 model = MTPMiniModel(
-    vocab_size=model_data["vocab_size"],
+    vocab_size=VOCAB_SIZE,
     d_model=config["model"]["d_model"],
     n_layers=config["model"]["n_layers"],
     n_heads=config["model"]["n_heads"],
@@ -64,11 +67,22 @@ model = MTPMiniModel(
     dropout=0.0
 )
 
-# Cargar pesos y mover a GPU
 model.load_state_dict(model_data["model_state_dict"])
-model.to(DEVICE)
 model.eval()
-print(f"🚀 MTP 1.1 listo y corriendo en: {DEVICE.upper()}")
+
+# ======================
+# ⚙️ CUANTIZACIÓN CPU
+# ======================
+if DEVICE == "cpu":
+    model = torch.quantization.quantize_dynamic(
+        model,
+        {torch.nn.Linear},
+        dtype=torch.qint8
+    )
+    print("⚙️ Modelo cuantizado para CPU")
+
+model.to(DEVICE)
+print(f"🚀 MTP 2 listo y corriendo en: {DEVICE.upper()}")
 
 # ======================
 # API FASTAPI
@@ -78,16 +92,31 @@ app = FastAPI(title="MTP 2 API")
 class Prompt(BaseModel):
     text: str
 
+# ======================
+# 🧠 PROMPT MEJORADO (MISMO FORMATO)
+# ======================
+def build_prompt(user_input: str) -> str:
+    return f"""Eres MTP, un modelo de lenguaje experimental.
+Responde de forma clara, directa y coherente.
+No inventes información.
+
+### Instrucción:
+{user_input}
+
+### Respuesta:
+"""
+
+# ======================
+# GENERACIÓN NORMAL (IGUAL QUE ANTES)
+# ======================
 @app.post("/generate")
 def generate(prompt: Prompt):
     user_input = prompt.text.strip()
     if not user_input:
         return {"reply": ""}
 
-    full_prompt = f"### Instrucción:\n{user_input}\n\n### Respuesta:\n"
+    full_prompt = build_prompt(user_input)
     tokens = [tokenizer.bos_id()] + tokenizer.encode(full_prompt)
-    
-    # IMPORTANTE: Mover los inputs también a la GPU
     input_ids = torch.tensor([tokens], device=DEVICE)
 
     with torch.no_grad():
@@ -101,21 +130,57 @@ def generate(prompt: Prompt):
 
     gen_tokens = output_ids[0, len(tokens):].tolist()
 
-    if tokenizer.eos_id() in gen_tokens:
-        gen_tokens = gen_tokens[:gen_tokens.index(tokenizer.eos_id())]
+    # 🔒 FILTRO DE SEGURIDAD
+    safe_tokens = [
+        t for t in gen_tokens
+        if 0 <= t < VOCAB_SIZE and t != tokenizer.eos_id()
+    ]
 
-    response = tokenizer.decode(gen_tokens).strip()
+    response = tokenizer.decode(safe_tokens).strip()
     if "###" in response:
         response = response.split("###")[0].strip()
 
     return {"reply": response}
 
 # ======================
-# INTERFAZ WEB (FRONTEND MEJORADO)
+# 📡 STREAMING SSE OFICIAL
+# ======================
+@app.get("/generate_sse")
+def generate_sse(text: str):
+    def event_stream():
+        full_prompt = build_prompt(text)
+        tokens = [tokenizer.bos_id()] + tokenizer.encode(full_prompt)
+        input_ids = torch.tensor([tokens], device=DEVICE)
+
+        for _ in range(150):
+            with torch.no_grad():
+                logits = model(input_ids)[:, -1, :VOCAB_SIZE]
+                probs = torch.softmax(logits / 0.7, dim=-1)
+                next_id = torch.argmax(probs, dim=-1).item()
+
+            if next_id == tokenizer.eos_id():
+                break
+
+            if 0 <= next_id < VOCAB_SIZE:
+                token_text = tokenizer.decode([next_id])
+                yield f"data:{token_text}\n\n"
+                input_ids = torch.cat(
+                    [input_ids, torch.tensor([[next_id]], device=DEVICE)],
+                    dim=1
+                )
+                time.sleep(0.015)
+
+        yield "data:[DONE]\n\n"
+
+    return StreamingResponse(event_stream(), media_type="text/event-stream")
+
+# ======================
+# INTERFAZ WEB (TU HTML COMPLETO, SIN QUITAR NADA)
 # ======================
 @app.get("/", response_class=HTMLResponse)
 def chat_ui():
     return """
+
 <!DOCTYPE html>
 <html lang="es">
 <head>
@@ -621,12 +686,11 @@ window.onload = () => userInput.focus();
 </script>
 </body>
 </html>
-"""
 
-# ======================
-# ENTRYPOINT
-# ======================
+"""
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
-
-
+    uvicorn.run(
+        app,
+        host="0.0.0.0",
+        port=int(os.environ.get("PORT", 7860))
+    )