Upload 6 files

config.yaml

@@ -0,0 +1,38 @@
+data:
+  augmentation_prob: 0.3
+  corpus_path: data.jsonl
+  max_text_length: 3000
+  min_text_length: 30
+  use_augmentation: true
+  validation_split: 0.15
+generation:
+  default_max_tokens: 200
+  default_repetition_penalty: 1.2
+  default_temperature: 0.8
+  default_top_k: 50
+  default_top_p: 0.95
+  min_response_length: 30
+model:
+  d_ff: 4096
+  d_model: 1024
+  dropout: 0.1
+  max_seq_len: 2048
+  n_heads: 16
+  n_layers: 24
+  vocab_size: 8000
+training:
+  accumulation_steps: 8
+  batch_size: 2
+  epochs: 30
+  label_smoothing: 0.1
+  learning_rate: 0.0003
+  max_grad_norm: 1.0
+  min_delta: 0.0005
+  min_lr: 1.0e-06
+  num_threads: 4
+  patience: 7
+  save_every: 3
+  use_amp: true
+  use_lr_scheduler: true
+  warmup_steps: 500
+  weight_decay: 0.1

dataset.py

@@ -0,0 +1,124 @@
+import torch
+from torch.utils.data import Dataset
+import json
+import random
+
+
+class MTPDataset(Dataset):
+    """Dataset optimizado con augmentación inteligente"""
+    
+    def __init__(self, corpus_path, tokenizer, max_seq_len=2048, 
+                 use_augmentation=True, augmentation_prob=0.3):
+        self.tokenizer = tokenizer
+        self.max_seq_len = max_seq_len
+        self.use_augmentation = use_augmentation
+        self.augmentation_prob = augmentation_prob
+        self.data = []
+        
+        # Load corpus
+        print(f"   → Cargando corpus: {corpus_path}")
+        with open(corpus_path, 'r', encoding='utf-8') as f:
+            for line in f:
+                line = line.strip()
+                if line:
+                    try:
+                        entry = json.loads(line)
+                        if 'instruction' in entry and 'response' in entry:
+                            # Validar que no estén vacíos
+                            if entry['instruction'].strip() and entry['response'].strip():
+                                self.data.append(entry)
+                    except json.JSONDecodeError:
+                        continue
+        
+        print(f"   ✓ Cargados {len(self.data)} ejemplos válidos")
+        if use_augmentation:
+            print(f"   ✓ Augmentación activada (prob={augmentation_prob})")
+    
+    def __len__(self):
+        return len(self.data)
+    
+    def augment_text(self, text):
+        """Augmentación mejorada de texto"""
+        if not self.use_augmentation or random.random() > self.augmentation_prob:
+            return text
+        
+        # 1. Variación en espacios y formato
+        if random.random() < 0.3:
+            text = text.strip()
+        
+        # 2. Variación en puntuación final
+        if random.random() < 0.25:
+            if text.endswith('.'):
+                # A veces remover punto final
+                if random.random() < 0.5:
+                    text = text[:-1]
+            elif not text.endswith(('.', '!', '?', ':')):
+                # A veces agregar punto
+                if random.random() < 0.5:
+                    text = text + '.'
+        
+        # 3. Variación en mayúsculas iniciales (muy ocasional)
+        if random.random() < 0.1 and len(text) > 0:
+            if text[0].isupper():
+                text = text[0].lower() + text[1:]
+            elif text[0].islower():
+                text = text[0].upper() + text[1:]
+        
+        return text
+    
+    def __getitem__(self, idx):
+        entry = self.data[idx]
+        
+        instruction = entry['instruction']
+        response = entry['response']
+        
+        # Aplicar augmentación
+        instruction = self.augment_text(instruction)
+        response = self.augment_text(response)
+        
+        # Formato optimizado para entrenamiento
+        full_text = f"### Instrucción:\n{instruction}\n\n### Respuesta:\n{response}"
+        
+        # Tokenize
+        tokens = self.tokenizer.encode(full_text)
+        
+        # Add BOS and EOS
+        tokens = [self.tokenizer.bos_id()] + tokens + [self.tokenizer.eos_id()]
+        
+        # Truncate if too long (mantener BOS y EOS)
+        if len(tokens) > self.max_seq_len:
+            tokens = [tokens[0]] + tokens[1:self.max_seq_len-1] + [self.tokenizer.eos_id()]
+        
+        # Pad token ID será -100 para ignorar en loss
+        input_ids = torch.tensor(tokens[:-1], dtype=torch.long)
+        target_ids = torch.tensor(tokens[1:], dtype=torch.long)
+        
+        return input_ids, target_ids
+
+
+def collate_fn(batch, pad_id=0):
+    """Collate function optimizada con padding dinámico"""
+    input_ids = [item[0] for item in batch]
+    target_ids = [item[1] for item in batch]
+    
+    # Find max length in this batch (dynamic padding)
+    max_len = max(len(ids) for ids in input_ids)
+    
+    # Pad sequences
+    input_ids_padded = []
+    target_ids_padded = []
+    
+    for inp, tgt in zip(input_ids, target_ids):
+        pad_len = max_len - len(inp)
+        
+        # Pad input with pad_id
+        input_ids_padded.append(
+            torch.cat([inp, torch.full((pad_len,), pad_id, dtype=torch.long)])
+        )
+        
+        # Pad target with -100 (ignore_index in CrossEntropyLoss)
+        target_ids_padded.append(
+            torch.cat([tgt, torch.full((pad_len,), -100, dtype=torch.long)])
+        )
+    
+    return torch.stack(input_ids_padded), torch.stack(target_ids_padded)

