Delete model.py

model.py

@@ -1,331 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import math
-
-
-class RotaryPositionalEmbedding(nn.Module):
-    """RoPE - Rotary Position Embedding con scaling mejorado"""
-    
-    def __init__(self, dim, max_seq_len=4096, base=10000):
-        super().__init__()
-        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
-        self.register_buffer('inv_freq', inv_freq)
-        self.max_seq_len = max_seq_len
-        
-    def forward(self, seq_len, device):
-        t = torch.arange(seq_len, device=device).type_as(self.inv_freq)
-        freqs = torch.einsum('i,j->ij', t, self.inv_freq)
-        emb = torch.cat((freqs, freqs), dim=-1)
-        return emb.cos(), emb.sin()
-
-
-def apply_rotary_pos_emb(q, k, cos, sin):
-    """Aplica RoPE a queries y keys"""
-    def rotate_half(x):
-        x1, x2 = x.chunk(2, dim=-1)
-        return torch.cat((-x2, x1), dim=-1)
-    
-    q_embed = (q * cos) + (rotate_half(q) * sin)
-    k_embed = (k * cos) + (rotate_half(k) * sin)
-    return q_embed, k_embed
-
-
-class MultiQueryAttention(nn.Module):
-    """Multi-Query Attention (MQA) - Más eficiente que MHA"""
-    
-    def __init__(self, d_model, n_heads, dropout=0.1, max_seq_len=4096):
-        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
-        
-        # Multi-query: Q tiene múltiples heads, K y V tienen 1 head
-        self.q_linear = nn.Linear(d_model, d_model, bias=False)
-        self.k_linear = nn.Linear(d_model, self.d_k, bias=False)
-        self.v_linear = nn.Linear(d_model, self.d_k, bias=False)
-        self.out_linear = nn.Linear(d_model, d_model, bias=False)
-        
-        self.dropout = nn.Dropout(dropout)
-        self.attn_dropout = nn.Dropout(dropout)
-        self.rope = RotaryPositionalEmbedding(self.d_k, max_seq_len)
-        
-        self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention')
-        
-    def forward(self, x, mask=None, use_cache=False, past_kv=None):
-        batch_size, seq_len, d_model = x.size()
-        
-        # Q: [batch, seq, n_heads, d_k]
-        Q = self.q_linear(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
-        
-        # K, V: [batch, seq, d_k] -> expandir a [batch, n_heads, seq, d_k]
-        K = self.k_linear(x).unsqueeze(1).expand(-1, self.n_heads, -1, -1)
-        V = self.v_linear(x).unsqueeze(1).expand(-1, self.n_heads, -1, -1)
-        
-        # Apply RoPE
-        cos, sin = self.rope(seq_len, x.device)
-        cos = cos[None, None, :, :]
-        sin = sin[None, None, :, :]
-        Q, K = apply_rotary_pos_emb(Q, K, cos, sin)
-        
-        # KV cache para inferencia
-        if use_cache:
-            if past_kv is not None:
-                K = torch.cat([past_kv[0], K], dim=2)
-                V = torch.cat([past_kv[1], V], dim=2)
-            cache = (K, V)
-        else:
-            cache = None
-        
-        # Attention
-        if self.flash and mask is None and not use_cache:
-            context = F.scaled_dot_product_attention(
-                Q, K, V, 
-                attn_mask=None,
-                dropout_p=self.dropout.p if self.training else 0.0,
-                is_causal=True
-            )
-        else:
-            scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
-            if mask is not None:
-                scores = scores.masked_fill(mask == 0, float('-inf'))
-            attn_weights = F.softmax(scores, dim=-1)
-            attn_weights = self.attn_dropout(attn_weights)
-            context = torch.matmul(attn_weights, V)
-        
-        context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, d_model)
-        output = self.out_linear(context)
-        return self.dropout(output), cache
-
-
-class SwiGLU(nn.Module):
-    """SwiGLU activation con eficiencia mejorada"""
-    
-    def __init__(self, d_model, d_ff, dropout=0.1):
-        super().__init__()
-        # FFN de GPT-3: 4x expansion
-        self.w1 = nn.Linear(d_model, d_ff, bias=False)
-        self.w2 = nn.Linear(d_ff, d_model, bias=False)
-        self.w3 = nn.Linear(d_model, d_ff, bias=False)
-        self.dropout = nn.Dropout(dropout)
-        
-    def forward(self, x):
-        return self.w2(self.dropout(F.silu(self.w1(x)) * self.w3(x)))
-
-
-class RMSNorm(nn.Module):
-    """RMSNorm - Más estable que LayerNorm"""
-    
-    def __init__(self, dim, eps=1e-6):
-        super().__init__()
-        self.eps = eps
-        self.weight = nn.Parameter(torch.ones(dim))
-        
-    def forward(self, x):
-        norm = torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
-        return x * norm * self.weight
-
-
-class TransformerBlock(nn.Module):
-    """Transformer Block optimizado estilo GPT-3"""
-    
-    def __init__(self, d_model, n_heads, d_ff, dropout=0.1, max_seq_len=4096):
-        super().__init__()
-        self.attention = MultiQueryAttention(d_model, n_heads, dropout, max_seq_len)
-        self.feed_forward = SwiGLU(d_model, d_ff, dropout)
-        
-        self.norm1 = RMSNorm(d_model)
-        self.norm2 = RMSNorm(d_model)
-        
-    def forward(self, x, mask=None, use_cache=False, past_kv=None):
-        # Pre-norm architecture (mejor que post-norm)
-        attn_out, cache = self.attention(self.norm1(x), mask, use_cache, past_kv)
-        x = x + attn_out
-        x = x + self.feed_forward(self.norm2(x))
-        return x, cache
-
-
-class MTPModel(nn.Module):
-    """MTP 3 - Arquitectura mejorada nivel GPT-3"""
-    
-    def __init__(self, vocab_size, d_model=1024, n_layers=24, n_heads=16, 
-                 d_ff=4096, max_seq_len=2048, dropout=0.1):
-        super().__init__()
-        
-        self.vocab_size = vocab_size
-        self.d_model = d_model
-        self.max_seq_len = max_seq_len
-        
-        # Embeddings con escalado
-        self.token_embedding = nn.Embedding(vocab_size, d_model)
-        self.dropout = nn.Dropout(dropout)
-        
-        # Transformer blocks
-        self.blocks = nn.ModuleList([
-            TransformerBlock(d_model, n_heads, d_ff, dropout, max_seq_len)
-            for _ in range(n_layers)
-        ])
-        
-        # Final norm y projection
-        self.norm_f = RMSNorm(d_model)
-        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
-        
-        # Weight tying (reduce parámetros)
-        self.token_embedding.weight = self.lm_head.weight
-        
-        # Inicialización mejorada (GPT-3 style)
-        self.apply(self._init_weights)
-        
-        # Escalado especial para residual connections
-        for pn, p in self.named_parameters():
-            if pn.endswith('w2.weight') or pn.endswith('out_linear.weight'):
-                torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * n_layers))
-        
-    def _init_weights(self, module):
-        if isinstance(module, nn.Linear):
-            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
-            if module.bias is not None:
-                torch.nn.init.zeros_(module.bias)
-        elif isinstance(module, nn.Embedding):
-            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
-    
-    def forward(self, input_ids, targets=None):
-        batch_size, seq_len = input_ids.size()
-        
-        # Causal mask
-        mask = torch.tril(torch.ones(seq_len, seq_len, device=input_ids.device)).view(1, 1, seq_len, seq_len)
-        
-        # Embeddings con escalado
-        x = self.dropout(self.token_embedding(input_ids) * math.sqrt(self.d_model))
-        
-        # Transformer blocks
-        for block in self.blocks:
-            x, _ = block(x, mask)
-        
-        # Final norm y projection
-        x = self.norm_f(x)
-        logits = self.lm_head(x)
-        
-        loss = None
-        if targets is not None:
-            # Label smoothing para mejor generalización
-            loss = F.cross_entropy(
-                logits.view(-1, self.vocab_size), 
-                targets.view(-1),
-                label_smoothing=0.1,
-                ignore_index=-100
-            )
-        
-        return logits, loss
-    
-    @torch.no_grad()
-    def generate(self, input_ids, max_new_tokens=200, temperature=0.8, 
-                 top_k=50, top_p=0.95, repetition_penalty=1.2,
-                 min_length=30, eos_token_id=3):
-        """Generación optimizada con KV cache"""
-        self.eval()
-        
-        device = input_ids.device
-        generated = input_ids.clone()
-        past_kvs = [None] * len(self.blocks)
-        generated_text_tokens = 0
-        
-        for step in range(max_new_tokens):
-            # Use cache para tokens ya procesados
-            if step == 0:
-                current_input = generated
-                use_cache = False
-            else:
-                current_input = generated[:, -1:]
-                use_cache = True
-            
-            # Truncate si excede max_seq_len
-            if current_input.size(1) > self.max_seq_len:
-                current_input = current_input[:, -self.max_seq_len:]
-                use_cache = False
-                past_kvs = [None] * len(self.blocks)
-            
-            # Forward pass
-            batch_size, seq_len = current_input.size()
-            mask = torch.tril(torch.ones(seq_len, seq_len, device=device)).view(1, 1, seq_len, seq_len)
-            
-            x = self.token_embedding(current_input) * math.sqrt(self.d_model)
-            
-            new_past_kvs = []
-            for i, block in enumerate(self.blocks):
-                x, cache = block(x, mask, use_cache, past_kvs[i] if use_cache else None)
-                new_past_kvs.append(cache)
-            
-            if use_cache:
-                past_kvs = new_past_kvs
-            
-            x = self.norm_f(x)
-            logits = self.lm_head(x[:, -1, :])
-            
-            # Repetition penalty
-            if repetition_penalty != 1.0:
-                for token_id in set(generated[0].tolist()):
-                    if logits[0, token_id] < 0:
-                        logits[0, token_id] *= repetition_penalty
-                    else:
-                        logits[0, token_id] /= repetition_penalty
-            
-            # Penalizar tokens muy repetidos
-            if generated.size(1) > 20:
-                recent = generated[0, -20:].tolist()
-                for token_id in set(recent):
-                    count = recent.count(token_id)
-                    if count > 3:
-                        logits[0, token_id] -= count * 3.0
-            
-            # Control de longitud mínima
-            if generated_text_tokens < min_length:
-                logits[0, eos_token_id] = float('-inf')
-            else:
-                # Boost EOS gradualmente
-                eos_boost = min((generated_text_tokens - min_length) * 0.15, 3.0)
-                logits[0, eos_token_id] += eos_boost
-            
-            # Temperature scaling
-            logits = logits / temperature
-            
-            # Top-k filtering
-            if top_k > 0:
-                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
-                logits[logits < v[:, [-1]]] = float('-inf')
-            
-            # Top-p (nucleus) filtering
-            if top_p < 1.0:
-                sorted_logits, sorted_indices = torch.sort(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_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
-                logits[indices_to_remove] = float('-inf')
-            
-            # Sample
-            probs = F.softmax(logits, dim=-1)
-            next_token = torch.multinomial(probs, num_samples=1)
-            
-            # Check EOS
-            if next_token.item() == eos_token_id and generated_text_tokens >= min_length:
-                break
-            
-            generated = torch.cat([generated, next_token], dim=1)
-            generated_text_tokens += 1
-        
-        return generated
-    
-    def count_parameters(self):
-        """Cuenta parámetros entrenables"""
-        return sum(p.numel() for p in self.parameters() if p.requires_grad)
-    
-    def get_num_params(self, non_embedding=True):
-        """Cuenta parámetros excluyendo embeddings si se requiere"""
-        n_params = sum(p.numel() for p in self.parameters())
-        if non_embedding:
-            n_params -= self.token_embedding.weight.numel()
-        return n_params