Upload model.py with huggingface_hub

model.py

@@ -0,0 +1,144 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, Tuple
+import math
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def _norm(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        return self._norm(x.float()).type_as(x) * self.weight
+
+class SSMBlock(nn.Module):
+    """
+    Una implementación simplificada de una capa de State Space Model (SSM) Selectivo.
+    A diferencia del Transformer, esta capa tiene memoria LINEAL.
+    """
+    def __init__(self, dim: int, state_dim: int = 16, expand: int = 2):
+        super().__init__()
+        self.dim = dim
+        self.state_dim = state_dim
+        self.inner_dim = expand * dim
+        
+        # Proyecciones de entrada
+        self.in_proj = nn.Linear(dim, self.inner_dim * 2, bias=False)
+        
+        # Convolución 1D para capturar contexto local (estilo Mamba)
+        self.conv1d = nn.Conv1d(
+            in_channels=self.inner_dim,
+            out_channels=self.inner_dim,
+            kernel_size=4,
+            groups=self.inner_dim,
+            padding=3
+        )
+        
+        # Parámetros del SSM (Simplificado)
+        # Delta (dt): El paso de tiempo reactivo
+        self.dt_proj = nn.Linear(self.inner_dim, self.inner_dim, bias=True)
+        
+        # matrices A y B (Simplificadas para ser selectivas)
+        self.A_log = nn.Parameter(torch.log(torch.arange(1, state_dim + 1).float().repeat(self.inner_dim, 1)))
+        self.B_proj = nn.Linear(self.inner_dim, state_dim, bias=False)
+        self.C_proj = nn.Linear(self.inner_dim, state_dim, bias=False)
+        
+        # Proyección de salida
+        self.out_proj = nn.Linear(self.inner_dim, dim, bias=False)
+
+    def forward(self, x: torch.Tensor):
+        # x: [Batch, SeqLen, Dim]
+        b, l, d = x.shape
+        
+        # 1. Proyección inicial y split
+        x_and_res = self.in_proj(x) # [B, L, 2 * InnerDim]
+        x_inner, res = x_and_res.split(self.inner_dim, dim=-1)
+        
+        # 2. Convolución local
+        x_inner = x_inner.transpose(1, 2) # [B, InnerDim, L]
+        x_inner = self.conv1d(x_inner)[:, :, :l] # Recortar padding
+        x_inner = x_inner.transpose(1, 2) # [B, L, InnerDim]
+        x_inner = F.silu(x_inner) # Activación Swish
+        
+        # 3. Mecanismo SSM Selectivo (Simplificado)
+        # En una implementación real (como Mamba), esto se hace con un kernel de GPU 
+        # para ser ultra rápido. Aquí usamos un bucle o aproximación para entenderlo.
+        
+        dt = F.softplus(self.dt_proj(x_inner)) # [B, L, InnerDim]
+        A = -torch.exp(self.A_log) # [InnerDim, StateDim]
+        B = self.B_proj(x_inner) # [B, L, StateDim]
+        C = self.C_proj(x_inner) # [B, L, StateDim]
+        
+        # El "Estado Oculto" del modelo (Memory)
+        # Aquí es donde ocurre la magia: el estado tiene tamaño fijo [B, InnerDim, StateDim]
+        state = torch.zeros(b, self.inner_dim, self.state_dim, device=x.device)
+        y = torch.zeros(b, l, self.inner_dim, device=x.device)
+        
+        # Escaneo Secuencial (Selective Scan)
+        # Esto reemplaza a la Atención del Transformer.
+        for t in range(l):
+            # dt_t: [B, InnerDim]
+            dt_t = dt[:, t, :].unsqueeze(-1)
+            # x_t: [B, InnerDim]
+            x_t = x_inner[:, t, :].unsqueeze(-1)
+            # B_t: [B, StateDim]
+            B_t = B[:, t, :].unsqueeze(1)
+            
+            # Discretización (Aproximación de Euler)
+            A_bar = torch.exp(A.unsqueeze(0) * dt_t) # [B, InnerDim, StateDim]
+            B_bar = dt_t * B_t
+            
+            # Actualizar Estado: h = A*h + B*x
+            state = A_bar * state + B_bar * x_t
+            
+            # Salida: y = C*h
+            C_t = C[:, t, :].unsqueeze(-1) # [B, StateDim, 1]
+            y[:, t, :] = torch.matmul(state, C_t).squeeze(-1)
+        
+        # 4. Combinar con el residuo y proyectar salida
+        out = y * F.silu(res)
+        return self.out_proj(out)
+
+class TransformerKiller(nn.Module):
+    """
+    Arquitectura basada en SSM (Mamba-like).
+    No usa Atención. Contexto teóricamente infinito.
+    """
+    def __init__(self, vocab_size: int, dim: int, n_layers: int, state_dim: int = 16):
+        super().__init__()
+        self.tok_embeddings = nn.Embedding(vocab_size, dim)
+        
+        self.layers = nn.ModuleList([
+            nn.ModuleDict({
+                'norm': RMSNorm(dim),
+                'ssm': SSMBlock(dim, state_dim=state_dim)
+            }) for _ in range(n_layers)
+        ])
+        
+        self.norm_f = RMSNorm(dim)
+        self.output = nn.Linear(dim, vocab_size, bias=False)
+
+    def forward(self, tokens: torch.Tensor):
+        x = self.tok_embeddings(tokens)
+        
+        for layer in self.layers:
+            # Conexión residual + SSM
+            x = x + layer['ssm'](layer['norm'](x))
+            
+        x = self.norm_f(x)
+        return self.output(x)
+
+if __name__ == "__main__":
+    # Test de dimensiones
+    model = TransformerKiller(vocab_size=100, dim=128, n_layers=4)
+    test_input = torch.randint(0, 100, (2, 50))
+    output = model(test_input)
+    print(f"Input shape: {test_input.shape}")
+    print(f"SSM Output shape: {output.shape}")
+    # Nota como no hay límite de BLOCK_SIZE en el forward, 
+    # solo el que dicte tu RAM, pero el coste es Lineal!