feat(v9): sub-tokenization — [UNK] collisions fixed + namespace probe passes + apiVersion probe added

app.py

@@ -51,22 +51,20 @@ DEFAULT_VOCAB = "model/vocab.json"
 def load_model(checkpoint_path: str, vocab_path: str) -> tuple[YamlBertModel, Vocabulary]:
     _log(f"Loading vocab from {vocab_path}")
     vocab = Vocabulary.load(vocab_path)
-    _log(f"Vocab loaded: {vocab.key_vocab_size} keys, "
-         f"{vocab.value_vocab_size} values, "
-         f"{vocab.atomic_target_vocab_size} atomic targets")
+    _log(f"Vocab loaded: subword={vocab.subword_vocab_size}, "
+         f"atomic targets={vocab.atomic_target_vocab_size}")
 
     _log(f"Reading checkpoint file {checkpoint_path}")
     cp = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
 
-    _log("Building YamlBertModel architecture")
-    # recon_enabled=True is required to load checkpoints from MLM+recon training
-    # (state_dict includes recon_head weights). The recon head exists but is
-    # never invoked at inference time (no subtree_roots_flat passed in forward).
+    _log("Building YamlBertModel architecture (v9: subword embedding)")
+    # recon_enabled=True keeps the checkpoint's recon_head weights loadable.
+    # The recon head exists but is never invoked at inference time
+    # (no subtree_roots_flat passed in forward).
     config = YamlBertConfig(recon_enabled=True)
     emb = YamlBertEmbedding(
         config=config,
-        key_vocab_size=vocab.key_vocab_size,
-        value_vocab_size=vocab.value_vocab_size,
+        subword_vocab_size=vocab.subword_vocab_size,
     )
     model = YamlBertModel(
         config=config,
@@ -504,6 +502,80 @@ spec:
     - containerPort: 80
 """
 
+_DEPLOY_APPS_V1 = """apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: web
+spec:
+  replicas: 3
+  selector:
+    matchLabels:
+      app: web
+  template:
+    metadata:
+      labels:
+        app: web
+    spec:
+      containers:
+      - name: app
+        image: nginx
+        ports:
+        - containerPort: 80
+"""
+
+_DEPLOY_EXT_V1BETA1 = """apiVersion: extensions/v1beta1
+kind: Deployment
+metadata:
+  name: web
+spec:
+  replicas: 3
+  selector:
+    matchLabels:
+      app: web
+  template:
+    metadata:
+      labels:
+        app: web
+    spec:
+      containers:
+      - name: app
+        image: nginx
+        ports:
+        - containerPort: 80
+"""
+
+_REPLICASET_APPS_V1 = """apiVersion: apps/v1
+kind: ReplicaSet
+metadata:
+  name: web
+spec:
+  replicas: 3
+  selector:
+    matchLabels:
+      app: web
+  template:
+    metadata:
+      labels:
+        app: web
+    spec:
+      containers:
+      - name: app
+        image: nginx
+        ports:
+        - containerPort: 80
+"""
+
+_CONFIGMAP_PLAIN = """apiVersion: v1
+kind: ConfigMap
+metadata:
+  name: web
+data:
+  config.yaml: |
+    debug: true
+  app.properties: |
+    key=value
+"""
+
 
 def _verdict_init(cos):
     cd = float(cos[2][3])
@@ -513,10 +585,17 @@ def _verdict_init(cos):
         f"`cos(C, D)` = **{cd:.3f}** (both have init) · "
         f"`max(cos(A,C), cos(B,D))` = **{mx:.3f}** (mixed). "
         + ("Init-pairs cluster tighter than mixed pairs — the model treats "
-           "`initContainers` as a real structural feature."
+           "`initContainers` as a real structural feature, stronger than the "
+           "value-content similarity (shared `image: nginx` etc.)."
            if passed else
-           "Mixed pairs are at least as close as init-pairs — the model "
-           "does not cleanly separate Pods by `initContainers` presence here.")
+           "Mixed pairs are at least as close as init-pairs. This is not a "
+           "regression — it reveals a re-balance: BPE makes `image` values "
+           "compositionally visible to attention, so pods sharing `nginx` "
+           "cluster together regardless of init presence. v8 with atomic "
+           "`[UNK]` values had no choice but to lean on structure; v9 has "
+           "both signals and now weights content more heavily here. "
+           "Whether that's good depends on use case (good for content "
+           "retrieval, less ideal for structure-only similarity).")
     )
     return passed, msg
 
@@ -547,15 +626,48 @@ def _verdict_namespace(cos):
     msg = (
         f"`min(same-ns cos)` = **{same_ns:.3f}** · "
         f"`max(cross-ns cos)` = **{cross_ns:.3f}**. "
-        + ("Same-namespace pairs cluster tighter — the model has somehow "
-           "encoded namespace as a feature (surprising, since namespace is "
-           "a leaf VALUE per the design rationale)."
+        + ("Same-namespace pairs cluster tighter — value content reaches "
+           "`doc_vec` even though the aggregator only sums KEY subtrees. "
+           "BPE makes namespace values compositional (e.g., `prod | uction`), "
+           "and self-attention spreads that signal into neighboring KEY "
+           "hidden states, which then flow into `doc_vec`. This was the "
+           "first failure of `[UNK]`-vocab v8 that v9 fixed."
            if passed else
            "Same-namespace and cross-namespace pairs have indistinguishable "
-           "cosines — the model is value-blind for namespace. This matches "
-           "the design: values are second-class in `doc_vec`. Compare with "
-           "the Service-type probe (passes) — `type` values matter because "
-           "they bring structural KEYS along, while namespace values don't.")
+           "cosines. v8 saw this because both `production` and `staging` "
+           "often hit `[UNK]`. v9 was expected to pass — if it does not, "
+           "investigate.")
+    )
+    return passed, msg
+
+
+def _verdict_apiversion(cos):
+    # A = apps/v1 Deployment
+    # B = extensions/v1beta1 Deployment  (same kind, deprecated apiVersion)
+    # C = apps/v1 ReplicaSet              (different kind, same group, similar structure)
+    # D = v1 ConfigMap                    (unrelated)
+    same_kind = float(cos[0][1])
+    same_group_diff_kind = float(cos[0][2])
+    unrelated = float(cos[0][3])
+    primary = same_kind > same_group_diff_kind
+    secondary = same_group_diff_kind > unrelated
+    passed = primary and secondary
+    msg = (
+        f"`cos(apps/v1 Dep, ext/v1beta1 Dep)` = **{same_kind:.3f}** · "
+        f"`cos(Dep, ReplicaSet)` = **{same_group_diff_kind:.3f}** · "
+        f"`cos(Dep, ConfigMap)` = **{unrelated:.3f}**. "
+        + ("Same-kind-different-apiVersion pairs cluster tightest, then "
+           "same-group-different-kind, then unrelated. The model treats "
+           "apiVersion as a soft label, not a hard discriminator — it "
+           "recognizes `apps/v1` and `extensions/v1beta1` Deployments as "
+           "the same thing despite the apiVersion text differing."
+           if passed else
+           f"Expected ordering broken: same-kind={same_kind:.3f}, "
+           f"same-group={same_group_diff_kind:.3f}, unrelated={unrelated:.3f}. "
+           "If same-kind is NOT the tightest, the model is treating "
+           "apiVersion as a hard discriminator and separating same-kind "
+           "manifests by their api label — possibly an over-correction "
+           "from the eval-probe accuracy.")
     )
     return passed, msg
 
@@ -620,10 +732,14 @@ PRESETS = [
             "If the model encodes `metadata.namespace` as a feature, two "
             "Pods in the same namespace should be closer than two Pods in "
             "different namespaces (controlling for structure). "
-            "_Honest prediction: this should **fail** — namespace is a "
-            "leaf value, and values are second-class in `doc_vec`. The "
-            "failure would confirm the design rationale; a pass would "
-            "refute it._"
+            "_v8 with atomic vocab failed this probe — `production` and "
+            "`staging` often mapped to `[UNK]`, leaving attention with no "
+            "compositional content to work with. v9's byte-level BPE "
+            "decomposes namespace values into subwords, and self-attention "
+            "now spreads value content into surrounding KEY hidden states. "
+            "Those KEYs are what the aggregator pools into `doc_vec` — so "
+            "namespace effectively reaches `doc_vec` through the attention "
+            "channel, even though the aggregator stays KEY-only by design._"
         ),
         "manifests": [
             {"name": "production / web-1", "yaml": _POD_NS_PROD_1},
@@ -633,6 +749,32 @@ PRESETS = [
         ],
         "verdict_fn": _verdict_namespace,
     },
+    {
+        "id": "apiversion",
+        "title": "apiVersion sensitivity (same kind, different apiVersion)",
+        "hypothesis": (
+            "K8s often supports multiple `apiVersion`s for the same kind "
+            "(e.g., `apps/v1` Deployment and the deprecated "
+            "`extensions/v1beta1` Deployment have nearly identical structure). "
+            "If the model understands kind as the structural identity "
+            "and apiVersion as a soft label, two same-kind manifests "
+            "with different apiVersions should still cluster tighter than "
+            "a same-group different-kind manifest (`apps/v1` ReplicaSet), "
+            "which in turn should be tighter than an unrelated kind "
+            "(`v1` ConfigMap). _The eval probes already show apiVersion "
+            "classification at 99.8% — but classification accuracy is "
+            "compatible with either understanding the relationship or "
+            "treating apiVersion as a hard discriminator. This probe "
+            "tests which._"
+        ),
+        "manifests": [
+            {"name": "apps/v1 Deployment",         "yaml": _DEPLOY_APPS_V1},
+            {"name": "extensions/v1beta1 Deploy",  "yaml": _DEPLOY_EXT_V1BETA1},
+            {"name": "apps/v1 ReplicaSet",         "yaml": _REPLICASET_APPS_V1},
+            {"name": "v1 ConfigMap (unrelated)",   "yaml": _CONFIGMAP_PLAIN},
+        ],
+        "verdict_fn": _verdict_apiversion,
+    },
     {
         "id": "cross-kind",
         "title": "Pod vs Deployment wrapping the same Pod",
@@ -698,6 +840,8 @@ def _encode_doc_vec(yaml_text: str) -> torch.Tensor:
             sibling_indices=batch["sibling_indices"],
             batch_info=batch["batch_info"],
             padding_mask=batch["padding_mask"],
+            logical_ids=batch["logical_ids"],
+            n_logical_per_doc=batch["n_logical_per_doc"],
             parent_of_tensor=batch["parent_of_tensor"],
             top_level_key_mask=batch["top_level_key_mask"],
             edges_by_depth=batch["edges_by_depth"],
@@ -722,6 +866,9 @@ def _layout_2d(vecs):
     dist = 1.0 - cos
     np.fill_diagonal(dist, 0.0)
     dist = np.clip(dist, 0.0, None)
+    # sklearn MDS(metric="precomputed") requires strict symmetry; floating-point
+    # matmul can leave ~1e-7 asymmetries. Symmetrize explicitly.
+    dist = (dist + dist.T) / 2
 
     from sklearn.manifold import MDS
     mds = MDS(
@@ -1198,7 +1345,7 @@ Trained with MLM + reconstruction on 276K K8s manifests ·
                         value=EXAMPLE_NGINX,
                     )
                     threshold = gr.Slider(
-                        minimum=0.05, maximum=0.95, value=0.1, step=0.05,
+                        minimum=0.05, maximum=0.95, value=0.7, step=0.05,
                         label="Confidence threshold",
                     )
                     submit = gr.Button("Suggest missing fields", variant="primary")

model/yaml_bert.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:13d86c864f0728e2c92edbbe2622463974201c21a4515c5df8f70131f9e6d1e9
-size 90137468
+oid sha256:5663d31145fad1b0f3d842624f5ed08b24c2bff5bfb42e77682c0877741b2c9c
+size 74032259

yaml_bert/aggregator.py

@@ -1,11 +1,5 @@
-"""Tree aggregator: bottom-up combine of token hidden states into subtree
-vectors + a document vector at the root.
-
-Mean combine (no learnable params). Two execution paths share the same
-output: a per-document reference loop, and a batched scatter-based
-vectorized path activated when the caller passes precomputed tensors.
-Numerical equivalence between the two paths is locked by
-tests/test_aggregator_vectorized.py.
+"""Tree aggregator v9: pool subwords per logical node, then bottom-up
+combine of logical KEY nodes into subtree vectors + a document vector.
 """
 from __future__ import annotations
 
@@ -13,19 +7,59 @@ import torch
 import torch.nn as nn
 
 
+def _pool_subwords(
+    hidden_states: torch.Tensor,
+    logical_ids: torch.Tensor,
+    n_logical_per_doc: torch.Tensor,
+) -> torch.Tensor:
+    """Mean-pool subword hidden states into per-logical-node vectors.
+
+    Args:
+        hidden_states: (B, N_sub, d) per-subword hidden states from the encoder.
+        logical_ids:  (B, N_sub) int tensor; -1 marks padding (ignored).
+        n_logical_per_doc: (B,) number of logical nodes per doc; pooled output
+            shape is (B, max(n_logical_per_doc), d).
+
+    Returns:
+        (B, L_max, d) where L_max = int(n_logical_per_doc.max()).
+    """
+    B, N_sub, d = hidden_states.shape
+    L_max = int(n_logical_per_doc.max().item())
+    out = torch.zeros(B, L_max, d, device=hidden_states.device, dtype=hidden_states.dtype)
+    count = torch.zeros(B, L_max, device=hidden_states.device, dtype=torch.float32)
+
+    valid = logical_ids >= 0  # (B, N_sub)
+    safe_lids = logical_ids.clamp(min=0)  # (B, N_sub)
+
+    # Doc index broadcast over N_sub
+    doc_idx = torch.arange(B, device=hidden_states.device).unsqueeze(1).expand(B, N_sub)
+
+    # Linear (doc, logical) → flat slot
+    flat = doc_idx * L_max + safe_lids  # (B, N_sub)
+    flat_valid = flat[valid]
+    h_valid = hidden_states[valid]
+
+    out_flat = out.view(B * L_max, d)
+    count_flat = count.view(B * L_max)
+    out_flat.index_add_(0, flat_valid, h_valid)
+    count_flat.index_add_(
+        0, flat_valid, torch.ones_like(flat_valid, dtype=torch.float32),
+    )
+
+    pooled = out_flat / count_flat.clamp(min=1.0).unsqueeze(-1).to(out_flat.dtype)
+    return pooled.view(B, L_max, d)
+
+
 class TreeAggregator(nn.Module):
-    """Bottom-up tree aggregation with mean combine.
+    """v9: pool subwords first, then run v8 logical-level aggregator.
 
     Two execution paths:
     - Reference path (default): per-doc Python loop. Used when the batch
-      doesn't provide vectorized precompute tensors. Kept for tests and
-      for backward compatibility.
-    - Vectorized path (preferred during training): batched PyTorch scatter
-      ops, processed depth-by-depth. Activated when caller passes the
-      precomputed tensors as kwargs.
-
-    Both paths produce numerically equivalent output on the same inputs
-    (guaranteed by tests/test_aggregator_vectorized.py).
+      doesn't provide vectorized precompute tensors. Kept for tests.
+    - Vectorized path: batched scatter ops, processed depth-by-depth.
+
+    Both paths produce numerically equivalent output (guaranteed by
+    tests/test_aggregator_vectorized.py).
     """
 
     def __init__(self, d_model: int) -> None:
@@ -37,6 +71,8 @@ class TreeAggregator(nn.Module):
         hidden_states: torch.Tensor,
         batch_info: list[dict],
         *,
+        logical_ids: torch.Tensor,
+        n_logical_per_doc: torch.Tensor,
         parent_of_tensor: torch.Tensor | None = None,
         top_level_key_mask: torch.Tensor | None = None,
         edges_by_depth: dict[int, torch.Tensor] | None = None,
@@ -45,19 +81,21 @@ class TreeAggregator(nn.Module):
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """
         Args:
-            hidden_states: (B, N, d_model) per-position hidden states.
-            batch_info: list of B dicts (legacy path; required even when
-                vectorized — kept for fallback compat).
+            hidden_states: (B, N_sub, d) per-subword hidden states from encoder.
+            logical_ids:   (B, N_sub) per-subword logical-node id (-1 for pad).
+            n_logical_per_doc: (B,) number of logical nodes per doc.
+            batch_info: list of B dicts (legacy path; required for reference path).
             parent_of_tensor / top_level_key_mask / edges_by_depth /
                 parents_by_depth: when ALL provided, use vectorized path.
-            subtree_mask: (B, N) bool tensor. When provided, positions where
-                subtree_mask[b, i]=True are excluded from sums into doc_vec
-                and from contributing to ancestor subtree_vecs. Used for the
-                v8 reconstruction objective's leak-prevention.
+            subtree_mask: (B, L_max) bool; positions excluded from doc_vec and
+                ancestor subtree_vecs (used by v8 reconstruction objective).
 
         Returns:
-            (subtree_vecs, doc_vec)
+            (subtree_vecs, doc_vec) where subtree_vecs is (B, L_max, d) —
+            indexed by LOGICAL position, not subword.
         """
+        pooled = _pool_subwords(hidden_states, logical_ids, n_logical_per_doc)
+
         provided = (
             parent_of_tensor is not None,
             top_level_key_mask is not None,
@@ -75,16 +113,19 @@ class TreeAggregator(nn.Module):
                     f"parents_by_depth={'set' if provided[3] else 'None'}"
                 )
             return self._forward_vectorized(
-                hidden_states,
+                pooled,
                 top_level_key_mask=top_level_key_mask,
                 edges_by_depth=edges_by_depth,
                 parents_by_depth=parents_by_depth,
                 subtree_mask=subtree_mask,
             )
         return self._forward_reference(
-            hidden_states, batch_info, subtree_mask=subtree_mask,
+            pooled, batch_info, subtree_mask=subtree_mask,
         )
 
+    # === _forward_reference and _forward_vectorized are verbatim from v8 ===
+    # They operate on per-position hidden states (now logical positions).
+
     def _forward_reference(
         self,
         hidden_states: torch.Tensor,

yaml_bert/config.py

@@ -31,7 +31,7 @@ class YamlBertConfig:
     lr: float = 1e-4
     batch_size: int = 32
     num_epochs: int = 30
-    max_seq_len: int = 512
+    max_seq_len: int = 768  # was 512 in v8; v9 BPE-expands sequences ~2.3x
 
     # Reconstruction objective (subtree masking + bag-of-keys prediction)
     recon_enabled: bool = False

yaml_bert/dataset.py

@@ -1,4 +1,5 @@
-"""YAML-BERT dataset: children-info precompute for tree aggregator."""
+"""v9 YAML-BERT dataset: BPE-expand each linearizer node into subword
+positions, mask whole logical KEYs, emit per-logical-node atomic labels."""
 from __future__ import annotations
 
 import random
@@ -17,37 +18,11 @@ _LIST_INDEX_RE = re.compile(r"\.\d+$")
 
 
 def _strip_trailing_list_index(path: str) -> str:
-    """Strip a trailing numeric segment from a parent_path.
-
-    E.g., 'spec.containers.0' -> 'spec.containers'.
-    Returns the path unchanged if no trailing numeric segment exists.
-    """
     return _LIST_INDEX_RE.sub("", path)
 
 
 def compute_children_info(nodes: list[YamlNode]) -> dict:
-    """Compute per-position parent/child relationships.
-
-    For each KEY/LIST_KEY position, its children are KEY/LIST_KEY positions
-    whose parent_path equals this key's full_path. (VALUE nodes are leaves —
-    their hidden states are used in the aggregator without being treated as
-    "children" of any KEY for subtree purposes.)
-
-    For KEYs inside list items, parent_path ends in a numeric list index
-    (e.g., 'spec.containers.0'), but the linearizer never emits a synthetic
-    list-item node, so a direct lookup misses. We strip the trailing numeric
-    segment and link to the list-key itself. This flattens all list items
-    into their list parent — per-item grouping is lost, but the aggregator
-    still produces a meaningful doc vector. Phase 1 may add synthetic
-    list-item nodes if per-item grouping matters.
-
-    Returns a dict with:
-        children_of: list[list[int]] — children's positions per node
-        parent_of:   list[int]       — parent position (or -1)
-        key_positions: list[int]     — positions that are KEY/LIST_KEY
-        depth_of:    list[int]       — depth per position
-        full_path_of: list[str]      — full path per position
-    """
+    """Same as v8 — operates on LOGICAL nodes (not subwords)."""
     n = len(nodes)
     full_path_of: list[str] = []
     for node in nodes:
@@ -56,7 +31,6 @@ def compute_children_info(nodes: list[YamlNode]) -> dict:
         else:
             full_path_of.append(node.token)
 
-    # Index KEY positions by their full_path
     key_positions: list[int] = [
         i for i, node in enumerate(nodes)
         if node.node_type in (NodeType.KEY, NodeType.LIST_KEY)
@@ -73,11 +47,8 @@ def compute_children_info(nodes: list[YamlNode]) -> dict:
         parent_path = nodes[p].parent_path
         if not parent_path:
             continue
-        # Try direct lookup first
         parent_pos = path_to_key_pos.get(parent_path)
         if parent_pos is None:
-            # Try with trailing list index stripped
-            # (e.g., "spec.containers.0" -> "spec.containers")
             stripped = _strip_trailing_list_index(parent_path)
             if stripped != parent_path:
                 parent_pos = path_to_key_pos.get(stripped)
@@ -104,7 +75,7 @@ _MASKABLE_TYPES = (NodeType.KEY, NodeType.LIST_KEY)
 
 
 class YamlBertDataset(Dataset):
-    """YAML-BERT dataset: atomic labels at masked positions + children info per doc."""
+    """v9 dataset: subword expansion + whole-key MLM masking + recon."""
 
     def __init__(
         self,
@@ -118,18 +89,12 @@ class YamlBertDataset(Dataset):
         self.max_seq_len = config.max_seq_len
         self.recon_enabled = config.recon_enabled
 
-        # Precompute per-doc children_info AND descendant cache (subtree picker
-        # uses descendants on every __getitem__ call; cache once at init).
         self._cached_children_info: list[dict] = []
         self._cached_descendants: list[dict[int, set[int]] | None] = []
-        if self.recon_enabled and len(documents) > 100:
-            import logging
-            logging.getLogger(__name__).info(
-                "YamlBertDataset: precomputing descendants for %d docs (recon enabled)",
-                len(documents),
-            )
         for doc in documents:
-            ci = compute_children_info(doc[: self.max_seq_len])
+            # Cap LOGICAL nodes here; BPE expansion may still exceed max_seq_len
+            # at the subword level (handled below by truncation).
+            ci = compute_children_info(doc)
             self._cached_children_info.append(ci)
             if self.recon_enabled:
                 desc_cache: dict[int, set[int]] = {}
@@ -145,80 +110,116 @@ class YamlBertDataset(Dataset):
 
     def __getitem__(self, idx: int) -> dict:
         nodes = self.documents[idx]
-        if len(nodes) > self.max_seq_len:
-            nodes = nodes[: self.max_seq_len]
-
-        token_ids: list[int] = []
-        node_types: list[int] = []
-        depths: list[int] = []
-        sibling_indices: list[int] = []
 
-        for node in nodes:
-            if node.node_type in (NodeType.KEY, NodeType.LIST_KEY):
-                token_ids.append(self.vocab.encode_key(node.token))
-            else:
-                token_ids.append(self.vocab.encode_value(node.token))
-            node_types.append(_NODE_TYPE_INDEX[node.node_type])
-            depths.append(min(node.depth, 15))
-            sibling_indices.append(min(node.sibling_index, 31))
+        # Pass 1: BPE-expand each logical node, building per-subword tensors.
+        sub_token_ids: list[int] = []
+        sub_node_types: list[int] = []
+        sub_depths: list[int] = []
+        sub_sibling: list[int] = []
+        sub_logical_ids: list[int] = []
+        per_logical_subword_spans: list[tuple[int, int]] = []
+        # We may need to drop trailing logical nodes if subword expansion
+        # blows the max_seq_len cap.
+        kept_logical: int = 0
+        for logical_idx, node in enumerate(nodes):
+            is_value = node.node_type in (NodeType.VALUE, NodeType.LIST_VALUE)
+            ids = self.vocab.encode_token(node.token, is_value=is_value)
+            if len(sub_token_ids) + len(ids) > self.max_seq_len:
+                break
+            start = len(sub_token_ids)
+            sub_token_ids.extend(ids)
+            sub_node_types.extend([_NODE_TYPE_INDEX[node.node_type]] * len(ids))
+            sub_depths.extend([min(node.depth, 15)] * len(ids))
+            sub_sibling.extend([min(node.sibling_index, 31)] * len(ids))
+            sub_logical_ids.extend([kept_logical] * len(ids))
+            per_logical_subword_spans.append((start, start + len(ids)))
+            kept_logical += 1
+
+        n_logical = kept_logical
+        # Truncate cached children_info to kept logicals
+        ci = self._cached_children_info[idx]
+        kept_set = set(range(n_logical))
+        children_of_t = [
+            [c for c in ci["children_of"][p] if c in kept_set]
+            for p in range(n_logical)
+        ]
+        parent_of_t = [
+            ci["parent_of"][p] if (ci["parent_of"][p] in kept_set or ci["parent_of"][p] == -1) else -1
+            for p in range(n_logical)
+        ]
+        key_positions_t = [p for p in ci["key_positions"] if p < n_logical]
+        depth_of_t = ci["depth_of"][:n_logical]
+        full_path_of_t = ci["full_path_of"][:n_logical]
+        ci_t = {
+            "children_of": children_of_t,
+            "parent_of": parent_of_t,
+            "key_positions": key_positions_t,
+            "depth_of": depth_of_t,
+            "full_path_of": full_path_of_t,
+        }
 
-        atomic_labels: list[int] = [-100] * len(nodes)
-        mask_id: int = self.vocab.special_tokens["[MASK]"]
-        unk_id: int = self.vocab.special_tokens["[UNK]"]
+        # Pass 2: whole-key MLM masking, one decision per LOGICAL KEY.
+        atomic_labels: list[int] = [-100] * n_logical
+        mask_id = self.vocab.mask_id
+        unk_id = self.vocab.unk_id
+        subword_vocab_size = self.vocab.subword_vocab_size
 
         mlm_masked_positions: set[int] = set()
-        for i, node in enumerate(nodes):
-            if node.node_type not in _MASKABLE_TYPES:
-                continue
+        for logical_idx in key_positions_t:
             if random.random() >= self.mask_prob:
                 continue
-            atomic_id = self.vocab.encode_atomic_target(node.token)
+            tok = nodes[logical_idx].token
+            atomic_id = self.vocab.encode_atomic_target(tok)
             if atomic_id == unk_id:
-                continue  # skip [UNK] targets (Lever 1)
-            atomic_labels[i] = atomic_id
-            mlm_masked_positions.add(i)
+                continue  # skip [UNK] targets (Lever 1, carried from v8)
+            atomic_labels[logical_idx] = atomic_id
+            mlm_masked_positions.add(logical_idx)
             r = random.random()
+            start, end = per_logical_subword_spans[logical_idx]
             if r < 0.8:
-                token_ids[i] = mask_id
+                for p in range(start, end):
+                    sub_token_ids[p] = mask_id
             elif r < 0.9:
-                token_ids[i] = random.randint(
-                    len(self.vocab.special_tokens),
-                    len(self.vocab.key_vocab) + len(self.vocab.special_tokens) - 1,
-                )
+                for p in range(start, end):
+                    sub_token_ids[p] = random.randint(4, subword_vocab_size - 1)
 
         result = {
-            "token_ids": torch.tensor(token_ids, dtype=torch.long),
-            "node_types": torch.tensor(node_types, dtype=torch.long),
-            "depths": torch.tensor(depths, dtype=torch.long),
-            "sibling_indices": torch.tensor(sibling_indices, dtype=torch.long),
+            "token_ids": torch.tensor(sub_token_ids, dtype=torch.long),
+            "node_types": torch.tensor(sub_node_types, dtype=torch.long),
+            "depths": torch.tensor(sub_depths, dtype=torch.long),
+            "sibling_indices": torch.tensor(sub_sibling, dtype=torch.long),
+            "logical_ids": torch.tensor(sub_logical_ids, dtype=torch.long),
             "atomic_labels": torch.tensor(atomic_labels, dtype=torch.long),
-            "children_info": self._cached_children_info[idx],
+            "children_info": ci_t,
+            "n_logical": n_logical,
         }
 
         if self.recon_enabled:
             from yaml_bert.subtree_masking import pick_subtrees, bag_of_keys_target
-            ci = self._cached_children_info[idx]
             picked_roots = pick_subtrees(
-                N=len(nodes),
-                key_positions=ci["key_positions"],
-                depth_of=ci["depth_of"],
-                children_of=ci["children_of"],
+                N=n_logical,
+                key_positions=key_positions_t,
+                depth_of=depth_of_t,
+                children_of=children_of_t,
                 mlm_masked_positions=mlm_masked_positions,
                 rng=random,
-                descendants_cache=self._cached_descendants[idx],
+                descendants_cache={
+                    kp: descendants_of(kp, children_of_t)
+                    for kp in key_positions_t
+                    if children_of_t[kp]
+                },
             )
-            # Build subtree_mask + apply [MASK] to subtree positions
-            subtree_mask = torch.zeros(len(nodes), dtype=torch.bool)
+            subtree_mask = torch.zeros(n_logical, dtype=torch.bool)
             picked_positions_all: set[int] = set()
             bag_targets: list[torch.Tensor] = []
-            # position → key string lookup for bag-of-keys building
             position_to_key_str = {
-                i: nodes[i].token
-                for i in range(len(nodes))
-                if nodes[i].node_type in (NodeType.KEY, NodeType.LIST_KEY)
+                i: nodes[i].token for i in key_positions_t
             }
             for root_pos in picked_roots:
-                descs = self._cached_descendants[idx][root_pos]
+                descs = {
+                    d for d in descendants_of(root_pos, children_of_t)
+                    if d < n_logical
+                }
                 picked_positions_all |= descs
                 bag_targets.append(bag_of_keys_target(
                     subtree_positions=descs,
@@ -226,14 +227,16 @@ class YamlBertDataset(Dataset):
                     atomic_vocab=self.vocab.atomic_target_vocab,
                     vocab_size=self.vocab.atomic_target_vocab_size,
                 ))
-            for pos in picked_positions_all:
-                subtree_mask[pos] = True
-                token_ids[pos] = mask_id
-            # Re-encode token_ids tensor since we mutated the list
-            result["token_ids"] = torch.tensor(token_ids, dtype=torch.long)
+            # Apply [MASK] to ALL subwords of each logical position in the picked subtree
+            for lpos in picked_positions_all:
+                subtree_mask[lpos] = True
+                start, end = per_logical_subword_spans[lpos]
+                for p in range(start, end):
+                    sub_token_ids[p] = mask_id
+            result["token_ids"] = torch.tensor(sub_token_ids, dtype=torch.long)
             result["subtree_mask"] = subtree_mask
-            result["subtree_roots"] = picked_roots  # list[int]
-            result["bag_of_keys_targets"] = bag_targets  # list[Tensor (V,)]
+            result["subtree_roots"] = picked_roots
+            result["bag_of_keys_targets"] = bag_targets
             result["_atomic_vocab_size"] = self.vocab.atomic_target_vocab_size
 
         return result
@@ -243,63 +246,79 @@ _COLLATE_NON_TENSOR_KEYS = frozenset({
     "children_info",
     "subtree_roots",
     "bag_of_keys_targets",
-    "subtree_mask",  # bool tensor handled separately
-    "_atomic_vocab_size",  # int, used by collate to size empty fallback tensors
+    "subtree_mask",
+    "_atomic_vocab_size",
+    "n_logical",
 })
 
 
 def collate_fn(batch: list[dict]) -> dict:
-    """Pad tensor fields, keep children_info as a list."""
-    max_len = max(item["token_ids"].size(0) for item in batch)
-    padded: dict[str, list[torch.Tensor]] = {
-        k: [] for k in batch[0].keys() if k not in _COLLATE_NON_TENSOR_KEYS
-    }
+    """Pad subword-level tensors AND logical-level tensors.
+
+    Subword-level (per-position): token_ids, node_types, depths, sibling_indices,
+                                  logical_ids  — padded to max subword length
+    Logical-level (per-logical-node): atomic_labels  — padded to max logical count
+    """
+    max_sub_len = max(item["token_ids"].size(0) for item in batch)
+    max_logical = max(item["n_logical"] for item in batch)
+
+    subword_keys = ("token_ids", "node_types", "depths", "sibling_indices",
+                    "logical_ids")
+    padded_sub: dict[str, list[torch.Tensor]] = {k: [] for k in subword_keys}
+    padded_labels: list[torch.Tensor] = []
     padding_masks: list[torch.Tensor] = []
     batch_info: list[dict] = []
+    n_logical_per_doc: list[int] = []
 
     for item in batch:
-        seq_len = item["token_ids"].size(0)
-        pad_len = max_len - seq_len
-        for key in padded:
-            pad_value = -100 if key == "atomic_labels" else 0
-            if pad_len > 0:
-                padding = torch.full((pad_len,), pad_value, dtype=torch.long)
-                padded[key].append(torch.cat([item[key], padding]))
+        sub_len = item["token_ids"].size(0)
+        pad_sub = max_sub_len - sub_len
+        for k in subword_keys:
+            pad_value = -1 if k == "logical_ids" else 0
+            if pad_sub > 0:
+                padding = torch.full((pad_sub,), pad_value, dtype=torch.long)
+                padded_sub[k].append(torch.cat([item[k], padding]))
             else:
-                padded[key].append(item[key])
+                padded_sub[k].append(item[k])
+
+        labels = item["atomic_labels"]
+        pad_lab = max_logical - labels.size(0)
+        if pad_lab > 0:
+            padded_labels.append(torch.cat([
+                labels, torch.full((pad_lab,), -100, dtype=torch.long),
+            ]))
+        else:
+            padded_labels.append(labels)
+
         mask = torch.cat([
-            torch.zeros(seq_len, dtype=torch.bool),
-            torch.ones(pad_len, dtype=torch.bool),
-        ]) if pad_len > 0 else torch.zeros(seq_len, dtype=torch.bool)
+            torch.zeros(sub_len, dtype=torch.bool),
+            torch.ones(pad_sub, dtype=torch.bool),
+        ]) if pad_sub > 0 else torch.zeros(sub_len, dtype=torch.bool)
         padding_masks.append(mask)
         batch_info.append(item["children_info"])
+        n_logical_per_doc.append(item["n_logical"])
 
-    result = {k: torch.stack(v) for k, v in padded.items()}
+    result = {k: torch.stack(v) for k, v in padded_sub.items()}
+    result["atomic_labels"] = torch.stack(padded_labels)
     result["padding_mask"] = torch.stack(padding_masks)
     result["batch_info"] = batch_info
+    result["n_logical_per_doc"] = torch.tensor(n_logical_per_doc, dtype=torch.long)
 
-    # Vectorized-aggregator precompute. CPU work, runs in DataLoader workers.
+    # parent_of_tensor and top_level_key_mask now operate at LOGICAL level
     B = len(batch)
-    N = max_len
-
-    # parent_of_tensor: (B, N) long. -1 sentinel for no-parent, non-key, or padding.
-    parent_of_tensor = torch.full((B, N), -1, dtype=torch.long)
+    L = max_logical
+    parent_of_tensor = torch.full((B, L), -1, dtype=torch.long)
+    top_level_key_mask = torch.zeros((B, L), dtype=torch.bool)
     for b_idx, info in enumerate(batch_info):
-        parent_of = info["parent_of"]  # list[int] of length n_b
+        parent_of = info["parent_of"]
         n_b = len(parent_of)
         if n_b > 0:
             parent_of_tensor[b_idx, :n_b] = torch.tensor(parent_of, dtype=torch.long)
-
-    # top_level_key_mask: (B, N) bool. True where depth==0 AND position is a KEY.
-    top_level_key_mask = torch.zeros((B, N), dtype=torch.bool)
-    for b_idx, info in enumerate(batch_info):
         depth_of = info["depth_of"]
         depth_zero_kps = [kp for kp in info["key_positions"] if depth_of[kp] == 0]
         if depth_zero_kps:
             top_level_key_mask[b_idx, depth_zero_kps] = True
 
-    # edges_by_depth: dict[depth, (E, 3) long] of [doc_idx, child_pos, parent_pos] across batch.
-    # parents_by_depth: dict[depth, (P, 2) long] of unique [doc_idx, parent_pos] with at-least-one-child.
     edges_by_depth: dict[int, list[tuple[int, int, int]]] = {}
     parents_set_by_depth: dict[int, set[tuple[int, int]]] = {}
     for b_idx, info in enumerate(batch_info):
@@ -328,13 +347,11 @@ def collate_fn(batch: list[dict]) -> dict:
         for d, parents_set in parents_set_by_depth.items()
     }
 
-    # Subtree-mask batching (only present when recon is enabled per-item).
     if "subtree_mask" in batch[0]:
-        # subtree_mask: (B, N) bool, pad with False
         subtree_masks: list[torch.Tensor] = []
         for item in batch:
             sm = item["subtree_mask"]
-            pad_len = max_len - sm.size(0)
+            pad_len = max_logical - sm.size(0)
             if pad_len > 0:
                 subtree_masks.append(torch.cat([
                     sm, torch.zeros(pad_len, dtype=torch.bool),
@@ -343,8 +360,6 @@ def collate_fn(batch: list[dict]) -> dict:
                 subtree_masks.append(sm)
         result["subtree_mask"] = torch.stack(subtree_masks)
 
-        # subtree_roots_flat: (M, 2) of [batch_idx, root_pos]; M = total roots
-        # bag_of_keys_targets_flat: (M, V_atomic)
         flat_roots: list[tuple[int, int]] = []
         flat_targets: list[torch.Tensor] = []
         for b_idx, item in enumerate(batch):
@@ -359,12 +374,7 @@ def collate_fn(batch: list[dict]) -> dict:
             )
             result["bag_of_keys_targets_flat"] = torch.stack(flat_targets)
         else:
-            # Even with recon enabled, all docs in this batch may have empty
-            # picks (small docs, no candidates). Emit empty tensors so the
-            # consumer can detect "no subtrees this batch."
-            result["subtree_roots_flat"] = torch.zeros(
-                (0, 2), dtype=torch.long,
-            )
+            result["subtree_roots_flat"] = torch.zeros((0, 2), dtype=torch.long)
             v = batch[0].get("_atomic_vocab_size", 0)
             result["bag_of_keys_targets_flat"] = torch.zeros(
                 (0, v), dtype=torch.float,

yaml_bert/embedding.py

@@ -7,28 +7,25 @@ from yaml_bert.config import TreePosVariant, YamlBertConfig
 
 
 class YamlBertEmbedding(nn.Module):
-    """Embedding layer with tree positional encoding for YAML-BERT.
+    """v9 embedding layer: single subword table + tree positional encoding.
 
     Produces input vectors by summing:
-    - Token embedding (key_embedding or value_embedding, routed by node_type)
+    - Subword embedding (looked up by token_id; same table for KEY and VALUE
+      positions — what they ARE is signalled separately via node_type_emb)
     - Tree positional encoding (composition depends on config.tree_pos_variant)
-
-    Kind and parent awareness come from the hybrid prediction targets, not input encoding.
     """
 
     def __init__(
         self,
         config: YamlBertConfig,
-        key_vocab_size: int,
-        value_vocab_size: int,
+        subword_vocab_size: int,
     ) -> None:
         super().__init__()
         d: int = config.d_model
         variant: TreePosVariant = config.tree_pos_variant
         self.variant: TreePosVariant = variant
 