model.py

@@ -0,0 +1,331 @@
+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

mtp_tokenizer.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e13a9965e5241f5f677f6689f568b4bbaf68c16c2605d4907d0d608dede95adb
+size 124448

tokenizer.py

@@ -0,0 +1,138 @@
+import sentencepiece as spm
+import os
+import json
+
+
+class MTPTokenizer:
+    """Tokenizer using SentencePiece BPE"""
+    
+    def __init__(self, model_path=None):
+        self.sp = None
+        self.model_path = model_path
+        
+        if model_path and os.path.exists(model_path):
+            self.load(model_path)
+    
+    def train(self, corpus_path, vocab_size=4000, model_prefix='mtp_tokenizer'):
+        """Train SentencePiece BPE tokenizer on corpus"""
+        
+        # Extract text from JSONL corpus
+        texts = []
+        with open(corpus_path, 'r', encoding='utf-8') as f:
+            for line in f:
+                data = json.loads(line)
+                if 'instruction' in data:
+                    texts.append(data['instruction'])
+                if 'response' in data:
+                    texts.append(data['response'])
+        
+        # Save temporary text file
+        temp_file = 'temp_corpus.txt'
+        with open(temp_file, 'w', encoding='utf-8') as f:
+            f.write('\n'.join(texts))
+        
+        # Calculate optimal vocab size based on corpus
+        total_chars = sum(len(text) for text in texts)
+        max_vocab = min(vocab_size, int(total_chars * 0.15))  # Heuristic: ~15% of chars
+        
+        print(f"   → Corpus stats: {len(texts)} texts, {total_chars} characters")
+        print(f"   → Adjusted vocab size: {max_vocab} (requested: {vocab_size})")
+        
+        # Train SentencePiece with adjusted parameters
+        try:
+            spm.SentencePieceTrainer.train(
+                input=temp_file,
+                model_prefix=model_prefix,
+                vocab_size=max_vocab,
+                model_type='bpe',
+                pad_id=0,
+                unk_id=1,
+                bos_id=2,
+                eos_id=3,
+                character_coverage=1.0,
+                normalization_rule_name='identity',
+                num_threads=4,
+                split_digits=True,
+                allow_whitespace_only_pieces=False,
+                byte_fallback=False,
+                max_sentencepiece_length=16
+            )
+        except RuntimeError as e:
+            if "Vocabulary size too high" in str(e):
+                # Extract suggested max from error and retry
+                import re
+                match = re.search(r'value <= (\d+)', str(e))
+                if match:
+                    suggested_max = int(match.group(1))
+                    print(f"   → Retrying with vocab size: {suggested_max}")
+                    spm.SentencePieceTrainer.train(
+                        input=temp_file,
+                        model_prefix=model_prefix,
+                        vocab_size=suggested_max,
+                        model_type='bpe',
+                        pad_id=0,
+                        unk_id=1,
+                        bos_id=2,
+                        eos_id=3,
+                        character_coverage=1.0,
+                        normalization_rule_name='identity',
+                        num_threads=4,
+                        split_digits=True,
+                        allow_whitespace_only_pieces=False,
+                        byte_fallback=False,
+                        max_sentencepiece_length=16
+                    )
+                else:
+                    raise
+            else:
+                raise
+        
+        # Clean up
+        os.remove(temp_file)
+        
+        # Load the trained model
+        self.model_path = f"{model_prefix}.model"
+        self.load(self.model_path)
+        
+        print(f"✓ Tokenizer trained: {self.vocab_size()} tokens")
+        print(f"✓ Model saved: {self.model_path}")
+    
+    def load(self, model_path):
+        """Load trained tokenizer"""
+        self.sp = spm.SentencePieceProcessor()
+        self.sp.load(model_path)
+        self.model_path = model_path
+    
+    def encode(self, text):
+        """Encode text to token IDs"""
+        if self.sp is None:
+            raise ValueError("Tokenizer not loaded. Train or load a model first.")
+        return self.sp.encode_as_ids(text)
+    
+    def decode(self, ids):
+        """Decode token IDs to text"""
+        if self.sp is None:
+            raise ValueError("Tokenizer not loaded. Train or load a model first.")
+        return self.sp.decode_ids(ids)
+    
+    def vocab_size(self):
+        """Get vocabulary size"""
+        if self.sp is None:
+            return 0
+        return self.sp.get_piece_size()
+    
+    def bos_id(self):
+        """Beginning of sentence token ID"""
+        return self.sp.bos_id()
+    
+    def eos_id(self):
+        """End of sentence token ID"""
+        return self.sp.eos_id()
+    
+    def pad_id(self):
+        """Padding token ID"""
+        return self.sp.pad_id()
+    
+    def unk_id(self):
+        """Unknown token ID"""
+        return self.sp.unk_id()