Upload model (1).py

model (1).py

@@ -0,0 +1,318 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+
+
+class RotaryPositionalEmbedding(nn.Module):
+    """RoPE - Rotary Position Embedding"""
+    
+    def __init__(self, dim, max_seq_len=2048, 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 MultiHeadSelfAttention(nn.Module):
+    """Multi-Head Self-Attention con RoPE y Flash Attention"""
+    
+    def __init__(self, d_model, n_heads, dropout=0.1, max_seq_len=2048):
+        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.q_linear = nn.Linear(d_model, d_model, bias=False)
+        self.k_linear = nn.Linear(d_model, d_model, bias=False)
+        self.v_linear = nn.Linear(d_model, d_model, 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):
+        batch_size, seq_len, d_model = x.size()
+        
+        Q = self.q_linear(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
+        K = self.k_linear(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
+        V = self.v_linear(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
+        
+        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)
+        
+        if self.flash and mask is None:
+            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)
+
+
+class SwiGLU(nn.Module):
+    """SwiGLU activation - Mejor que GELU"""
+    
+    def __init__(self, d_model, d_ff, dropout=0.1):
+        super().__init__()
+        hidden_dim = int(d_ff * 2 / 3)
+        self.w1 = nn.Linear(d_model, hidden_dim, bias=False)
+        self.w2 = nn.Linear(hidden_dim, d_model, bias=False)
+        self.w3 = nn.Linear(d_model, hidden_dim, 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 eficiente 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 mejorado"""
+    
+    def __init__(self, d_model, n_heads, d_ff, dropout=0.1, max_seq_len=2048, use_swiglu=True):
+        super().__init__()
+        self.attention = MultiHeadSelfAttention(d_model, n_heads, dropout, max_seq_len)
+        
+        if use_swiglu:
+            self.feed_forward = SwiGLU(d_model, d_ff, dropout)
+        else:
+            self.feed_forward = nn.Sequential(
+                nn.Linear(d_model, d_ff),
+                nn.GELU(),
+                nn.Dropout(dropout),
+                nn.Linear(d_ff, d_model)
+            )
+        
+        self.norm1 = RMSNorm(d_model)
+        self.norm2 = RMSNorm(d_model)
+        
+    def forward(self, x, mask=None):
+        x = x + self.attention(self.norm1(x), mask)
+        x = x + self.feed_forward(self.norm2(x))
+        return x
+
+
+class MTPMiniModel(nn.Module):
+    """MTP Mini - Modelo 20x más grande e inteligente con anti-alucinación"""
+    
+    def __init__(self, vocab_size, d_model=1024, n_layers=24, n_heads=16, 
+                 d_ff=4096, max_seq_len=2048, dropout=0.15, use_swiglu=True,
+                 use_confidence_scoring=True, use_gradient_checkpointing=False):
+        super().__init__()
+        
+        self.vocab_size = vocab_size
+        self.d_model = d_model
+        self.max_seq_len = max_seq_len
+        self.use_confidence_scoring = use_confidence_scoring
+        self.use_gradient_checkpointing = use_gradient_checkpointing
+        
+        self.token_embedding = nn.Embedding(vocab_size, d_model)
+        self.dropout = nn.Dropout(dropout)
+        
+        self.blocks = nn.ModuleList([
+            TransformerBlock(d_model, n_heads, d_ff, dropout, max_seq_len, use_swiglu)
+            for _ in range(n_layers)
+        ])
+        
+        self.norm_f = RMSNorm(d_model)
+        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
+        
+        # Weight tying
+        self.lm_head.weight = self.token_embedding.weight
+        
+        # Confidence scoring head
+        if use_confidence_scoring:
+            self.confidence_head = nn.Sequential(
+                nn.Linear(d_model, d_model // 2),
+                nn.ReLU(),
+                nn.Dropout(dropout),
+                nn.Linear(d_model // 2, 1),
+                nn.Sigmoid()
+            )
+        
+        self.apply(self._init_weights)
+        
+    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, use_eos_weight=False, eos_weight=2.0, return_confidence=False):
+        batch_size, seq_len = input_ids.size()
+        
+        mask = torch.tril(torch.ones(seq_len, seq_len, device=input_ids.device)).view(1, 1, seq_len, seq_len)
+        
+        x = self.dropout(self.token_embedding(input_ids))
+        
+        for block in self.blocks:
+            if self.use_gradient_checkpointing and self.training:
+                x = torch.utils.checkpoint.checkpoint(block, x, mask, use_reentrant=False)
+            else:
+                x = block(x, mask)
+        
+        x = self.norm_f(x)
+        logits = self.lm_head(x)
+        
+        confidence = None
+        if self.use_confidence_scoring and return_confidence:
+            confidence = self.confidence_head(x)
+        
+        loss = None
+        if targets is not None:
+            if use_eos_weight:
+                weights = torch.ones(self.vocab_size, device=logits.device)
+                weights[3] = eos_weight
+                loss = F.cross_entropy(
+                    logits.view(-1, self.vocab_size), 
+                    targets.view(-1),
+                    weight=weights,
+                    label_smoothing=0.15
+                )
+            else:
+                loss = F.cross_entropy(
+                    logits.view(-1, self.vocab_size), 
+                    targets.view(-1),
+                    label_smoothing=0.15
+                )
+        
+        if return_confidence:
+            return logits, loss, confidence
+        return logits, loss
+    
+    def generate(self, input_ids, max_new_tokens=300, temperature=0.65, 
+                 top_k=50, top_p=0.9, repetition_penalty=1.2,
+                 min_length=30, eos_token_id=3,
+                 use_confidence_filter=True, min_confidence=0.3,
+                 use_entropy_threshold=True, max_entropy=4.0):
+        """Generación con anti-alucinación"""
+        self.eval()
+        
+        generated = input_ids.clone()
+        generated_text_tokens = 0
+        
+        with torch.no_grad():
+            for step in range(max_new_tokens):
+                input_ids_cond = generated if generated.size(1) <= self.max_seq_len else generated[:, -self.max_seq_len:]
+                
+                logits, _, confidence = self(input_ids_cond, return_confidence=True)
+                logits = logits[:, -1, :].clone()
+                
+                # Confidence filtering
+                if use_confidence_filter and confidence is not None:
+                    conf_score = confidence[:, -1, :].item()
+                    if conf_score < min_confidence:
+                        temperature = min(temperature * 1.2, 1.0)
+                
+                # 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 repeticiones recientes
+                if generated.size(1) > 15:
+                    recent_tokens = generated[0, -15:].tolist()
+                    for token_id in set(recent_tokens):
+                        count = recent_tokens.count(token_id)
+                        if count > 3:
+                            logits[0, token_id] -= count * 3.0
+                
+                # Control de longitud
+                if generated_text_tokens < min_length:
+                    logits[0, eos_token_id] = float('-inf')
+                else:
+                    eos_boost = (generated_text_tokens - min_length) * 0.15
+                    logits[0, eos_token_id] += eos_boost
+                
+                # Temperature
+                logits = logits / temperature
+                
+                # Top-k
+                if top_k > 0:
+                    v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                    logits[logits < v[:, [-1]]] = float('-inf')
+                
+                # Top-p
+                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')
+                
+                # Entropy threshold
+                probs = F.softmax(logits, dim=-1)
+                if use_entropy_threshold:
+                    entropy = -(probs * torch.log(probs + 1e-10)).sum(dim=-1)
+                    if entropy.item() > max_entropy:
+                        temperature = max(temperature * 0.7, 0.3)
+                        logits = logits / temperature
+                        probs = F.softmax(logits, dim=-1)
+                
+                # Sample
+                next_token = torch.multinomial(probs, num_samples=1)
+                
+                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):
+        return sum(p.numel() for p in self.parameters() if p.requires_grad)