Create pipeline/mdlm/decoder.py

pipeline/mdlm/decoder.py

@@ -0,0 +1,293 @@
+"""
+Phase 5: Constrained Decoder — EXECUTE phase of the GGP.
+
+Takes a committed governed structure (from PROMOTE) and generates
+natural language within the validity envelope.
+
+Architecture: Small transformer decoder conditioned on governed operator
+tokens. The governed structure is the prompt; the output is prose that
+expresses the structure in natural language.
+
+This is NOT a general-purpose LLM. It generates governed prose —
+text whose semantic content is constrained to what the governed permits.
+The decoder cannot introduce implicit authority structures because
+the governed frame doesn't encode them.
+
+Training data: (structure tokens, source text) pairs extracted from
+the decomposition pipeline.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Iterator
+
+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 (
+    encode as encode_gov, VOCAB_SIZE as STRUCT_VOCAB_SIZE,
+    BOS, EOS, PAD, TOKEN_NAMES,
+)
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# PAIRED DATA EXTRACTION
+# ═══════════════════════════════════════════════════════════════════════════════
+
+@dataclass
+class FrameProsePair:
+    """A (governed structure, source prose) pair for decoder training."""
+    gov_tokens: list[int]  # Encoded governed structure
+    prose: str               # Original source text
+    source_id: str
+
+
+def extract_pairs_from_pipeline(
+    corpus_dir: str | Path,
+    theory_dir: str | Path | None = None,
+) -> list[FrameProsePair]:
+    """Extract structure-prose pairs by re-running the pipeline with text capture.
+
+    Since the emitted JSONL doesn't store the original text, we re-run
+    the decomposition and capture both the governed and the source segment.
+    """
+    import sys
+    sys.path.insert(0, ".")
+
+    from pipeline.ingest.chat_archive import ingest_conversation_file
+    from pipeline.stages.s2_classify import classify, Classification
+    from pipeline.stages.s3_decompose import decompose
+    from pipeline.stages.s4_validate import validate_and_score, TigStatus, Verdict
+
+    pairs = []
+
+    if theory_dir:
+        theory_path = Path(theory_dir)
+        for conv_file in sorted(theory_path.glob("conv_*.json")):
+            try:
+                for seg in ingest_conversation_file(conv_file):
+                    c = classify(seg)
+                    if c.classification != Classification.TECHNICAL:
+                        continue
+                    ex = decompose(c)
+                    if ex is None:
+                        continue
+                    r = validate_and_score(ex)
+                    if r.tig_status != TigStatus.TRUE:
+                        continue
+
+                    # Build pair
+                    struct_dict = {
+                        "channel_a": {"operators": [
+                            {"operator": e.operator.canonical_name, "evidence": e.evidence}
+                            for e in ex.channel_a.operators.expressions
+                        ]},
+                        "channel_b": {"operators": [
+                            {"operator": e.operator.canonical_name, "evidence": e.evidence}
+                            for e in ex.channel_b.operators.expressions
+                        ]},
+                        "channel_c": {"operators": [
+                            {"operator": e.operator.canonical_name, "evidence": e.evidence}
+                            for e in ex.channel_c.operators.expressions
+                        ]},
+                        "witnesses": {
+                            w.canonical_name: {"attested": a.attested, "evidence": a.evidence}
+                            for w, a in ex.witnesses.attestations.items()
+                        },
+                    }
+                    gov_tokens = encode_gov(struct_dict)
+
+                    pairs.append(FrameProsePair(
+                        gov_tokens=gov_tokens,
+                        prose=seg.text[:512],  # Cap at 512 chars for training
+                        source_id=ex.provenance.source_id,
+                    ))
+            except Exception:
+                continue
+
+    return pairs
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# PROSE TOKENIZER (character-level for simplicity)
+# ══════════════════��════════════════════════════════════════════════════════════
+
+PROSE_PAD = 0
+PROSE_BOS = 1
+PROSE_EOS = 2
+PROSE_UNK = 3
+PROSE_VOCAB_OFFSET = 4
+
+# Build vocab from printable ASCII + common unicode
+PROSE_CHARS = (
+    " !\"#$%&'()*+,-./0123456789:;<=>?@"
+    "ABCDEFGHIJKLMNOPQRSTUVWXYZ[\\]^_`"
+    "abcdefghijklmnopqrstuvwxyz{|}~"
+)
+PROSE_VOCAB_SIZE = PROSE_VOCAB_OFFSET + len(PROSE_CHARS)
+_CHAR_TO_ID = {c: i + PROSE_VOCAB_OFFSET for i, c in enumerate(PROSE_CHARS)}
+
+
+def encode_prose(text: str, max_len: int = 256) -> list[int]:
+    """Encode prose as character-level token IDs."""
+    tokens = [PROSE_BOS]
+    for ch in text[:max_len - 2]:
+        tokens.append(_CHAR_TO_ID.get(ch, PROSE_UNK))
+    tokens.append(PROSE_EOS)
+    return tokens
+
+
+def decode_prose(token_ids: list[int]) -> str:
+    """Decode character-level token IDs back to text."""
+    id_to_char = {v: k for k, v in _CHAR_TO_ID.items()}
+    chars = []
+    for tid in token_ids:
+        if tid in (PROSE_PAD, PROSE_BOS, PROSE_EOS):
+            continue
+        if tid == PROSE_UNK:
+            chars.append("?")
+        else:
+            chars.append(id_to_char.get(tid, "?"))
+    return "".join(chars)
+
+
+def pad_prose(tokens: list[int], max_len: int) -> list[int]:
+    """Pad or truncate prose tokens to fixed length."""
+    if len(tokens) >= max_len:
+        return tokens[:max_len]
+    return tokens + [PROSE_PAD] * (max_len - len(tokens))
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# CONSTRAINED DECODER MODEL
+# ═══════════════════════════════════════════════════════════════════════════════
+
+if HAS_TORCH:
+
+    class ConstrainedDecoder(nn.Module):
+        """Transformer decoder conditioned on governed structure.
+
+        Encoder: processes governed token sequence (the committed structure)
+        Decoder: generates prose character-by-character within the envelope
+
+        The structure tokens serve as cross-attention keys — the decoder
+        can only attend to the committed structure, not to arbitrary context.
+        """
+
+        def __init__(
+            self,
+            gov_vocab: int = STRUCT_VOCAB_SIZE,
+            prose_vocab: int = PROSE_VOCAB_SIZE,
+            d_model: int = 128,
+            nhead: int = 4,
+            num_encoder_layers: int = 2,
+            num_decoder_layers: int = 4,
+            max_struct_len: int = 40,
+            max_prose_len: int = 256,
+            dropout: float = 0.1,
+        ):
+            super().__init__()
+            self.d_model = d_model
+            self.max_prose_len = max_prose_len
+
+            # Encoder (governed structure)
+            self.struct_embedding = nn.Embedding(gov_vocab, d_model, padding_idx=PAD)
+            self.struct_pos = nn.Embedding(max_struct_len, d_model)
+            encoder_layer = nn.TransformerEncoderLayer(
+                d_model=d_model, nhead=nhead, dim_feedforward=d_model * 4,
+                dropout=dropout, batch_first=True,
+            )
+            self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_encoder_layers)
+
+            # Decoder (prose generation)
+            self.prose_embedding = nn.Embedding(prose_vocab, d_model, padding_idx=PROSE_PAD)
+            self.prose_pos = nn.Embedding(max_prose_len, d_model)
+            decoder_layer = nn.TransformerDecoderLayer(
+                d_model=d_model, nhead=nhead, dim_feedforward=d_model * 4,
+                dropout=dropout, batch_first=True,
+            )
+            self.decoder = nn.TransformerDecoder(decoder_layer, num_layers=num_decoder_layers)
+            self.output_proj = nn.Linear(d_model, prose_vocab)
+
+        def forward(
+            self,
+            gov_tokens: torch.Tensor,   # (B, struct_len)
+            prose_tokens: torch.Tensor,    # (B, prose_len)
+        ) -> torch.Tensor:
+            """Forward pass. Returns logits (B, prose_len, prose_vocab)."""
+            B = gov_tokens.size(0)
+
+            # Encode governed structure
+            struct_len = gov_tokens.size(1)
+            struct_pos = torch.arange(struct_len, device=gov_tokens.device).unsqueeze(0).expand(B, -1)
+            struct_h = self.struct_embedding(gov_tokens) + self.struct_pos(struct_pos)
+            struct_pad_mask = (gov_tokens == PAD)
+            memory = self.encoder(struct_h, src_key_padding_mask=struct_pad_mask)
+
+            # Decode prose
+            prose_len = prose_tokens.size(1)
+            prose_pos = torch.arange(prose_len, device=prose_tokens.device).unsqueeze(0).expand(B, -1)
+            prose_h = self.prose_embedding(prose_tokens) + self.prose_pos(prose_pos)
+
+            # Causal mask for autoregressive generation
+            causal_mask = nn.Transformer.generate_square_subsequent_mask(prose_len, device=prose_tokens.device)
+            prose_pad_mask = (prose_tokens == PROSE_PAD)
+
+            decoded = self.decoder(
+                prose_h, memory,
+                tgt_mask=causal_mask,
+                tgt_key_padding_mask=prose_pad_mask,
+                memory_key_padding_mask=struct_pad_mask,
+            )
+            return self.output_proj(decoded)
+
+        def generate(
+            self,
+            gov_tokens: torch.Tensor,  # (B, struct_len)
+            max_len: int = 200,
+            temperature: float = 0.8,
+        ) -> list[str]:
+            """Generate prose from governed structure."""
+            self.eval()
+            B = gov_tokens.size(0)
+            device = gov_tokens.device
+
+            # Encode governed
+            struct_len = gov_tokens.size(1)
+            struct_pos = torch.arange(struct_len, device=device).unsqueeze(0).expand(B, -1)
+            struct_h = self.struct_embedding(gov_tokens) + self.struct_pos(struct_pos)
+            struct_pad_mask = (gov_tokens == PAD)
+            memory = self.encoder(struct_h, src_key_padding_mask=struct_pad_mask)
+
+            # Autoregressive generation
+            generated = torch.full((B, 1), PROSE_BOS, dtype=torch.long, device=device)
+
+            with torch.no_grad():
+                for _ in range(max_len):
+                    prose_len = generated.size(1)
+                    prose_pos = torch.arange(prose_len, device=device).unsqueeze(0).expand(B, -1)
+                    prose_h = self.prose_embedding(generated) + self.prose_pos(prose_pos)
+                    causal_mask = nn.Transformer.generate_square_subsequent_mask(prose_len, device=device)
+
+                    decoded = self.decoder(prose_h, memory, tgt_mask=causal_mask, memory_key_padding_mask=struct_pad_mask)
+                    logits = self.output_proj(decoded[:, -1, :]) / temperature
+                    probs = F.softmax(logits, dim=-1)
+                    next_token = torch.multinomial(probs, 1)
+                    generated = torch.cat([generated, next_token], dim=1)
+
+                    if (next_token == PROSE_EOS).all():
+                        break
+
+            results = []
+            for b in range(B):
+                results.append(decode_prose(generated[b].tolist()))
+            return results