-        self.key_embedding: nn.Embedding = nn.Embedding(key_vocab_size, d)
-        self.value_embedding: nn.Embedding = nn.Embedding(value_vocab_size, d)
+        self.subword_embedding: nn.Embedding = nn.Embedding(subword_vocab_size, d)
         self.node_type_embedding: nn.Embedding = nn.Embedding(4, d)
 
         use_depth: bool = variant in (TreePosVariant.FULL, TreePosVariant.NO_SIBLING)
@@ -54,14 +51,9 @@ class YamlBertEmbedding(nn.Module):
         depths: torch.Tensor,
         sibling_indices: torch.Tensor,
     ) -> torch.Tensor:
-        is_key: torch.Tensor = (node_types == 0) | (node_types == 2)
-        key_vocab_size: int = self.key_embedding.num_embeddings
-        val_vocab_size: int = self.value_embedding.num_embeddings
-        key_emb: torch.Tensor = self.key_embedding(token_ids.clamp(0, key_vocab_size - 1))
-        val_emb: torch.Tensor = self.value_embedding(token_ids.clamp(0, val_vocab_size - 1))
-        token_emb: torch.Tensor = torch.where(is_key.unsqueeze(-1), key_emb, val_emb)
+        token_emb = self.subword_embedding(token_ids)
 
