Create pipeline/mdlm/train.py

pipeline/mdlm/train.py

@@ -0,0 +1,179 @@
+"""
+MDLM Training Loop — Train on governed corpus, measure convergence.
+
+Usage:
+  python -m pipeline.mdlm.train --corpus corpus/CORPUS-FINAL --schedule A --epochs 50
+
+Convergence metric: per-epoch accuracy on masked token prediction.
+PHI_CONTRACTION_RATE threshold: loss ratio between consecutive epochs.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import random
+import sys
+import time
+from pathlib import Path
+
+sys.stdout.reconfigure(encoding="utf-8", errors="replace")
+
+try:
+    import torch
+    import torch.optim as optim
+    from torch.utils.data import DataLoader, TensorDataset
+    HAS_TORCH = True
+except ImportError:
+    HAS_TORCH = False
+    print("PyTorch not available. Install with: pip install torch")
+
+from pipeline.mdlm.tokenizer import load_corpus, pad_sequence, VOCAB_SIZE, decode
+from pipeline.mdlm.model import (
+    MaskingSchedule, StructureModel, compute_loss, generate,
+)
+
+PHI_CONTRACTION_RATE = 0.381966  # (3 - sqrt(5)) / 2
+
+
+def train(
+    corpus_dir: str,
+    schedule: MaskingSchedule,
+    epochs: int = 50,
+    batch_size: int = 64,
+    lr: float = 1e-3,
+    total_timesteps: int = 100,
+    max_len: int = 40,
+    device_name: str = "cpu",
+):
+    """Train MDLM on governed corpus."""
+    print(f"=== MDLM Training ===")
+    print(f"Schedule: {schedule.value} ({schedule.name})")
+    print(f"Corpus: {corpus_dir}")
+    print(f"Epochs: {epochs}, Batch: {batch_size}, LR: {lr}")
+    print()
+
+    # Load and encode corpus
+    sequences = load_corpus(corpus_dir)
+    print(f"Loaded {len(sequences)} sequences")
+
+    # Pad to fixed length
+    padded = [pad_sequence(seq, max_len) for seq in sequences]
+    data = torch.tensor(padded, dtype=torch.long)
+    dataset = TensorDataset(data)
+    loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
+
+    # Model
+    device = torch.device(device_name)
+    model = StructureModel(
+        vocab_size=VOCAB_SIZE, d_model=128, nhead=4,
+        num_layers=4, max_len=max_len,
+    ).to(device)
+
+    param_count = sum(p.numel() for p in model.parameters())
+    print(f"Model parameters: {param_count:,}")
+    print()
+
+    optimizer = optim.Adam(model.parameters(), lr=lr)
+
+    # Training loop
+    history = []
+    prev_loss = float("inf")
+
+    for epoch in range(1, epochs + 1):
+        model.train()
+        epoch_loss = 0.0
+        epoch_batches = 0
+
+        for (batch,) in loader:
+            batch = batch.to(device)
+            # Random timestep per batch
+            timestep = random.randint(0, total_timesteps)
+
+            loss = compute_loss(model, batch, schedule, timestep, total_timesteps)
+
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+            epoch_loss += loss.item()
+            epoch_batches += 1
+
+        avg_loss = epoch_loss / max(epoch_batches, 1)
+
+        # Convergence ratio: loss_t / loss_{t-1}
+        ratio = avg_loss / prev_loss if prev_loss > 0 and prev_loss != float("inf") else 1.0
+        prev_loss = avg_loss
+
+        # Regime classification
+        if ratio < PHI_CONTRACTION_RATE:
+            regime = "CONTRACTING"
+        elif ratio < 0.8:
+            regime = "BOUNDARY"
+        elif ratio < 1.0:
+            regime = "CRITICAL"
+        else:
+            regime = "DIVERGENT"
+
+        history.append({
+            "epoch": epoch,
+            "loss": avg_loss,
+            "ratio": ratio,
+            "regime": regime,
+        })
+
+        if epoch <= 5 or epoch % 5 == 0 or epoch == epochs:
+            print(f"  Epoch {epoch:3d}: loss={avg_loss:.4f}  ratio={ratio:.4f}  [{regime}]")
+
+    print()
+
+    # Generate samples
+    print("Generating 5 samples...")
+    model.eval()
+    samples = generate(model, 5, max_len, schedule, total_timesteps)
+    for i in range(5):
+        seq = samples[i].tolist()
+        print(f"  Sample {i}: {decode(seq)}")
+    print()
+
+    # Summary
+    final = history[-1]
+    contracting_epochs = sum(1 for h in history if h["regime"] == "CONTRACTING")
+    print(f"=== SUMMARY ===")
+    print(f"Final loss: {final['loss']:.4f}")
+    print(f"Final ratio: {final['ratio']:.4f} ({final['regime']})")
+    print(f"Contracting epochs: {contracting_epochs}/{len(history)}")
+    print(f"Schedule: {schedule.value} ({schedule.name})")
+
+    return history
+
+
+def main():
+    if not HAS_TORCH:
+        sys.exit(1)
+
+    parser = argparse.ArgumentParser(description="MDLM Training")
+    parser.add_argument("--corpus", required=True, help="Corpus directory")
+    parser.add_argument("--schedule", default="A", choices=["A", "B", "C", "D"])
+    parser.add_argument("--epochs", type=int, default=50)
+    parser.add_argument("--batch-size", type=int, default=64)
+    parser.add_argument("--lr", type=float, default=1e-3)
+    parser.add_argument("--device", default="cpu")
+    parser.add_argument("--output", default=None, help="Save history to JSON")
+    args = parser.parse_args()
+
+    schedule = MaskingSchedule(args.schedule)
+    history = train(
+        args.corpus, schedule, args.epochs, args.batch_size,
+        args.lr, device_name=args.device,
+    )
+
+    if args.output:
+        with open(args.output, "w") as f:
+            json.dump(history, f, indent=2)
+        print(f"History saved to {args.output}")
+
+
+if __name__ == "__main__":
+    main()