Update handler.py

handler.py

@@ -1,9 +1,10 @@
+from typing import Dict, List, Any
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import os
 
-# REPETIMOS TU ARQUITECTURA AQUÍ (Exactamente igual que en tu script)
+# Configuración idéntica a tu script
 embed_size = 256
 num_heads = 8
 num_layers = 4
@@ -23,33 +24,32 @@ class MiniGPT(nn.Module):
         self.fc_out = nn.Linear(embed_size, v_size)
 
     def forward(self, idx):
-        B, T = idx.shape
+        T = idx.shape[1]
         x = self.token_embedding(idx) + self.pos_embedding(torch.arange(T, device="cpu"))[None, :, :]
         mask = torch.triu(torch.ones(T, T, device="cpu"), diagonal=1).bool()
         for block in self.blocks: x = block(x, src_mask=mask)
-        logits = self.fc_out(self.ln(x))
-        return logits
+        return self.fc_out(self.ln(x))
 
 class EndpointHandler:
     def __init__(self, path=""):
-        # Cargar el modelo
         self.model = MiniGPT()
-        checkpoint = os.path.join(path, "pytorch_model.bin")
-        self.model.load_state_dict(torch.load(checkpoint, map_location="cpu"))
+        # Buscamos el archivo de pesos
+        checkpoint_path = os.path.join(path, "pytorch_model.bin")
+        self.model.load_state_dict(torch.load(checkpoint_path, map_location="cpu"))
         self.model.eval()
 
-    def __call__(self, data):
-        # Procesar la entrada
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         inputs = data.get("inputs", "")
-        tokens = [ord(c) if ord(c) < 256 else 32 for c in inputs]
+        if not inputs: return [{"generated_text": ""}]
         
-        # Generar (versión simplificada de tu función generate)
+        tokens = [ord(c) if ord(c) < 256 else 32 for c in inputs]
         res = ""
-        for _ in range(50): # Generamos 50 caracteres
+        
+        for _ in range(30): # Generamos 30 caracteres para probar
             idx = torch.tensor([tokens[-block_size:]])
             with torch.no_grad():
                 logits = self.model(idx)
-            logits = logits[:, -1, :] / 0.7 # temp fija 0.7
+            logits = logits[:, -1, :] / 0.7
             probs = F.softmax(logits, dim=-1)
             nxt = torch.multinomial(probs, 1).item()
             if nxt == ord('\n'): break