Create handler.py

handler.py

@@ -0,0 +1,59 @@
+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)
+embed_size = 256
+num_heads = 8
+num_layers = 4
+block_size = 256
+
+class MiniGPT(nn.Module):
+    def __init__(self, v_size=256):
+        super().__init__()
+        self.token_embedding = nn.Embedding(v_size, embed_size)
+        self.pos_embedding = nn.Embedding(block_size, embed_size)
+        self.blocks = nn.ModuleList([
+            nn.TransformerEncoderLayer(d_model=embed_size, nhead=num_heads, 
+                                       dim_feedforward=embed_size*4, batch_first=True, 
+                                       dropout=0.1, norm_first=True) for _ in range(num_layers)
+        ])
+        self.ln = nn.LayerNorm(embed_size)
+        self.fc_out = nn.Linear(embed_size, v_size)
+
+    def forward(self, idx):
+        B, T = idx.shape
+        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
+
+class EndpointHandler:
+    def __init__(self, path=""):
+        # Cargar el modelo
+        self.model = MiniGPT()
+        checkpoint = os.path.join(path, "mini_gpt.pth")
+        self.model.load_state_dict(torch.load(checkpoint, map_location="cpu"))
+        self.model.eval()
+
+    def __call__(self, data):
+        # Procesar la entrada
+        inputs = data.get("inputs", "")
+        tokens = [ord(c) if ord(c) < 256 else 32 for c in inputs]
+        
+        # Generar (versión simplificada de tu función generate)
+        res = ""
+        for _ in range(50): # Generamos 50 caracteres
+            idx = torch.tensor([tokens[-block_size:]])
+            with torch.no_grad():
+                logits = self.model(idx)
+            logits = logits[:, -1, :] / 0.7 # temp fija 0.7
+            probs = F.softmax(logits, dim=-1)
+            nxt = torch.multinomial(probs, 1).item()
+            if nxt == ord('\n'): break
+            tokens.append(nxt)
+            res += chr(nxt)
+            
+        return [{"generated_text": res}]