Upload 5 files

config.yaml

@@ -0,0 +1,24 @@
+# MTP Mini Configuration - Optimized
+
+model:
+  vocab_size: 4000
+  d_model: 384          # Aumentado de 256 para más capacidad
+  n_layers: 6           # Aumentado de 4 para más profundidad
+  n_heads: 6            # Aumentado de 4 para mejor atención
+  d_ff: 1536            # Aumentado de 1024 (4x d_model)
+  max_seq_len: 256      # Aumentado de 128 para contexto más largo
+  dropout: 0.1
+
+training:
+  batch_size: 8         # Aumentado de 4 para mejor gradiente
+  epochs: 50            # Aumentado de 20 para más entrenamiento
+  learning_rate: 0.0001 # Reducido de 0.0003 para estabilidad
+  weight_decay: 0.01
+  max_grad_norm: 0.5    # Reducido de 1.0 para mejor estabilidad
+  num_threads: 4
+  save_every: 10
+  warmup_steps: 100     # Nuevo: warmup del learning rate
+  use_lr_scheduler: true # Nuevo: learning rate decay
+
+data:
+  corpus_path: corpus/mtp_mini_corpus.jsonl

model.py

@@ -0,0 +1,202 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+
+
+class MultiHeadSelfAttention(nn.Module):
+    """Multi-Head Self-Attention mechanism"""
+    
+    def __init__(self, d_model, n_heads, dropout=0.1):
+        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)
+        self.k_linear = nn.Linear(d_model, d_model)
+        self.v_linear = nn.Linear(d_model, d_model)
+        self.out_linear = nn.Linear(d_model, d_model)
+        
+        self.dropout = nn.Dropout(dropout)
+        
+    def forward(self, x, mask=None):
+        batch_size, seq_len, d_model = x.size()
+        
+        # Linear projections
+        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)
+        
+        # Scaled dot-product attention
+        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.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 output
+
+
+class FeedForward(nn.Module):
+    """Position-wise Feed-Forward Network"""
+    
+    def __init__(self, d_model, d_ff, dropout=0.1):
+        super().__init__()
+        self.linear1 = nn.Linear(d_model, d_ff)
+        self.linear2 = nn.Linear(d_ff, d_model)
+        self.dropout = nn.Dropout(dropout)
+        
+    def forward(self, x):
+        return self.linear2(self.dropout(F.gelu(self.linear1(x))))
+
+
+class TransformerBlock(nn.Module):
+    """Single Transformer Decoder Block"""
+    
+    def __init__(self, d_model, n_heads, d_ff, dropout=0.1):
+        super().__init__()
+        self.attention = MultiHeadSelfAttention(d_model, n_heads, dropout)
+        self.feed_forward = FeedForward(d_model, d_ff, dropout)
+        self.ln1 = nn.LayerNorm(d_model)
+        self.ln2 = nn.LayerNorm(d_model)
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+        
+    def forward(self, x, mask=None):
+        # Self-attention with residual connection
+        attn_output = self.attention(self.ln1(x), mask)
+        x = x + self.dropout1(attn_output)
+        
+        # Feed-forward with residual connection
+        ff_output = self.feed_forward(self.ln2(x))
+        x = x + self.dropout2(ff_output)
+        
+        return x
+
+
+class MTPMiniModel(nn.Module):
+    """MTP Mini - GPT-style Transformer Language Model"""
+    
+    def __init__(self, vocab_size, d_model=256, n_layers=4, n_heads=4, 
+                 d_ff=1024, max_seq_len=128, dropout=0.1):
+        super().__init__()
+        
+        self.vocab_size = vocab_size
+        self.d_model = d_model
+        self.max_seq_len = max_seq_len
+        
+        # Token embeddings
+        self.token_embedding = nn.Embedding(vocab_size, d_model)
+        
+        # Positional embeddings (learnable)
+        self.position_embedding = nn.Embedding(max_seq_len, d_model)
+        
+        # Transformer blocks
+        self.blocks = nn.ModuleList([
+            TransformerBlock(d_model, n_heads, d_ff, dropout)
+            for _ in range(n_layers)
+        ])
+        
+        # Final layer norm
+        self.ln_f = nn.LayerNorm(d_model)
+        
+        # Output projection to vocabulary
+        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
+        
+        # Weight tying
+        self.lm_head.weight = self.token_embedding.weight
+        
+        self.dropout = nn.Dropout(dropout)
+        
+        # Initialize weights
+        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)
+        elif isinstance(module, nn.LayerNorm):
+            torch.nn.init.zeros_(module.bias)
+            torch.nn.init.ones_(module.weight)
+    
+    def forward(self, input_ids, targets=None):
+        batch_size, seq_len = input_ids.size()
+        
+        # Create causal mask
+        mask = torch.tril(torch.ones(seq_len, seq_len, device=input_ids.device)).view(1, 1, seq_len, seq_len)
+        
+        # Token embeddings + positional embeddings
+        positions = torch.arange(0, seq_len, device=input_ids.device).unsqueeze(0)
+        tok_emb = self.token_embedding(input_ids)
+        pos_emb = self.position_embedding(positions)
+        x = self.dropout(tok_emb + pos_emb)
+        
+        # Pass through transformer blocks
+        for block in self.blocks:
+            x = block(x, mask)
+        
+        # Final layer norm
+        x = self.ln_f(x)
+        
+        # Project to vocabulary
+        logits = self.lm_head(x)
+        
+        # Calculate loss if targets provided
+        loss = None
+        if targets is not None:
+            loss = F.cross_entropy(logits.view(-1, self.vocab_size), targets.view(-1))
+        
+        return logits, loss
+    
+    def generate(self, input_ids, max_new_tokens=50, temperature=1.0, top_k=50, top_p=0.9):
+        """Autoregressive generation with sampling"""
+        self.eval()
+        
+        with torch.no_grad():
+            for _ in range(max_new_tokens):
+                # Crop to max_seq_len
+                input_ids_cond = input_ids if input_ids.size(1) <= self.max_seq_len else input_ids[:, -self.max_seq_len:]
+                
+                # Forward pass
+                logits, _ = self(input_ids_cond)
+                logits = logits[:, -1, :] / 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 from distribution
+                probs = F.softmax(logits, dim=-1)
+                next_token = torch.multinomial(probs, num_samples=1)
+                
+                # Append to sequence
+                input_ids = torch.cat([input_ids, next_token], dim=1)
+        
+        return input_ids
+    
+    def count_parameters(self):
+        """Count trainable parameters"""
+        return sum(p.numel() for p in self.parameters() if p.requires_grad)

mtp_mini.pkl

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

mtp_tokenizer.model

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

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()