TinyTorch/tests/13_transformers/test_training_simple.py

"""
Simple end-to-end training test for transformers.

This test validates that a transformer can successfully learn from a tiny dataset,
demonstrating that the entire training pipeline (forward, loss, backward, update) works.
"""

import numpy as np
import sys
import time
from pathlib import Path

# Add parent directory to path for imports
sys.path.insert(0, str(Path(__file__).parent.parent.parent))

from tinytorch.core.tensor import Tensor
from tinytorch.core.autograd import enable_autograd
from tinytorch.core.optimizers import Adam
from tinytorch.core.losses import CrossEntropyLoss
from tinytorch.models.transformer import GPT
from tinytorch.text.tokenization import CharTokenizer


def test_transformer_memorization():
    """
    Test that a transformer can memorize a tiny dataset.
    
    Success criteria:
    - Loss decreases by at least 80% in 500 steps
    - No NaN/Inf losses
    - All parameters receive gradients
    - Training completes in reasonable time (<120s)
    """
    print("\n" + "="*70)
    print("TEST: Transformer Memorization Capability")
    print("="*70)
    
    # Tiny dataset (5 patterns)
    patterns = [
        "def add(a, b):\n    return a + b",
        "def sub(a, b):\n    return a - b",
        "for i in range(10):\n    print(i)",
        "if x > 0:\n    print('positive')",
        "numbers = [1, 2, 3, 4, 5]",
    ]
    
    # Create tokenizer
    tokenizer = CharTokenizer()
    tokenizer.build_vocab(patterns)
    print(f"   Vocabulary size: {tokenizer.vocab_size}")
    
    # Create model (small for fast testing)
    model = GPT(
        vocab_size=tokenizer.vocab_size,
        embed_dim=32,
        num_layers=1,
        num_heads=4,
        max_seq_len=64
    )
    
    num_params = sum(np.prod(p.shape) for p in model.parameters())
    print(f"   Model parameters: {num_params:,}")
    
    # Optimizer and loss
    optimizer = Adam(model.parameters(), lr=0.001)
    loss_fn = CrossEntropyLoss()
    
    # Encode and pad patterns
    max_len = 64
    encoded = []
    for p in patterns:
        tokens = tokenizer.encode(p)
        if len(tokens) > max_len:
            tokens = tokens[:max_len]
        else:
            tokens = tokens + [0] * (max_len - len(tokens))
        encoded.append(tokens)
    
    # Training
    print("   Training for 500 steps...")
    losses = []
    start_time = time.time()
    
    for step in range(500):
        # Sample random pattern
        tokens = encoded[np.random.randint(len(encoded))]
        x = Tensor(np.array([tokens[:-1]], dtype=np.int32))
        y = Tensor(np.array([tokens[1:]], dtype=np.int32))
        
        # Forward pass
        logits = model.forward(x)
        logits_flat = logits.reshape(len(tokens)-1, tokenizer.vocab_size)
        y_flat = y.reshape(len(tokens)-1)
        loss = loss_fn(logits_flat, y_flat)
        
        # Check for NaN/Inf
        assert not np.isnan(loss.data).any(), f"NaN loss at step {step}"
        assert not np.isinf(loss.data).any(), f"Inf loss at step {step}"
        
        # Backward pass
        optimizer.zero_grad()
        loss.backward()
        
        # Check gradients on first step
        if step == 0:
            params_with_grad = sum(1 for p in model.parameters() 
                                   if p.grad is not None and np.abs(p.grad).max() > 1e-10)
            total_params = len(model.parameters())
            assert params_with_grad == total_params, \
                f"Only {params_with_grad}/{total_params} parameters have gradients"
        
        # Gradient clipping
        for p in model.parameters():
            if p.grad is not None:
                p.grad = np.clip(p.grad, -1.0, 1.0)
        
        # Update
        optimizer.step()
        
        # Track loss
        losses.append(loss.data.item())
    
    elapsed = time.time() - start_time
    
    # Compute statistics
    initial_loss = losses[0]
    final_loss = np.mean(losses[-100:])
    loss_decrease_pct = ((initial_loss - final_loss) / initial_loss) * 100
    
    print(f"\n   Results:")
    print(f"   ├─ Initial loss: {initial_loss:.3f}")
    print(f"   ├─ Final loss: {final_loss:.3f}")
    print(f"   ├─ Loss decrease: {loss_decrease_pct:.1f}%")
    print(f"   └─ Training time: {elapsed:.1f}s")
    
    # Assertions
    assert elapsed < 120, f"Training too slow: {elapsed:.1f}s > 120s"
    assert loss_decrease_pct > 80, \
        f"Insufficient learning: loss decreased only {loss_decrease_pct:.1f}% (expected >80%)"
    assert final_loss < 0.5, \
        f"Final loss too high: {final_loss:.3f} (expected <0.5 for memorization)"
    
    print(f"\n✅ Transformer successfully memorized dataset!")
    print(f"   Loss decreased {loss_decrease_pct:.1f}% in {elapsed:.1f}s")
    return True


def test_transformer_convergence_rate():
    """
    Test that transformer converges at expected rate.
    
    This is a regression test to catch training instabilities.
    """
    print("\n" + "="*70)
    print("TEST: Transformer Convergence Rate")
    print("="*70)
    
    # Setup (same as memorization test)
    patterns = [
        "def add(a, b):\n    return a + b",
        "def sub(a, b):\n    return a - b",
    ]
    
    tokenizer = CharTokenizer()
    tokenizer.build_vocab(patterns)
    
    model = GPT(
        vocab_size=tokenizer.vocab_size,
        embed_dim=32,
        num_layers=1,
        num_heads=4,
        max_seq_len=64
    )
    
    optimizer = Adam(model.parameters(), lr=0.001)
    loss_fn = CrossEntropyLoss()
    
    # Encode patterns
    max_len = 64
    encoded = []
    for p in patterns:
        tokens = tokenizer.encode(p)
        if len(tokens) > max_len:
            tokens = tokens[:max_len]
        else:
            tokens = tokens + [0] * (max_len - len(tokens))
        encoded.append(tokens)
    
    # Train until loss < 0.1
    step = 0
    loss_val = float('inf')
    
    print(f"   Training until loss < 0.1...")
    
    while loss_val > 0.1 and step < 1000:
        tokens = encoded[np.random.randint(len(encoded))]
        x = Tensor(np.array([tokens[:-1]], dtype=np.int32))
        y = Tensor(np.array([tokens[1:]], dtype=np.int32))
        
        logits = model.forward(x)
        logits_flat = logits.reshape(len(tokens)-1, tokenizer.vocab_size)
        y_flat = y.reshape(len(tokens)-1)
        loss = loss_fn(logits_flat, y_flat)
        
        optimizer.zero_grad()
        loss.backward()
        
        for p in model.parameters():
            if p.grad is not None:
                p.grad = np.clip(p.grad, -1.0, 1.0)
        
        optimizer.step()
        
        loss_val = loss.data.item()
        step += 1
    
    print(f"   Reached loss < 0.1 in {step} steps")
    
    # Regression check: should converge in < 500 steps for 2 patterns
    assert step < 500, \
        f"Convergence too slow: {step} steps (expected <500). Training may be unstable."
    
    print(f"✅ Convergence rate is acceptable ({step} steps)")
    return True


if __name__ == "__main__":
    print("\n" + "="*70)
    print("TRANSFORMER TRAINING TEST SUITE")
    print("="*70)
    
    test_transformer_memorization()
    test_transformer_convergence_rate()
    
    print("\n" + "="*70)
    print("✅ ALL TRAINING TESTS PASSED")
    print("="*70 + "\n")