Create pipeline/mdlm/model.py

pipeline/mdlm/model.py

@@ -0,0 +1,292 @@
+"""
+MDLM — Masked Diffusion Language Model for governed structures.
+
+Architecture:
+  - Small transformer encoder (4 layers, 128 dim, 4 heads)
+  - Absorbing-state masking: tokens → <MASK> at rate alpha(t)
+  - Denoising: predict original token from masked sequence
+  - Loss: cross-entropy on masked positions (reweighted MLM)
+
+Masking schedules:
+  A: hierarchical hierarchical (Tier 1 → Tier 2 → Tier 3+readiness)
+  B: flat hierarchical (operators only, no readiness staging)
+  C: Uniform random
+  D: inverted inverted
+
+Per PLAN-GHA-002 §4.4: A > B > C > D predicted.
+"""
+
+from __future__ import annotations
+
+import math
+import random
+from enum import Enum
+
+try:
+    import torch
+    import torch.nn as nn
+    import torch.nn.functional as F
+    HAS_TORCH = True
+except ImportError:
+    HAS_TORCH = False
+
+from pipeline.mdlm.tokenizer import (
+    VOCAB_SIZE, MASK, PAD, NEVER_MASKED,
+    TIER_1_TOKENS, TIER_2_TOKENS, TIER_3_TOKENS,
+    get_tier, pad_sequence,
+)
+
+
+class MaskingSchedule(str, Enum):
+    """Masking schedule variants for the hierarchical hypothesis test."""
+    HIERARCHICAL = "A"  # hierarchical: Tier 1 → Tier 2 → CL+PreAttest
+    FLAT = "B"  # flat: operators only, uniform within tiers
+    UNIFORM = "C"           # uniform random over all maskable tokens
+    INVERTED = "D"      # inverted: CL first, Tier 1 last
+
+
+if HAS_TORCH:
+
+    class StructureModel(nn.Module):
+        """Small transformer for governed structure denoising."""
+
+        def __init__(
+            self,
+            vocab_size: int = VOCAB_SIZE,
+            d_model: int = 128,
+            nhead: int = 4,
+            num_layers: int = 4,
+            max_len: int = 40,
+            dropout: float = 0.1,
+        ):
+            super().__init__()
+            self.d_model = d_model
+            self.embedding = nn.Embedding(vocab_size, d_model, padding_idx=PAD)
+            self.pos_embedding = nn.Embedding(max_len, d_model)
+            self.timestep_embedding = nn.Embedding(1000, d_model)  # diffusion timestep
+
+            encoder_layer = nn.TransformerEncoderLayer(
+                d_model=d_model, nhead=nhead, dim_feedforward=d_model * 4,
+                dropout=dropout, batch_first=True,
+            )
+            self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
+            self.output_proj = nn.Linear(d_model, vocab_size)
+
+        def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+            """
+            x: (batch, seq_len) — token ids with some positions masked
+            t: (batch,) — diffusion timestep (0 = clean, T = fully masked)
+            Returns: (batch, seq_len, vocab_size) — logits for each position
+            """
+            B, L = x.shape
+            positions = torch.arange(L, device=x.device).unsqueeze(0).expand(B, -1)
+
+            h = self.embedding(x) + self.pos_embedding(positions)
+            h = h + self.timestep_embedding(t).unsqueeze(1)
+
+            pad_mask = (x == PAD)
+            h = self.transformer(h, src_key_padding_mask=pad_mask)
+            return self.output_proj(h)
+
+
+    def apply_mask(
+        tokens: torch.Tensor,
+        mask_rate: float,
+        schedule: MaskingSchedule,
+        timestep: int = 0,
+        total_timesteps: int = 100,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Apply masking schedule to a batch of token sequences.
+
+        Returns:
+            masked_tokens: tokens with some positions replaced by MASK
+            mask_positions: boolean tensor (True = was masked)
+        """
+        B, L = tokens.shape
+        masked = tokens.clone()
+        mask_positions = torch.zeros(B, L, dtype=torch.bool, device=tokens.device)
+
+        for b in range(B):
+            for i in range(L):
+                tok = tokens[b, i].item()
+                if tok in NEVER_MASKED:
+                    continue
+
+                tier = get_tier(tok)
+                if tier == 0:
+                    continue
+
+                # Compute per-tier mask probability based on schedule
+                p = _tier_mask_prob(tier, mask_rate, schedule, timestep, total_timesteps)
+
+                if random.random() < p:
+                    masked[b, i] = MASK
+                    mask_positions[b, i] = True
+
+        return masked, mask_positions
+
+
+    def _tier_mask_prob(
+        tier: int,
+        base_rate: float,
+        schedule: MaskingSchedule,
+        timestep: int,
+        total_timesteps: int,
+    ) -> float:
+        """Compute mask probability for a token based on its tier and the schedule."""
+        t_frac = timestep / max(total_timesteps, 1)  # 0 = clean, 1 = fully masked
+
+        if schedule == MaskingSchedule.UNIFORM:
+            return base_rate
+
+        if schedule == MaskingSchedule.HIERARCHICAL:
+            # Tier 1 (Tier 1): masked last, unmasked first
+            # Tier 3 (CL+PreAttest): masked first, unmasked last
+            if tier == 1:
+                return base_rate * max(0.0, (t_frac - 0.66) / 0.34) if t_frac > 0.66 else 0.0
+            elif tier == 2:
+                return base_rate * max(0.0, (t_frac - 0.33) / 0.34) if t_frac > 0.33 else 0.0
+            else:  # tier 3
+                return base_rate * min(1.0, t_frac / 0.33)
+
+        if schedule == MaskingSchedule.FLAT:
+            # Same as 369 but witness tokens are tier 2 priority
+            if tier == 1:
+                return base_rate * max(0.0, (t_frac - 0.66) / 0.34) if t_frac > 0.66 else 0.0
+            elif tier == 2:
+                return base_rate * max(0.0, (t_frac - 0.33) / 0.34) if t_frac > 0.33 else 0.0
+            else:
+                return base_rate * min(1.0, t_frac / 0.33)
+
+        if schedule == MaskingSchedule.INVERTED:
+            # Inverted: Tier 1 masked first
+            if tier == 1:
+                return base_rate * min(1.0, t_frac / 0.33)
+            elif tier == 2:
+                return base_rate * max(0.0, (t_frac - 0.33) / 0.34) if t_frac > 0.33 else 0.0
+            else:
+                return base_rate * max(0.0, (t_frac - 0.66) / 0.34) if t_frac > 0.66 else 0.0
+
+        return base_rate
+
+
+    def compute_loss(
+        model: StructureModel,
+        batch: torch.Tensor,
+        schedule: MaskingSchedule,
+        timestep: int,
+        total_timesteps: int = 100,
+        mask_rate: float = 0.5,
+    ) -> torch.Tensor:
+        """Compute MDLM denoising loss on a batch.
+
+        Loss = cross-entropy on masked positions only.
+        Returns zero loss if no positions were masked (avoids NaN).
+        """
+        device = next(model.parameters()).device
+        batch = batch.to(device)
+        t_tensor = torch.full((batch.size(0),), timestep, dtype=torch.long, device=device)
+
+        masked, mask_pos = apply_mask(batch, mask_rate, schedule, timestep, total_timesteps)
+
+        # If nothing was masked, return zero loss
+        if not mask_pos.any():
+            return torch.tensor(0.0, device=device, requires_grad=True)
+
+        logits = model(masked, t_tensor)
+
+        # Loss only on masked positions
+        loss = F.cross_entropy(
+            logits[mask_pos],
+            batch[mask_pos],
+            ignore_index=PAD,
+        )
+        return loss
+
+
+    def generate(
+        model: StructureModel,
+        num_samples: int,
+        seq_len: int,
+        schedule: MaskingSchedule,
+        total_timesteps: int = 50,
+        g_slots: int = 3,
+        s_slots: int = 4,
+        f_slots: int = 3,
+    ) -> torch.Tensor:
+        """Generate governed structures by template-guided iterative unmasking.
+
+        The channel_b frame is IMPOSED (governance), not learned:
+          <BOS> <G> [MASK slots] </G> <S> [MASK slots] </S> <F> [MASK slots] </F>
+                [witness MASK slots] <EOS>
+
+        The model fills in operator tokens and witness attestation status.
+        This respects PROPOSE ≠ PROMOTE: the frame is governance,
+        the content is what the kernel crystallizes.
+
+        g_slots, s_slots, f_slots: number of operator MASK slots per modality.
+        Should match the corpus distribution.
+        """
+        device = next(model.parameters()).device
+        from pipeline.mdlm.tokenizer import (
+            BOS, EOS, G_OPEN, G_CLOSE, S_OPEN, S_CLOSE, F_OPEN, F_CLOSE,
+            WIT_OFFSET, ATTESTED,
+        )
+
+        # Build template with configurable slot counts
+        template = [BOS, G_OPEN] + [MASK] * g_slots + [G_CLOSE,
+                    S_OPEN] + [MASK] * s_slots + [S_CLOSE,
+                    F_OPEN] + [MASK] * f_slots + [F_CLOSE]
+        # 7 witness pairs: WIT_TOKEN MASK
+        for w in range(7):
+            template.extend([WIT_OFFSET + w, MASK])
+        template.append(EOS)
+
+        # Pad to seq_len
+        while len(template) < seq_len:
+            template.append(PAD)
+        template = template[:seq_len]
+
+        samples = torch.tensor([template] * num_samples, dtype=torch.long, device=device)
+
+        model.eval()
+        with torch.no_grad():
+            for step in range(total_timesteps, -1, -1):
+                t_tensor = torch.full((num_samples,), step, dtype=torch.long, device=device)
+                logits = model(samples, t_tensor)
+                probs = F.softmax(logits, dim=-1)
+
+                t_frac = step / total_timesteps
+
+                for b in range(num_samples):
+                    for i in range(seq_len):
+                        if samples[b, i].item() != MASK:
+                            continue
+
+                        pred = torch.multinomial(probs[b, i], 1).item()
+                        tier = get_tier(pred)
+
+                        # Tier-based unmasking schedule
+                        should_unmask = False
+                        if schedule == MaskingSchedule.HIERARCHICAL:
+                            should_unmask = (tier == 1 and t_frac < 0.33) or \
+                                          (tier == 2 and 0.33 <= t_frac < 0.66) or \
+                                          (tier == 3 and t_frac >= 0.66) or \
+                                          (step == 0)  # unmask everything at final step
+                        else:
+                            should_unmask = True
+
+                        if should_unmask:
+                            samples[b, i] = pred
+
+            # Final pass: force-unmask any remaining MASK tokens
+            remaining = (samples == MASK)
+            if remaining.any():
+                t_tensor = torch.zeros((num_samples,), dtype=torch.long, device=device)
+                logits = model(samples, t_tensor)
+                for b in range(num_samples):
+                    for i in range(seq_len):
+                        if samples[b, i].item() == MASK:
+                            samples[b, i] = logits[b, i].argmax().item()
+
+        return samples