Create app.py

app.py

@@ -0,0 +1,696 @@
+import os
+import sys
+import torch
+import json
+import time
+import gc
+import re
+from fastapi import FastAPI, Request
+from fastapi.responses import HTMLResponse, StreamingResponse
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel, Field
+from huggingface_hub import snapshot_download
+import uvicorn
+import math
+import torch.nn as nn
+import torch.nn.functional as F
+import sentencepiece as spm
+
+# ======================
+# CONFIGURACIÓN DE DISPOSITIVO
+# ======================
+if torch.cuda.is_available():
+    DEVICE = "cuda"
+    print("✅ GPU NVIDIA detectada. Usando CUDA.")
+else:
+    DEVICE = "cpu"
+    print("⚠️ GPU no detectada. Usando CPU (puede ser más lento).")
+
+if DEVICE == "cpu":
+    torch.set_num_threads(max(1, os.cpu_count() // 2))
+
+torch.set_grad_enabled(False)
+
+# CAMBIA ESTO POR EL NOMBRE DE TU REPO EN HUGGING FACE
+MODEL_REPO = "TeszenAI/MTP-3.2"
+
+# ======================
+# FUNCIONES DE LIMPIEZA Y CONTROL DE CALIDAD
+# ======================
+
+def truncate_greeting_response(text: str) -> str:
+    """
+    Para respuestas de saludo, trunca SOLO en el primer PUNTO (.)
+    No usa signos de exclamación o interrogación.
+    """
+    if not text:
+        return text
+    
+    # Buscar el primer PUNTO (.)
+    end_match = re.search(r'\.', text)
+    
+    if end_match:
+        # Cortar justo después del punto
+        end_pos = end_match.end()
+        truncated = text[:end_pos].strip()
+        return truncated
+    
+    # Si no hay punto, devolver solo primeras 80 caracteres
+    if len(text) > 80:
+        return text[:80] + "..."
+    return text
+
+def clean_response(text: str, user_input: str = "") -> str:
+    """Limpia la respuesta del modelo"""
+    if not text:
+        return ""
+    
+    # Eliminar repeticiones excesivas
+    words = text.split()
+    cleaned_words = []
+    last_word = ""
+    repeat_count = 0
+    
+    for word in words:
+        if word == last_word:
+            repeat_count += 1
+            if repeat_count > 2:
+                continue
+        else:
+            last_word = word
+            repeat_count = 0
+        cleaned_words.append(word)
+    
+    text = " ".join(cleaned_words)
+    
+    # Eliminar caracteres raros
+    text = re.sub(r'(.)\1{4,}', r'\1\1', text)
+    
+    # Detectar si es un saludo
+    is_greeting = user_input.lower().strip() in ["hola", "hola!", "hola.", "buenas", "saludos", "hola?"]
+    
+    if is_greeting and text:
+        # Para saludos, truncar SOLO en el primer PUNTO (.)
+        punct_match = re.search(r'\.', text)
+        if punct_match:
+            text = text[:punct_match.end()].strip()
+        else:
+            # Si no hay punto, tomar solo la primera oración o 60 caracteres
+            first_sentence = text.split('.')[0].strip()
+            if len(first_sentence) > 5:
+                text = first_sentence
+            elif len(text) > 60:
+                text = text[:60]
+    
+    # Si la respuesta es muy corta o vacía
+    if len(text.strip()) < 5:
+        if is_greeting:
+            return "¡Hola! ¿En qué puedo ayudarte?"
+        return "Lo siento, no pude generar una respuesta clara. ¿Podrías reformular tu pregunta?"
+    
+    # Eliminar espacios múltiples
+    text = re.sub(r'\s+', ' ', text).strip()
+    
+    return text
+
+# ======================
+# DEFINIR ARQUITECTURA DEL MODELO (MTP)
+# ======================
+class LayerNorm(nn.Module):
+    def __init__(self, d_model: int, eps: float = 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):
+        mean = x.mean(-1, keepdim=True)
+        std = x.std(-1, keepdim=True)
+        return self.weight * (x - mean) / (std + self.eps) + self.bias
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self, d_model: int, n_heads: int, dropout: float = 0.1):
+        super().__init__()
+        assert d_model % n_heads == 0
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.d_k = d_model // n_heads
+        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)
+        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)
+        K = self.w_k(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
+        V = self.w_v(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
+        scores = torch.matmul(Q, K.transpose(-2, -1)) / self.scale
+        if mask is not None:
+            scores = scores.masked_fill(mask == 0, float('-inf'))
+        attn_weights = F.softmax(scores, dim=-1)
+        attn_weights = self.dropout(attn_weights)
+        attn_output = torch.matmul(attn_weights, V)
+        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)
+        return self.w_o(attn_output)
+
+class FeedForward(nn.Module):
+    def __init__(self, d_model: int, d_ff: int, dropout: float = 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: int, n_heads: int, d_ff: int, dropout: float = 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):
+        attn_output = self.attention(x, mask)
+        x = x + self.dropout1(attn_output)
+        x = self.norm1(x)
+        ff_output = self.feed_forward(x)
+        x = x + self.dropout2(ff_output)
+        x = self.norm2(x)
+        return x
+
+class PositionalEncoding(nn.Module):
+    def __init__(self, d_model: int, max_len: int = 5000):
+        super().__init__()
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
+        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 = 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.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)
+        logits = self.lm_head(x)
+        return logits
+    
+    def generate(self, input_ids, max_new_tokens=150, temperature=0.8, top_k=50, top_p=0.9, repetition_penalty=1.1):
+        """Genera texto token por token"""
+        generated = input_ids
+        
+        for step 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()
+            
+            # EOS ID común para SentencePiece
+            if next_token == 2 or next_token == 3:
+                break
+            
+            generated = torch.cat([generated, torch.tensor([[next_token]], device=generated.device)], dim=1)
+        
+        return generated
+
+# ======================
+# DESCARGA Y CARGA DEL MODELO
+# ======================
+print(f"📦 Descargando modelo desde {MODEL_REPO}...")
+repo_path = snapshot_download(
+    repo_id=MODEL_REPO,
+    repo_type="model",
+    local_dir="mtp_repo"
+)
+
+# Cargar configuraci��n
+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)
+else:
+    config = {
+        "vocab_size": 5000,
+        "d_model": 256,
+        "n_heads": 8,
+        "n_layers": 6,
+        "d_ff": 1024,
+        "dropout": 0.1,
+        "max_len": 512
+    }
+
+# Cargar tokenizador
+tokenizer_path = os.path.join(repo_path, "mtp_tokenizer.model")
+if not os.path.exists(tokenizer_path):
+    print(f"❌ Tokenizador no encontrado en {tokenizer_path}")
+    sys.exit(1)
+
+sp = spm.SentencePieceProcessor()
+sp.load(tokenizer_path)
+VOCAB_SIZE = sp.get_piece_size()
+
+# Actualizar vocab_size en config
+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']}")
+
+model = MTPModel(**config)
+model.to(DEVICE)
+
+# Cargar pesos del modelo
+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 del modelo cargados")
+else:
+    print(f"⚠️ No se encontró {model_path}, usando pesos aleatorios")
+
+model.eval()
+
+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 API",
+    description="API para modelo de lenguaje MTP",
+    version="1.0"
+)
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+class PromptRequest(BaseModel):
+    text: str = Field(..., max_length=2000, description="Texto de entrada")
+    max_tokens: int = Field(default=150, ge=10, le=250, description="Tokens máximos a generar")
+    temperature: float = Field(default=0.3, 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")
+
+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"
+
+# ======================
+# GESTIÓN DE CARGA
+# ======================
+ACTIVE_REQUESTS = 0
+
+class MTPTokenizer:
+    """Wrapper para el tokenizador de SentencePiece"""
+    def __init__(self, sp_model):
+        self.sp = sp_model
+    
+    def encode(self, text):
+        return self.sp.encode(text)
+    
+    def decode(self, tokens):
+        return self.sp.decode(tokens)
+    
+    def bos_id(self):
+        return self.sp.bos_id()
+    
+    def eos_id(self):
+        return self.sp.eos_id()
+    
+    def pad_id(self):
+        return self.sp.pad_id()
+
+tokenizer_wrapper = MTPTokenizer(sp)
+
+@app.post("/generate")
+async def generate(req: PromptRequest):
+    """Endpoint principal de generación de texto"""
+    global ACTIVE_REQUESTS
+    ACTIVE_REQUESTS += 1
+    
+    user_input = req.text.strip()
+    if not user_input:
+        ACTIVE_REQUESTS -= 1
+        return {"reply": "", "tokens_generated": 0}
+    
+    # Detectar si es un saludo
+    is_greeting = user_input.lower().strip() in ["hola", "hola!", "hola.", "buenas", "saludos", "hola?"]
+    
+    # Si es saludo, usar menos tokens
+    max_tokens = 30 if is_greeting else req.max_tokens
+
+    full_prompt = build_prompt(user_input)
+    tokens = 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=max_tokens,
+                temperature=req.temperature,
+                top_k=req.top_k,
+                top_p=req.top_p,
+                repetition_penalty=req.repetition_penalty
+            )
+
+        gen_tokens = output_ids[0, len(tokens):].tolist()
+        
+        # Filtrar tokens inválidos
+        safe_tokens = [t for t in gen_tokens if 0 <= t < VOCAB_SIZE]
+        
+        if safe_tokens:
+            response = tokenizer_wrapper.decode(safe_tokens).strip()
+        else:
+            response = ""
+        
+        # Limpiar respuesta
+        response = clean_response(response, user_input)
+        
+        # Si la respuesta sigue vacía o es muy corta, usar respuesta por defecto
+        if len(response) < 3:
+            if is_greeting:
+                response = "¡Hola! ¿En qué puedo ayudarte?"
+            else:
+                response = "Lo siento, no pude generar una respuesta. ¿Podrías reformular tu pregunta?"
+
+        return {
+            "reply": response,
+            "tokens_generated": len(safe_tokens),
+            "model": "MTP"
+        }
+    
+    except Exception as e:
+        print(f"❌ Error durante generación: {e}")
+        if is_greeting:
+            fallback = "¡Hola! ¿En qué puedo ayudarte?"
+        else:
+            fallback = "Lo siento, ocurrió un error al procesar tu solicitud."
+        return {
+            "reply": fallback,
+            "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",
+        "device": DEVICE,
+        "active_requests": ACTIVE_REQUESTS,
+        "vocab_size": VOCAB_SIZE
+    }
+
+@app.get("/info")
+def model_info():
+    return {
+        "model_name": "MTP",
+        "version": "1.0",
+        "architecture": config,
+        "parameters": sum(p.numel() for p in model.parameters()),
+        "device": DEVICE
+    }
+
+# ======================
+# INTERFAZ WEB
+# ======================
+@app.get("/", response_class=HTMLResponse)
+def chat_ui():
+    return """
+<!DOCTYPE html>
+<html lang="es">
+<head>
+<meta charset="UTF-8">
+<meta name="viewport" content="width=device-width, initial-scale=1.0">
+<title>MTP - Asistente IA</title>
+<style>
+* { margin: 0; padding: 0; box-sizing: border-box; }
+body {
+    background: #131314;
+    font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
+    height: 100vh;
+    display: flex;
+    flex-direction: column;
+}
+.chat-header {
+    padding: 16px 20px;
+    background: #1E1F20;
+    border-bottom: 1px solid #2a2b2e;
+}
+.chat-header h1 {
+    color: white;
+    font-size: 1.2rem;
+    font-weight: 500;
+}
+.chat-messages {
+    flex: 1;
+    overflow-y: auto;
+    padding: 20px;
+    display: flex;
+    flex-direction: column;
+    gap: 16px;
+}
+.message {
+    display: flex;
+    gap: 12px;
+    max-width: 80%;
+}
+.message.user {
+    align-self: flex-end;
+    flex-direction: row-reverse;
+}
+.message-content {
+    padding: 10px 16px;
+    border-radius: 18px;
+    font-size: 0.95rem;
+    line-height: 1.4;
+}
+.user .message-content {
+    background: #4a9eff;
+    color: white;
+    border-radius: 18px 4px 18px 18px;
+}
+.bot .message-content {
+    background: #1E1F20;
+    color: #e3e3e3;
+    border-radius: 4px 18px 18px 18px;
+}
+.chat-input-container {
+    padding: 16px 20px;
+    background: #1E1F20;
+    border-top: 1px solid #2a2b2e;
+}
+.input-wrapper {
+    display: flex;
+    gap: 12px;
+    max-width: 800px;
+    margin: 0 auto;
+}
+#messageInput {
+    flex: 1;
+    padding: 12px 16px;
+    background: #2a2b2e;
+    border: none;
+    border-radius: 24px;
+    color: white;
+    font-size: 0.95rem;
+    outline: none;
+}
+#messageInput::placeholder {
+    color: #888;
+}
+#sendBtn {
+    padding: 12px 24px;
+    background: #4a9eff;
+    border: none;
+    border-radius: 24px;
+    color: white;
+    font-weight: 500;
+    cursor: pointer;
+    transition: opacity 0.2s;
+}
+#sendBtn:hover { opacity: 0.9; }
+#sendBtn:disabled {
+    opacity: 0.5;
+    cursor: not-allowed;
+}
+.typing {
+    display: flex;
+    gap: 4px;
+    padding: 10px 16px;
+}
+.typing span {
+    width: 8px;
+    height: 8px;
+    background: #888;
+    border-radius: 50%;
+    animation: bounce 1.4s infinite ease-in-out;
+}
+.typing span:nth-child(1) { animation-delay: -0.32s; }
+.typing span:nth-child(2) { animation-delay: -0.16s; }
+@keyframes bounce {
+    0%, 80%, 100% { transform: scale(0); }
+    40% { transform: scale(1); }
+}
+</style>
+</head>
+<body>
+<div class="chat-header">
+    <h1>🤖 MTP - Asistente IA</h1>
+</div>
+<div class="chat-messages" id="chatMessages">
+    <div class="message bot">
+        <div class="message-content">¡Hola! Soy MTP, tu asistente de IA. ¿En qué puedo ayudarte hoy?</div>
+    </div>
+</div>
+<div class="chat-input-container">
+    <div class="input-wrapper">
+        <input type="text" id="messageInput" placeholder="Escribe tu mensaje..." autocomplete="off">
+        <button id="sendBtn">Enviar</button>
+    </div>
+</div>
+<script>
+const chatMessages = document.getElementById('chatMessages');
+const messageInput = document.getElementById('messageInput');
+const sendBtn = document.getElementById('sendBtn');
+let isLoading = false;
+
+function addMessage(text, isUser) {
+    const div = document.createElement('div');
+    div.className = `message ${isUser ? 'user' : 'bot'}`;
+    div.innerHTML = `<div class="message-content">${text}</div>`;
+    chatMessages.appendChild(div);
+    chatMessages.scrollTop = chatMessages.scrollHeight;
+    return div;
+}
+
+function addTypingIndicator() {
+    const div = document.createElement('div');
+    div.className = 'message bot';
+    div.id = 'typingIndicator';
+    div.innerHTML = `<div class="typing"><span></span><span></span><span></span></div>`;
+    chatMessages.appendChild(div);
+    chatMessages.scrollTop = chatMessages.scrollHeight;
+}
+
+function removeTypingIndicator() {
+    const indicator = document.getElementById('typingIndicator');
+    if (indicator) indicator.remove();
+}
+
+async function sendMessage() {
+    const text = messageInput.value.trim();
+    if (!text || isLoading) return;
+    
+    messageInput.value = '';
+    addMessage(text, true);
+    isLoading = true;
+    sendBtn.disabled = true;
+    addTypingIndicator();
+    
+    try {
+        const response = await fetch('/generate', {
+            method: 'POST',
+            headers: { 'Content-Type': 'application/json' },
+            body: JSON.stringify({ text: text })
+        });
+        const data = await response.json();
+        removeTypingIndicator();
+        addMessage(data.reply, false);
+    } catch (error) {
+        removeTypingIndicator();
+        addMessage('Error de conexión. Intenta de nuevo.', false);
+    } finally {
+        isLoading = false;
+        sendBtn.disabled = false;
+        messageInput.focus();
+    }
+}
+
+messageInput.addEventListener('keypress', (e) => {
+    if (e.key === 'Enter') sendMessage();
+});
+sendBtn.addEventListener('click', sendMessage);
+messageInput.focus();
+</script>
+</body>
+</html>
+"""
+
+if __name__ == "__main__":
+    port = int(os.environ.get("PORT", 7860))
+    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")
+    
+    uvicorn.run(
+        app,
+        host="0.0.0.0",
+        port=port,
+        log_level="info"
+    )