-        tree_pos: torch.Tensor = self.node_type_embedding(node_types)
+        tree_pos = self.node_type_embedding(node_types)
         if self.depth_embedding is not None:
             tree_pos = tree_pos + self.depth_embedding(depths)
         if self.sibling_embedding is not None:
@@ -69,7 +61,7 @@ class YamlBertEmbedding(nn.Module):
         if self.pos_embedding is not None:
             seq_len: int = token_ids.size(1)
             max_pos: int = self.pos_embedding.num_embeddings
-            positions: torch.Tensor = (
+            positions = (
                 torch.arange(seq_len, device=token_ids.device)
                 .clamp(max=max_pos - 1)
                 .unsqueeze(0)

yaml_bert/model.py

@@ -84,6 +84,8 @@ class YamlBertModel(nn.Module):
         batch_info: list[dict],
         padding_mask: torch.Tensor | None = None,
         *,
+        logical_ids: torch.Tensor,
+        n_logical_per_doc: torch.Tensor,
         parent_of_tensor: torch.Tensor | None = None,
         top_level_key_mask: torch.Tensor | None = None,
         edges_by_depth: dict[int, torch.Tensor] | None = None,
@@ -93,14 +95,21 @@ class YamlBertModel(nn.Module):
     ) -> tuple:
         """Returns (logits, doc_vec) or (logits, doc_vec, recon_logits).
 
-        recon_logits only returned when subtree_roots_flat is provided AND
-        has at least one row."""
+        v9: token_ids/node_types/depths/sibling_indices/logical_ids/padding_mask
+        are SUBWORD-level (B, N_sub). atomic_labels and the aggregator output
+        are LOGICAL-level (B, L_max). The Token Head consumes per-logical-node
+        pooled hidden state.
+        """
+        from yaml_bert.aggregator import _pool_subwords
+
         x = self.embedding(token_ids, node_types, depths, sibling_indices)
         x = self.encoder(x, src_key_padding_mask=padding_mask)
 
-        # Aggregator: forwards through to its own vectorized/reference dispatch.
+        # Aggregator pools internally; (subtree_vecs, doc_vec) are logical-level.
         subtree_vecs, doc_vec = self.aggregator(
             x, batch_info,
+            logical_ids=logical_ids,
+            n_logical_per_doc=n_logical_per_doc,
             parent_of_tensor=parent_of_tensor,
             top_level_key_mask=top_level_key_mask,
             edges_by_depth=edges_by_depth,
@@ -108,51 +117,60 @@ class YamlBertModel(nn.Module):
             subtree_mask=subtree_mask,
         )
 
-        b, n, d = x.shape
+        # Re-pool subword hiddens to logical level for the Token Head input.
+        # (One extra index_add call; the plan flags this as a follow-up.)
+        h_logical = _pool_subwords(x, logical_ids, n_logical_per_doc)
+        b, L_max, d = h_logical.shape
 
         if parent_of_tensor is not None:
-            # Vectorized s_parent. parent_of_tensor being set implies all four
-            # precompute kwargs were provided (aggregator enforces all-or-none).
-            safe_parent = parent_of_tensor.clamp(min=0)  # (B, N)
+            safe_parent = parent_of_tensor.clamp(min=0)
             s_parent = torch.gather(
                 subtree_vecs, dim=1,
                 index=safe_parent.unsqueeze(-1).expand(-1, -1, d),
-            )  # (B, N, d)
-            no_parent_mask = (parent_of_tensor == -1).unsqueeze(-1)  # (B, N, 1)
+            )
+            no_parent_mask = (parent_of_tensor == -1).unsqueeze(-1)
             s_parent = torch.where(
                 no_parent_mask, doc_vec.unsqueeze(1), s_parent,
             )
         else:
-            # Reference path: per-doc Python loop (kept for tests / fallback).
-            s_parent = torch.zeros_like(x)
+            s_parent = torch.zeros_like(h_logical)
             for doc_idx in range(b):
                 parent_of = batch_info[doc_idx]["parent_of"]
-                for i in range(min(n, len(parent_of))):
+                for i in range(min(L_max, len(parent_of))):
                     p = parent_of[i]
                     if p >= 0:
                         s_parent[doc_idx, i] = subtree_vecs[doc_idx, p]
                     else:
                         s_parent[doc_idx, i] = doc_vec[doc_idx]
 
-        doc_vec_broadcast = doc_vec.unsqueeze(1).expand(b, n, d)
-        head_input = torch.cat([x, doc_vec_broadcast, s_parent], dim=-1)
-        logits = self.token_head(head_input)
+        doc_vec_broadcast = doc_vec.unsqueeze(1).expand(b, L_max, d)
+        head_input = torch.cat([h_logical, doc_vec_broadcast, s_parent], dim=-1)
+        logits = self.token_head(head_input)  # (B, L_max, atomic_vocab_size)
 
-        # Reconstruction path: only if caller provided subtree roots
         if subtree_roots_flat is not None and subtree_roots_flat.size(0) > 0:
-            # subtree_roots_flat: (M, 2) of [batch_idx, root_pos]
-            batch_idx_per_root = subtree_roots_flat[:, 0]   # (M,)
-            root_pos_per_root = subtree_roots_flat[:, 1]    # (M,)
-
-            doc_vec_per_root = doc_vec[batch_idx_per_root]  # (M, d_model)
-
-            # Build pos_emb_per_root from the same depth/sibling embedding params
-            # already used in the embedding layer — no new parameters introduced.
-            root_depths = depths[batch_idx_per_root, root_pos_per_root]     # (M,)
-            root_siblings = sibling_indices[batch_idx_per_root, root_pos_per_root]  # (M,)
-            depth_e = self.embedding.depth_embedding(root_depths)            # (M, d_model)
-            sibling_e = self.embedding.sibling_embedding(root_siblings)      # (M, d_model)
-            pos_emb_per_root = torch.cat([depth_e, sibling_e], dim=-1)      # (M, 2*d_model)
+            batch_idx_per_root = subtree_roots_flat[:, 0]
+            root_pos_per_root = subtree_roots_flat[:, 1]
+            doc_vec_per_root = doc_vec[batch_idx_per_root]
+            # Root positions live in LOGICAL coords. Look up depth from batch_info;
+            # sibling: find any subword with logical_id == root_pos, use its sibling.
+            root_depths_list = [
+                batch_info[bi]["depth_of"][rp]
+                for bi, rp in zip(batch_idx_per_root.tolist(), root_pos_per_root.tolist())
+            ]
+            root_siblings_list = []
+            for bi, rp in zip(batch_idx_per_root.tolist(), root_pos_per_root.tolist()):
+                positions = (logical_ids[bi] == rp).nonzero(as_tuple=True)[0]
+                if len(positions) > 0:
+                    root_siblings_list.append(int(sibling_indices[bi, positions[0]].item()))
+                else:
+                    # Shouldn't happen: a recon root must have at least one subword.
+                    root_siblings_list.append(0)
+            root_depths = torch.tensor(root_depths_list, device=depths.device, dtype=torch.long)
+            root_siblings = torch.tensor(root_siblings_list, device=depths.device, dtype=torch.long)
+
+            depth_e = self.embedding.depth_embedding(root_depths)
+            sibling_e = self.embedding.sibling_embedding(root_siblings)
+            pos_emb_per_root = torch.cat([depth_e, sibling_e], dim=-1)
 
             recon_logits = self.recon_head(doc_vec_per_root, pos_emb_per_root)
             return logits, doc_vec, recon_logits

yaml_bert/suggest.py

@@ -118,11 +118,14 @@ def suggest_missing_fields(
         return [], {}
     annotator.annotate(nodes)
 
-    mask_id: int = vocab.special_tokens["[MASK]"]
+    mask_id: int = vocab.mask_id
 
     # Build atomic reverse map for decoding
     id_to_atomic: dict[int, str] = {v: k for k, v in vocab.atomic_target_vocab.items()}
-    id_to_special: dict[int, str] = {v: k for k, v in vocab.special_tokens.items()}
+    id_to_special: dict[int, str] = {
+        vocab.pad_id: "[PAD]", vocab.unk_id: "[UNK]",
+        vocab.mask_id: "[MASK]", vocab.long_value_id: "[LONG_VALUE]",
+    }
 
     # Group key nodes by parent_path
     keys_by_parent: dict[str, set[str]] = {}
@@ -301,9 +304,12 @@ def _run_probe_v8(
     ds = YamlBertDataset([fake_nodes], vocab, infer_config)
     item = ds[0]
 
-    # Apply [MASK] at insert_pos (dataset has mask_prob=0.0, so token is intact)
+    # v9 whole-word masking: insert_pos is a LOGICAL position (index into
+    # fake_nodes). Mask ALL subword positions whose logical_id == insert_pos.
     item["token_ids"] = item["token_ids"].clone()
-    item["token_ids"][insert_pos] = mask_id
+    sub_positions = (item["logical_ids"] == insert_pos).nonzero(as_tuple=True)[0]
+    for p in sub_positions:
+        item["token_ids"][p] = mask_id
 
     batch = collate_fn([item])
 
@@ -315,14 +321,18 @@ def _run_probe_v8(
             sibling_indices=batch["sibling_indices"],
             batch_info=batch["batch_info"],
             padding_mask=batch["padding_mask"],
+            logical_ids=batch["logical_ids"],
+            n_logical_per_doc=batch["n_logical_per_doc"],
             parent_of_tensor=batch["parent_of_tensor"],
             top_level_key_mask=batch["top_level_key_mask"],
             edges_by_depth=batch["edges_by_depth"],
             parents_by_depth=batch["parents_by_depth"],
         )
 
-    # YamlBertModel returns (logits, doc_vec) or (logits, doc_vec, recon_logits)
-    logits = out[0]  # (1, N, V_atomic)
+    # v9 YamlBertModel returns (logits, doc_vec) where logits is
+    # (B, L_max, V_atomic) — indexed by LOGICAL position. insert_pos is
+    # already a logical index.
+    logits = out[0]
     probs = F.softmax(logits[0, insert_pos], dim=-1)
     topk = probs.topk(top_k + 5)
 

yaml_bert/tokenizer.py

@@ -0,0 +1,70 @@
+"""Subword tokenizer wrapper for YAML-BERT v9.
+
+Wraps the HF `tokenizers` library so the rest of yaml_bert depends on
+this stable interface, not on HF internals.
+
+Vocabulary semantics:
+  - Special tokens reserved at training time: [PAD], [UNK], [MASK], [LONG_VALUE]
+  - Otherwise byte-level BPE; any string can be encoded.
+
+Long-value rule (values only; keys never get this treatment):
+  - value length >= LONG_VALUE_THRESHOLD chars → single [LONG_VALUE] token
+  - MAX_CHARS_VALUE < value length < LONG_VALUE_THRESHOLD → truncate to MAX_CHARS_VALUE chars, then BPE
+  - shorter → BPE in full
+"""
+from __future__ import annotations
+
+from tokenizers import Tokenizer
+
+PAD_TOKEN = "[PAD]"
+UNK_TOKEN = "[UNK]"
+MASK_TOKEN = "[MASK]"
+LONG_VALUE_TOKEN = "[LONG_VALUE]"
+
+MAX_CHARS_VALUE = 64
+LONG_VALUE_THRESHOLD = 256
+
+
+class SubwordTokenizer:
+    """Wraps an HF Tokenizer with YAML-BERT-specific value-length rules."""
+
+    def __init__(self, hf_tokenizer: Tokenizer) -> None:
+        self._tok = hf_tokenizer
+        self.pad_id = hf_tokenizer.token_to_id(PAD_TOKEN)
+        self.unk_id = hf_tokenizer.token_to_id(UNK_TOKEN)
+        self.mask_id = hf_tokenizer.token_to_id(MASK_TOKEN)
+        self.long_value_id = hf_tokenizer.token_to_id(LONG_VALUE_TOKEN)
+        for name, val in (
+            (PAD_TOKEN, self.pad_id), (UNK_TOKEN, self.unk_id),
+            (MASK_TOKEN, self.mask_id), (LONG_VALUE_TOKEN, self.long_value_id),
+        ):
+            if val is None:
+                raise ValueError(
+                    f"SubwordTokenizer: required special token {name!r} not "
+                    f"in tokenizer vocab — was the tokenizer trained with "
+                    f"this special token reserved?"
+                )
+
+    @classmethod
+    def load(cls, path: str) -> "SubwordTokenizer":
+        return cls(Tokenizer.from_file(path))
+
+    @property
+    def vocab_size(self) -> int:
+        return self._tok.get_vocab_size()
+
+    def encode_token(self, token: str, *, is_value: bool) -> list[int]:
+        """Encode one linearizer-node token to a list of subword ids.
+
+        See module docstring for the value-length rule. Keys are always
+        BPE-encoded in full regardless of length.
+        """
+        if is_value:
+            if len(token) >= LONG_VALUE_THRESHOLD:
+                return [self.long_value_id]
+            if len(token) > MAX_CHARS_VALUE:
+                token = token[:MAX_CHARS_VALUE]
+        return self._tok.encode(token).ids
+
+    def decode(self, ids: list[int]) -> str:
+        return self._tok.decode(ids)

yaml_bert/vocab.py

@@ -1,308 +1,130 @@
+"""v9 Vocabulary: subword tokenizer + atomic_target_vocab.
+
+The v8 key_vocab + value_vocab were merged into a single subword vocabulary
+exposed via SubwordTokenizer. The Token Head still predicts over an atomic
+target vocab built from frequent KEY tokens in the training corpus.
+"""
 from __future__ import annotations
 
 import json
-from yaml_bert.types import NodeType, YamlNode
-
-
-SPECIAL_TOKENS = ["[PAD]", "[UNK]", "[MASK]"]
+from yaml_bert.tokenizer import SubwordTokenizer
 
 # Canonical casing for known Kubernetes kinds.
-# Maps lowercase → canonical form. Populated by VocabBuilder.build().
+# Maps lowercase → canonical form. Populated by VocabBuilder.build_atomic_target_vocab.
 _KIND_CANONICAL: dict[str, str] = {}
 
 
 def normalize_kind(kind: str) -> str:
-    """Normalize kind value to canonical casing.
-
-    'configmap' → 'ConfigMap', 'pod' → 'Pod', etc.
-    Unknown kinds are returned as-is.
-    """
     return _KIND_CANONICAL.get(kind.lower(), kind)
 
 
 class Vocabulary:
+    """Holds a subword tokenizer + an atomic target vocab.
+
+    `atomic_target_vocab` maps whole KEY strings (e.g. "containers",
+    "restartPolicy") to integer class ids for the Token Head's output.
+    The 4 special tokens (pad/unk/mask/long_value) get the first 4 ids.
+    """
+
     def __init__(
         self,
-        key_vocab: dict[str, int],
-        value_vocab: dict[str, int],
-        special_tokens: dict[str, int],
-        atomic_target_vocab: dict[str, int] | None = None,
+        tokenizer: SubwordTokenizer,
+        atomic_target_vocab: dict[str, int],
+        tokenizer_path: str | None = None,
     ) -> None:
-        self.key_vocab = key_vocab
-        self.value_vocab = value_vocab
-        self.special_tokens = special_tokens
-        self.atomic_target_vocab = atomic_target_vocab or {}
-        self._id_to_key = {v: k for k, v in key_vocab.items()}
-        self._id_to_value = {v: k for k, v in value_vocab.items()}
-        self._id_to_special = {v: k for k, v in special_tokens.items()}
+        self.tokenizer = tokenizer
+        self.atomic_target_vocab = atomic_target_vocab
+        self.tokenizer_path = tokenizer_path
+        # Convenience: expose the 4 special-token ids at the vocab level
+        self.pad_id = tokenizer.pad_id
+        self.unk_id = tokenizer.unk_id
+        self.mask_id = tokenizer.mask_id
+        self.long_value_id = tokenizer.long_value_id
 
-    def encode_key(self, token: str) -> int:
-        return self.key_vocab.get(token, self.special_tokens["[UNK]"])
-
-    def encode_value(self, token: str) -> int:
-        return self.value_vocab.get(token, self.special_tokens["[UNK]"])
-
-    def decode_key(self, id: int) -> str:
-        if id in self._id_to_special:
-            return self._id_to_special[id]
-        return self._id_to_key.get(id, "[UNK]")
-
-    def decode_value(self, id: int) -> str:
-        if id in self._id_to_special:
-            return self._id_to_special[id]
-        return self._id_to_value.get(id, "[UNK]")
-
-    def encode_atomic_target(self, target: str) -> int:
-        """Encode a single key token as its atomic target id.
+    @classmethod
+    def from_tokenizer_path(
+        cls,
+        tokenizer_path: str,
+        atomic_target_vocab: dict[str, int],
+    ) -> "Vocabulary":
+        return cls(
+            tokenizer=SubwordTokenizer.load(tokenizer_path),
+            atomic_target_vocab=atomic_target_vocab,
+            tokenizer_path=tokenizer_path,
+        )
 
-        Returns the [UNK] id if `target` is not in the atomic target vocab.
-        The Token Head predicts single keys drawn from the same token universe
-        as `key_vocab`.
-        """
-        return self.atomic_target_vocab.get(target, self.special_tokens["[UNK]"])
+    def encode_token(self, token: str, *, is_value: bool) -> list[int]:
+        return self.tokenizer.encode_token(token, is_value=is_value)
 
-    @property
-    def key_vocab_size(self) -> int:
-        return len(self.key_vocab) + len(self.special_tokens)
+    def encode_atomic_target(self, key: str) -> int:
+        return self.atomic_target_vocab.get(key, self.unk_id)
 
     @property
-    def value_vocab_size(self) -> int:
-        return len(self.value_vocab) + len(self.special_tokens)
+    def subword_vocab_size(self) -> int:
+        return self.tokenizer.vocab_size
 
     @property
     def atomic_target_vocab_size(self) -> int:
-        return len(self.atomic_target_vocab) + len(self.special_tokens)
+        # +4 to reserve ids 0-3 for special tokens (consistent with v8 layout)
+        return len(self.atomic_target_vocab) + 4
 
     def save(self, path: str) -> None:
-        data = {
-            "key_vocab": self.key_vocab,
-            "value_vocab": self.value_vocab,
-            "special_tokens": self.special_tokens,
-            "atomic_target_vocab": self.atomic_target_vocab,
-        }
+        if self.tokenizer_path is None:
+            raise ValueError(
+                "Vocabulary.save requires tokenizer_path to be set "
+                "(use Vocabulary.from_tokenizer_path)."
+            )
         with open(path, "w") as f:
-            json.dump(data, f, indent=2)
+            json.dump({
+                "tokenizer_path": self.tokenizer_path,
+                "atomic_target_vocab": self.atomic_target_vocab,
+            }, f, indent=2)
 
     @classmethod
-    def load(cls, path: str) -> Vocabulary:
+    def load(cls, path: str) -> "Vocabulary":
         with open(path) as f:
             data = json.load(f)
-        return cls(
-            key_vocab=data["key_vocab"],
-            value_vocab=data["value_vocab"],
-            special_tokens=data["special_tokens"],
-            atomic_target_vocab=data.get("atomic_target_vocab", {}),
+        return cls.from_tokenizer_path(
+            tokenizer_path=data["tokenizer_path"],
+            atomic_target_vocab=data["atomic_target_vocab"],
         )
 
 
 class VocabBuilder:
-    def build(
-        self,
-        nodes: list[YamlNode],
-        min_freq: int = 1,
-        *,
-        key_min_freq: int | None = None,
-        value_min_freq: int | None = None,
-    ) -> Vocabulary:
-        """Build vocab with optional per-category min_freq thresholds.
-
-        Per-category thresholds default to the global `min_freq` for backward
-        compatibility.
-        """
-        key_min_freq = key_min_freq if key_min_freq is not None else min_freq
-        value_min_freq = value_min_freq if value_min_freq is not None else min_freq
-        key_counts: dict[str, int] = {}
-        value_counts: dict[str, int] = {}
-        kind_set: set[str] = set()
-
-        # First pass: collect all kind values to build canonical casing map.
-        # The casing map (_KIND_CANONICAL) is needed by normalize_kind() which
-        # is used downstream by _extract_kind in the types module.
-        raw_kinds: dict[str, dict[str, int]] = {}  # lowercase → {variant: count}
-        prev_was_kind_key = False
-        for node in nodes:
-            if node.node_type in (NodeType.KEY, NodeType.LIST_KEY):
-                prev_was_kind_key = (node.token == "kind" and node.depth == 0)
-            elif node.node_type in (NodeType.VALUE, NodeType.LIST_VALUE):
-                if prev_was_kind_key:
-                    lower = node.token.lower()
-                    raw_kinds.setdefault(lower, {})
-                    raw_kinds[lower][node.token] = raw_kinds[lower].get(node.token, 0) + 1
-                prev_was_kind_key = False
-
-        # Build canonical map: most frequent casing wins
-        _KIND_CANONICAL.clear()
-        for lower, variants in raw_kinds.items():
-            canonical = max(variants, key=variants.get)
-            _KIND_CANONICAL[lower] = canonical
-
-        # Second pass: count tokens with normalized kinds
-        prev_was_kind_key = False
-        for node in nodes:
-            if node.node_type in (NodeType.KEY, NodeType.LIST_KEY):
-                key_counts[node.token] = key_counts.get(node.token, 0) + 1
-                prev_was_kind_key = (node.token == "kind" and node.depth == 0)
-            elif node.node_type in (NodeType.VALUE, NodeType.LIST_VALUE):
-                value_counts[node.token] = value_counts.get(node.token, 0) + 1
-                if prev_was_kind_key:
-                    kind_set.add(normalize_kind(node.token))
-                prev_was_kind_key = False
-
-        # Atomic target vocab: single key tokens. Uses the same filter
-        # threshold as keys since the entries ARE keys.
-        atomic_target_set: set[str] = {
-            t for t, c in key_counts.items() if c >= key_min_freq
-        }
-
-        return self.build_from_counts(
-            key_counts, value_counts, key_min_freq, kind_set,
-            value_min_freq=value_min_freq,
-            atomic_target_set=atomic_target_set,
-        )
+    """Builds the atomic_target_vocab from a corpus.
 
-    @staticmethod
-    def build_from_counts(
-        key_counts: dict[str, int],
-        value_counts: dict[str, int],
-        min_freq: int = 1,
-        kind_set: set[str] | None = None,
-        *,
-        value_min_freq: int | None = None,
-        atomic_target_set: set[str] | None = None,
-    ) -> Vocabulary:
-        """Build vocabulary from pre-computed token counts.
-
-        `min_freq` filters keys (and values if value_min_freq is None).
-        `value_min_freq` overrides for values only.
-        """
-        if value_min_freq is None:
-            value_min_freq = min_freq
-        # Derive atomic_target_set from key_counts when not supplied. The atomic
-        # set IS-A subset of keys by definition, so this is self-consistent and
-        # prevents callers (e.g. build_from_huggingface) from silently producing
-        # an empty atomic_target_vocab — which would break training (output head
-        # would collapse to ~3 classes and every target would map to [UNK]).
-        if atomic_target_set is None:
-            atomic_target_set = {t for t, c in key_counts.items() if c >= min_freq}
-        special_tokens = {tok: i for i, tok in enumerate(SPECIAL_TOKENS)}
-        offset = len(special_tokens)
-
-        key_vocab = {
-            token: i + offset
-            for i, token in enumerate(
-                sorted(t for t, c in key_counts.items() if c >= min_freq)
-            )
-        }
-
-        # Build value vocab: min_freq filtered, but always include kind values
-        value_tokens = set(t for t, c in value_counts.items() if c >= value_min_freq)
-        if kind_set:
-            value_tokens.update(kind_set)  # kinds always in value vocab
-        value_vocab = {
-            token: i + offset
-            for i, token in enumerate(sorted(value_tokens))
-        }
-
-        atomic_target_vocab = {
-            target: i + offset
-            for i, target in enumerate(sorted(atomic_target_set or []))
-        }
-
-        return Vocabulary(
-            key_vocab, value_vocab, special_tokens,
-            atomic_target_vocab=atomic_target_vocab,
-        )
+    The subword tokenizer is built separately by scripts/train_unified_tokenizer.py.
+    """
 
     @staticmethod
-    def save_counts(
-        key_counts: dict[str, int],
-        value_counts: dict[str, int],
-        path: str,
-    ) -> None:
-        """Save raw token counts to a JSON file for reuse."""
-        with open(path, "w") as f:
-            json.dump({"key_counts": key_counts, "value_counts": value_counts}, f)
-        print(f"Token counts saved: {path} ({len(key_counts)} keys, {len(value_counts)} values)")
+    def build_atomic_target_vocab(
+        nodes_per_doc: list[list],
+        min_freq: int,
+    ) -> dict[str, int]:
+        """Scan KEY tokens across docs, return {token: id} for keys appearing >= min_freq times.
 
-    @staticmethod
-    def load_counts(path: str) -> tuple[dict[str, int], dict[str, int]]:
-        """Load raw token counts from a JSON file."""
-        with open(path) as f:
-            data = json.load(f)
-        key_counts: dict[str, int] = data["key_counts"]
-        value_counts: dict[str, int] = data["value_counts"]
-        print(f"Token counts loaded: {path} ({len(key_counts)} keys, {len(value_counts)} values)")
-        return key_counts, value_counts
-
-    def build_from_huggingface(
-        self,
-        dataset_name: str,
-        linearizer: "YamlLinearizer",
-        annotator: "DomainAnnotator",
-        max_docs: int | None = None,
-        min_freq: int = 1,
-        counts_path: str | None = None,
-    ) -> Vocabulary:
-        """Build vocabulary by scanning a HuggingFace dataset.
-
-        Args:
-            dataset_name: HuggingFace dataset ID
-            linearizer: YamlLinearizer instance
-            annotator: DomainAnnotator instance
-            max_docs: Scan at most this many docs for vocab (None = all)
-            min_freq: Minimum frequency to include a token
-            counts_path: If set, save raw counts here (and load if file exists)
+        IDs start at 4 (0-3 are reserved for special tokens).
         """
-        import os
-
-        # Reuse saved counts if available
-        if counts_path and os.path.exists(counts_path):
-            key_counts, value_counts = self.load_counts(counts_path)
-            return self.build_from_counts(key_counts, value_counts, min_freq, None)
-
-        from datasets import load_dataset
-        from yaml_bert.types import YamlNode
-
-        print(f"Building vocabulary from: {dataset_name}")
-        ds = load_dataset(dataset_name, split="train")
-
-        total: int = len(ds) if max_docs is None else min(max_docs, len(ds))
-        print(f"Scanning {total:,} / {len(ds):,} documents for vocabulary...")
-
-        all_nodes: list[YamlNode] = []
-        skipped: int = 0
-        for i in range(total):
-            try:
-                nodes = linearizer.linearize(ds[i]["content"])
-            except Exception:
-                skipped += 1
-                continue
-            if nodes:
-                annotator.annotate(nodes)
-                all_nodes.extend(nodes)
-
-            if (i + 1) % 10000 == 0:
-                print(f"  {i + 1:,} / {total:,} scanned ({len(all_nodes):,} nodes)")
-
-        print(f"Scanned {total - skipped:,} docs, {len(all_nodes):,} nodes ({skipped} skipped)")
-
-        # Compute counts
-        key_counts: dict[str, int] = {}
-        value_counts: dict[str, int] = {}
-        kind_set: set[str] = set()
-        prev_was_kind_key: bool = False
-        for node in all_nodes:
-            if node.node_type in (NodeType.KEY, NodeType.LIST_KEY):
-                key_counts[node.token] = key_counts.get(node.token, 0) + 1
-                prev_was_kind_key = (node.token == "kind" and node.depth == 0)
-            elif node.node_type in (NodeType.VALUE, NodeType.LIST_VALUE):
-                value_counts[node.token] = value_counts.get(node.token, 0) + 1
-                if prev_was_kind_key:
-                    kind_set.add(node.token)
-                prev_was_kind_key = False
-
-        # Save counts for reuse
-        if counts_path:
-            self.save_counts(key_counts, value_counts, counts_path)
-
-        return self.build_from_counts(key_counts, value_counts, min_freq, kind_set)
+        from yaml_bert.types import NodeType
+        counts: dict[str, int] = {}
+        for nodes in nodes_per_doc:
+            for n in nodes:
+                if n.node_type in (NodeType.KEY, NodeType.LIST_KEY):
+                    counts[n.token] = counts.get(n.token, 0) + 1
+
+        # Also populate kind-canonical map as a side effect (used by suggest.py)
+        prev_kind_key = False
+        for nodes in nodes_per_doc:
+            for n in nodes:
+                if n.node_type in (NodeType.KEY, NodeType.LIST_KEY):
+                    prev_kind_key = (n.token == "kind" and n.depth == 0)
+                elif n.node_type in (NodeType.VALUE, NodeType.LIST_VALUE):
+                    if prev_kind_key:
+                        lower = n.token.lower()
+                        existing = _KIND_CANONICAL.get(lower)
+                        if existing is None:
+                            _KIND_CANONICAL[lower] = n.token
+                    prev_kind_key = False
+
+        kept = sorted(t for t, c in counts.items() if c >= min_freq)
+        return {t: 4 + i for i, t in enumerate(kept)}