Initial app: Gradio missing-field suggester (v6.1 model)

README.md

@@ -12,4 +12,36 @@ license: mit
 short_description: Tree-aware transformer for structured (YAML) data
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# YAML-BERT — Kubernetes missing-field suggester
+
+A small (7.8M-param) BERT-style encoder trained on 276K Kubernetes YAML
+manifests with **tree-aware positional encoding** and **hybrid bigram/
+trigram prediction targets**. Paste a YAML manifest below; the model
+identifies fields it expects to see but that are absent, ranked by
+confidence.
+
+This Space runs the **v6.1 checkpoint** — same architecture as v5, with
+a 5-line bug-fix to selective masking. See the project repo for full
+details.
+
+## What the model does
+
+- Predicts standard Kubernetes structural fields (`spec`, `replicas`,
+  `containers.image`, etc.)
+- Distinguishes kind-specific fields (`Deployment.replicas`,
+  `Service.ports`)
+- Calibrated confidence: strong on common patterns, weaker on ambiguous
+  positions
+
+## Known limitations
+
+- Status-side fields not yet predicted (addressed in v6.2)
+- Novel CRD instances and very rare annotation keys may collapse to
+  `[UNK]`
+- Trained on `substratusai/the-stack-yaml-k8s`
+
+## Links
+
+- Code & docs: <https://github.com/vimalk78/yaml-bert>
+- Evaluation results: <https://github.com/vimalk78/yaml-bert/blob/main/docs/evaluation-results.md>
+- v6 plan: <https://github.com/vimalk78/yaml-bert/blob/main/docs/v6-plan.md>

app.py

@@ -0,0 +1,213 @@
+"""YAML-BERT missing-field suggester — Gradio demo.
+
+Paste a Kubernetes YAML manifest; the model identifies fields it expects
+to see but that are absent. Runs the v6.1 checkpoint (Lever 1 — selective
+masking applied during training).
+
+Run locally:
+    pip install gradio
+    PYTHONPATH=. python app.py
+"""
+from __future__ import annotations
+
+import os
+import sys
+
+import gradio as gr
+import torch
+
+from yaml_bert.config import YamlBertConfig
+from yaml_bert.embedding import YamlBertEmbedding
+from yaml_bert.model import YamlBertModel
+from yaml_bert.suggest import suggest_missing_fields
+from yaml_bert.vocab import Vocabulary
+
+
+# ----- Model loading (once at startup) -----
+
+DEFAULT_CHECKPOINT = "output_v6.1_lever1_only_seed42/checkpoints/yaml_bert_v4_epoch_30.pt"
+DEFAULT_VOCAB = "output_v6.1_lever1_only_seed42/vocab.json"
+
+
+def load_model(checkpoint_path: str, vocab_path: str) -> tuple[YamlBertModel, Vocabulary]:
+    vocab = Vocabulary.load(vocab_path)
+    config = YamlBertConfig()
+    emb = YamlBertEmbedding(
+        config=config,
+        key_vocab_size=vocab.key_vocab_size,
+        value_vocab_size=vocab.value_vocab_size,
+    )
+    model = YamlBertModel(
+        config=config,
+        embedding=emb,
+        simple_vocab_size=vocab.simple_target_vocab_size,
+        kind_vocab_size=vocab.kind_target_vocab_size,
+    )
+    cp = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
+    model.load_state_dict(cp["model_state_dict"])
+    model.eval()
+    return model, vocab
+
+
+checkpoint_path = os.environ.get("YAML_BERT_CHECKPOINT", DEFAULT_CHECKPOINT)
+vocab_path = os.environ.get("YAML_BERT_VOCAB", DEFAULT_VOCAB)
+
+print(f"Loading model from {checkpoint_path} (vocab: {vocab_path})", file=sys.stderr)
+MODEL, VOCAB = load_model(checkpoint_path, vocab_path)
+n_params = sum(p.numel() for p in MODEL.parameters())
+print(f"Loaded {n_params:,} parameters", file=sys.stderr)
+
+
+# ----- Inference -----
+
+def suggest(yaml_text: str, threshold: float, top_k: int) -> str:
+    yaml_text = (yaml_text or "").strip()
+    if not yaml_text:
+        return "_Paste a YAML manifest above to see missing-field suggestions._"
+
+    try:
+        suggestions = suggest_missing_fields(
+            MODEL, VOCAB, yaml_text,
+            threshold=threshold, top_k=top_k,
+        )
+    except Exception as e:
+        return f"**Error parsing YAML:**\n```\n{e}\n```"
+
+    if not suggestions:
+        return ("_No suggestions above threshold. Model thinks this YAML is "
+                "either complete, or it doesn't have strong opinions at this confidence level._")
+
+    # Group by parent_path for readability
+    by_parent: dict[str, list[dict]] = {}
+    for s in suggestions:
+        by_parent.setdefault(s.get("parent_path") or "(root)", []).append(s)
+
+    blocks: list[str] = []
+    for parent in sorted(by_parent.keys(), key=lambda p: -max(s["confidence"] for s in by_parent[p])):
+        rows = ["| Missing key | Confidence | Strength |", "|---|---:|---|"]
+        for s in sorted(by_parent[parent], key=lambda x: -x["confidence"]):
+            conf = s["confidence"]
+            strength = "**STRONG**" if conf >= 0.7 else ("MODERATE" if conf >= 0.5 else "weak")
+            rows.append(f"| `{s['missing_key']}` | {conf:.1%} | {strength} |")
+        blocks.append(f"### `{parent}`\n" + "\n".join(rows))
+
+    return "\n\n".join(blocks)
+
+
+# ----- UI -----
+
+EXAMPLE_NGINX = """apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: nginx-deployment
+spec:
+  selector:
+    matchLabels:
+      app: nginx
+  replicas: 2
+  template:
+    metadata:
+      labels:
+        app: nginx
+    spec:
+      containers:
+      - name: nginx
+        image: nginx:1.8
+        resources:
+          limits:
+            memory: "128Mi"
+            cpu: "250m"
+        ports:
+        - containerPort: 80
+"""
+
+EXAMPLE_INCOMPLETE_SERVICE = """apiVersion: v1
+kind: Service
+metadata:
+  name: my-svc
+spec:
+  selector:
+    app: web
+"""
+
+EXAMPLE_CONFIGMAP = """apiVersion: v1
+kind: ConfigMap
+metadata:
+  name: app-config
+data:
+  key1: value1
+"""
+
+with gr.Blocks(title="YAML-BERT — missing-field suggester") as demo:
+    gr.Markdown(
+        f"""
+# YAML-BERT — Kubernetes missing-field suggester
+
+A small (7.8M-param) BERT-style encoder trained on 276K Kubernetes YAML manifests with
+**tree-aware positional encoding** and **hybrid bigram/trigram prediction targets**.
+This page runs the **v6.1 checkpoint** — same architecture as v5, with a 5-line bug
+fix to selective masking ([details](https://github.com/vimalk78/yaml-bert/blob/main/docs/evaluation-results.md#7-v61--lever-1-selective-masking)).
+
+**How it works:** the model walks each parent level in your YAML, inserts a fake
+`[MASK]` node, and reports the keys it expects to see there but that aren't present.
+Confidence reflects how strongly the model "expects" each missing field.
+
+Code: [github.com/vimalk78/yaml-bert](https://github.com/vimalk78/yaml-bert) ·
+{n_params:,} params
+"""
+    )
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            yaml_input = gr.Code(
+                language="yaml",
+                lines=22,
+                label="YAML input",
+                value=EXAMPLE_NGINX,
+            )
+            with gr.Row():
+                threshold = gr.Slider(
+                    minimum=0.1, maximum=0.95, value=0.3, step=0.05,
+                    label="Confidence threshold",
+                )
+                top_k = gr.Slider(
+                    minimum=3, maximum=20, value=10, step=1,
+                    label="Top-K predictions per position",
+                )
+            submit = gr.Button("Suggest missing fields", variant="primary")
+
+        with gr.Column(scale=1):
+            output = gr.Markdown(label="Suggestions", value="")
+
+    submit.click(fn=suggest, inputs=[yaml_input, threshold, top_k], outputs=output)
+    yaml_input.change(fn=suggest, inputs=[yaml_input, threshold, top_k], outputs=output)
+
+    gr.Examples(
+        examples=[
+            [EXAMPLE_NGINX, 0.3, 10],
+            [EXAMPLE_INCOMPLETE_SERVICE, 0.3, 10],
+            [EXAMPLE_CONFIGMAP, 0.3, 10],
+        ],
+        inputs=[yaml_input, threshold, top_k],
+        label="Example YAMLs",
+    )
+
+    gr.Markdown(
+        """
+---
+
+### What this model does well
+- Predicts standard Kubernetes structural fields (`spec`, `replicas`, `containers.image`, etc.)
+- Distinguishes kind-specific fields (`Deployment.replicas`, `Service.ports`)
+- Calibrated confidence: strong on common patterns, weaker on ambiguous positions
+
+### Known limitations
+- **Status-side fields not yet predicted** (still being addressed in v6.2)
+- Novel CRD instances and very rare annotation keys may collapse to `[UNK]`
+- Trained on `substratusai/the-stack-yaml-k8s` (276K manifests from HF)
+"""
+    )
+
+
+if __name__ == "__main__":
+    demo.launch()

output_v6.1_lever1_only_seed42/checkpoints/yaml_bert_v4_epoch_30.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dc057c85c31db30efafca5d89486f629ce0831efc9b9560dfd0d7dfca483b198
+size 93942115

requirements.txt

@@ -0,0 +1,2 @@
+torch>=2.0
+pyyaml>=6.0

yaml_bert/__init__.py

@@ -0,0 +1,26 @@
+"""YAML-BERT: Attention on Kubernetes Structured Data."""
+
+from yaml_bert.config import YamlBertConfig
+from yaml_bert.types import NodeType, YamlNode
+from yaml_bert.linearizer import YamlLinearizer
+from yaml_bert.vocab import Vocabulary, VocabBuilder
+from yaml_bert.annotator import DomainAnnotator
+from yaml_bert.embedding import YamlBertEmbedding
+from yaml_bert.model import YamlBertModel
+from yaml_bert.dataset import YamlDataset, collate_fn
+from yaml_bert.trainer import YamlBertTrainer
+
+__all__ = [
+    "YamlBertConfig",
+    "NodeType",
+    "YamlNode",
+    "YamlLinearizer",
+    "Vocabulary",
+    "VocabBuilder",
+    "DomainAnnotator",
+    "YamlBertEmbedding",
+    "YamlBertModel",
+    "YamlDataset",
+    "collate_fn",
+    "YamlBertTrainer",
+]

yaml_bert/annotator.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+from yaml_bert.types import NodeType, YamlNode
+
+
+class DomainAnnotator:
+    ORDERED_LISTS = {"initContainers"}
+
+    def annotate(self, nodes: list[YamlNode]) -> list[YamlNode]:
+        list_parent_ids = self._find_list_parent_ids(nodes)
+        for node in nodes:
+            if id(node) in list_parent_ids:
+                node.annotations["list_ordered"] = (
+                    node.token in self.ORDERED_LISTS
+                )
+        return nodes
+
+    def _find_list_parent_ids(self, nodes: list[YamlNode]) -> set[int]:
+        list_parent_ids: set[int] = set()
+        parent_paths_with_list_items: set[str] = set()
+
+        for node in nodes:
+            if node.node_type in (NodeType.LIST_KEY, NodeType.LIST_VALUE):
+                parts = node.parent_path.split(".")
+                for i, part in enumerate(parts):
+                    if part.isdigit():
+                        list_parent_path = ".".join(parts[:i])
+                        parent_paths_with_list_items.add(list_parent_path)
+                        break
+
+        for node in nodes:
+            if node.node_type in (NodeType.KEY, NodeType.LIST_KEY):
+                node_full_path = (
+                    f"{node.parent_path}.{node.token}" if node.parent_path else node.token
+                )
+                if node_full_path in parent_paths_with_list_items:
+                    list_parent_ids.add(id(node))
+
+        return list_parent_ids

yaml_bert/attention_pooling.py

@@ -0,0 +1,24 @@
+import math
+
+import torch
+import torch.nn as nn
+
+
+class AttentionPooling(nn.Module):
+    def __init__(self, d_model: int) -> None:
+        super().__init__()
+        self.d_model = d_model
+        self.W_doc = nn.Linear(in_features=d_model, out_features=1)  # (d_model,1)
+
+    def forward(
+        self,
+        h: torch.Tensor,  # Shape(Batch, Seq_Len, d_model)
+    ):
+        scores: torch.Tensor = self.W_doc(h)  # (Batch, Seq_Len, 1)
+        scores = scores / math.sqrt(self.d_model)
+        scores = scores.view(scores.shape[0], 1, scores.shape[1])  # (Batch, 1, Seq_Len)
+        weights: torch.Tensor = scores.softmax(dim=-1)  # (Batch,1, Seq_Len)
+        return (
+            (weights @ h).squeeze(1),
+            weights,
+        )  # returning weights for testing/visualization (Batch,d_model), (Batch,1,Seq_Len)

yaml_bert/cache.py

@@ -0,0 +1,104 @@
+"""Cache linearized YAML documents to disk.
+
+Linearizes all documents once (with multiprocessing), saves to pickle.
+Both VocabBuilder and YamlDataset load from cache — no re-parsing.
+
+Usage:
+    from yaml_bert.cache import build_or_load_cache
+
+    docs = build_or_load_cache(
+        dataset_name="substratusai/the-stack-yaml-k8s",
+        cache_path="output_v3/doc_cache.pkl",
+        max_docs=276520,
+    )
+"""
+from __future__ import annotations
+
+import os
+import pickle
+import time
+from multiprocessing import Pool, cpu_count
+from typing import Any
+
+from yaml_bert.annotator import DomainAnnotator
+from yaml_bert.linearizer import YamlLinearizer
+from yaml_bert.types import YamlNode
+
+
+def _linearize_one(yaml_content: str) -> list[YamlNode] | None:
+    """Linearize a single YAML string. Used by multiprocessing pool."""
+    try:
+        linearizer = YamlLinearizer()
+        annotator = DomainAnnotator()
+        nodes: list[YamlNode] = linearizer.linearize(yaml_content)
+        if nodes:
+            annotator.annotate(nodes)
+            return nodes
+    except Exception:
+        pass
+    return None
+
+
+def build_or_load_cache(
+    dataset_name: str,
+    cache_path: str,
+    max_docs: int | None = None,
+    num_workers: int | None = None,
+) -> list[list[YamlNode]]:
+    """Build or load cached linearized documents.
+
+    Args:
+        dataset_name: HuggingFace dataset ID
+        cache_path: Path to save/load the cache pickle
+        max_docs: Max documents to process (None = all)
+        num_workers: Number of parallel workers (None = cpu_count)
+
+    Returns:
+        List of linearized documents (each is a list of YamlNode)
+    """
+    if os.path.exists(cache_path):
+        print(f"Loading cached documents from {cache_path}")
+        start: float = time.time()
+        with open(cache_path, "rb") as f:
+            documents: list[list[YamlNode]] = pickle.load(f)
+        print(f"Loaded {len(documents):,} documents in {time.time() - start:.1f}s")
+        return documents
+
+    from datasets import load_dataset
+
+    print(f"Building document cache 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"Linearizing {total:,} documents with {num_workers or cpu_count()} workers...")
+
+    # Extract YAML content strings
+    yaml_contents: list[str] = [ds[i]["content"] for i in range(total)]
+
+    # Parallel linearization with progress bar
+    from tqdm import tqdm
+
+    start = time.time()
+    workers: int = num_workers or cpu_count()
+    with Pool(workers) as pool:
+        results: list[list[YamlNode] | None] = list(tqdm(
+            pool.imap(_linearize_one, yaml_contents, chunksize=500),
+            total=len(yaml_contents),
+            desc="Linearizing",
+        ))
+
+    documents = [doc for doc in results if doc is not None]
+    skipped: int = sum(1 for doc in results if doc is None)
+    elapsed: float = time.time() - start
+
+    print(f"Linearized {len(documents):,} documents in {elapsed:.1f}s ({skipped} skipped)")
+
+    # Save cache
+    os.makedirs(os.path.dirname(cache_path) or ".", exist_ok=True)
+    with open(cache_path, "wb") as f:
+        pickle.dump(documents, f)
+
+    cache_size: float = os.path.getsize(cache_path) / 1024 / 1024
+    print(f"Cache saved: {cache_path} ({cache_size:.1f} MB)")
+
+    return documents

yaml_bert/config.py

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import Enum
+
+
+class TreePosVariant(str, Enum):
+    """Tree positional encoding ablation variants."""
+
+    FULL = "full"             # depth + sibling + node_type (default)
+    NO_DEPTH = "no_depth"     # sibling + node_type
+    NO_SIBLING = "no_sibling" # depth + node_type
+    SEQUENTIAL = "sequential" # learned pos[seq_idx] + node_type
+
+
+@dataclass
+class YamlBertConfig:
+    """Central configuration for YAML-BERT hyperparameters."""
+
+    # Model architecture
+    d_model: int = 256
+    num_layers: int = 6
+    num_heads: int = 8
+    d_ff: int = 0  # 0 means "auto: 4 * d_model"
+    max_depth: int = 16
+    max_sibling: int = 32
+    tree_pos_variant: TreePosVariant = TreePosVariant.FULL
+
+    # Training
+    mask_prob: float = 0.15
+    lr: float = 1e-4
+    batch_size: int = 32
+    num_epochs: int = 30
+    max_seq_len: int = 512
+    skip_unk_targets: bool = True  # v6: don't supervise on positions whose target is [UNK]
+
+    # Auxiliary loss weights
+    aux_kind_weight: float = 0.1     # α: kind classification loss weight
+    aux_parent_weight: float = 0.1   # β: parent_key classification loss weight
+
+    def __post_init__(self) -> None:
+        if self.d_ff == 0:
+            self.d_ff = 4 * self.d_model

yaml_bert/dataset.py

@@ -0,0 +1,367 @@
+from __future__ import annotations
+
+import glob
+import os
+import random
+
+import torch
+from torch.utils.data import Dataset
+
+from yaml_bert.annotator import DomainAnnotator
+from yaml_bert.config import YamlBertConfig
+from yaml_bert.linearizer import YamlLinearizer
+from yaml_bert.types import NodeType, YamlNode
+from yaml_bert.vocab import Vocabulary, compute_target
+
+
+def _extract_kind(nodes: list[YamlNode]) -> str:
+    """Extract the kind value from a document's node list, normalized to canonical casing."""
+    from yaml_bert.vocab import normalize_kind
+    for i, node in enumerate(nodes):
+        if (node.token == "kind"
+            and node.depth == 0
+            and node.node_type == NodeType.KEY
+            and i + 1 < len(nodes)
+            and nodes[i + 1].node_type == NodeType.VALUE):
+            return normalize_kind(nodes[i + 1].token)
+    return ""
+
+
+# NodeType to integer index mapping
+_NODE_TYPE_INDEX: dict[NodeType, int] = {
+    NodeType.KEY: 0,
+    NodeType.VALUE: 1,
+    NodeType.LIST_KEY: 2,
+    NodeType.LIST_VALUE: 3,
+}
+
+# Node types eligible for masking
+_MASKABLE_TYPES: set[NodeType] = {NodeType.KEY, NodeType.LIST_KEY}
+
+
+class YamlDataset(Dataset):
+    """Dataset of linearized, masked YAML documents for YAML-BERT training."""
+
+    def __init__(
+        self,
+        yaml_dir: str,
+        vocab: Vocabulary,
+        linearizer: YamlLinearizer,
+        annotator: DomainAnnotator,
+        config: YamlBertConfig | None = None,
+    ) -> None:
+        config = config or YamlBertConfig()
+        self.vocab: Vocabulary = vocab
+        self.linearizer: YamlLinearizer = linearizer
+        self.annotator: DomainAnnotator = annotator
+        self.mask_prob: float = config.mask_prob
+        self.max_seq_len: int = config.max_seq_len
+
+        # Load and linearize all YAML files. Skip any files the linearizer
+        # can't handle (e.g., uncommon tags like !!value) rather than failing
+        # the whole load.
+        yaml_files: list[str] = sorted(
+            glob.glob(os.path.join(yaml_dir, "**", "*.yaml"), recursive=True)
+        )
+        self.documents: list[list[YamlNode]] = []
+        self.document_kinds: list[str] = []
+        skipped: int = 0
+        for path in yaml_files:
+            try:
+                nodes: list[YamlNode] = linearizer.linearize_file(path)
+            except Exception:
+                skipped += 1
+                continue
+            if nodes:
+                annotator.annotate(nodes)
+                self.documents.append(nodes)
+                self.document_kinds.append(_extract_kind(nodes))
+        if skipped:
+            print(f"Skipped {skipped} unparseable YAML files")
+
+    @classmethod
+    def from_huggingface(
+        cls,
+        dataset_name: str,
+        vocab: Vocabulary,
+        linearizer: YamlLinearizer,
+        annotator: DomainAnnotator,
+        config: YamlBertConfig | None = None,
+        max_docs: int | None = None,
+    ) -> "YamlDataset":
+        """Load dataset from HuggingFace Hub.
+
+        Args:
+            dataset_name: HuggingFace dataset ID, e.g. "substratusai/the-stack-yaml-k8s"
+            vocab: Pre-built vocabulary
+            linearizer: YamlLinearizer instance
+            annotator: DomainAnnotator instance
+            config: Optional config (uses defaults if None)
+            max_docs: Load at most this many documents (None = all)
+        """
+        from datasets import load_dataset
+
+        config = config or YamlBertConfig()
+        instance: YamlDataset = cls.__new__(cls)
+        instance.vocab = vocab
+        instance.linearizer = linearizer
+        instance.annotator = annotator
+        instance.mask_prob = config.mask_prob
+        instance.max_seq_len = config.max_seq_len
+        instance.documents = []
+        instance.document_kinds = []
+
+        print(f"Loading dataset: {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"Linearizing {total:,} / {len(ds):,} documents...")
+
+        skipped: int = 0
+        for i in range(total):
+            yaml_content: str = ds[i]["content"]
+            try:
+                nodes: list[YamlNode] = linearizer.linearize(yaml_content)
+            except Exception:
+                skipped += 1
+                continue
+            if nodes:
+                annotator.annotate(nodes)
+                instance.documents.append(nodes)
+                instance.document_kinds.append(_extract_kind(nodes))
+
+            if (i + 1) % 10000 == 0:
+                print(f"  {i + 1:,} / {total:,} processed ({skipped} skipped)")
+
+        print(f"Loaded {len(instance.documents):,} documents ({skipped} skipped)")
+        return instance
+
+    @classmethod
+    def from_cached_docs(
+        cls,
+        documents: list[list[YamlNode]],
+        vocab: Vocabulary,
+        config: YamlBertConfig | None = None,
+    ) -> "YamlDataset":
+        """Build dataset from pre-linearized cached documents. No parsing needed."""
+        config = config or YamlBertConfig()
+        instance: YamlDataset = cls.__new__(cls)
+        instance.vocab = vocab
+        instance.linearizer = None
+        instance.annotator = None
+        instance.mask_prob = config.mask_prob
+        instance.max_seq_len = config.max_seq_len
+        instance.documents = documents
+        instance.document_kinds = [_extract_kind(doc) for doc in documents]
+        return instance
+
+    @classmethod
+    def from_cached_docs_v4(
+        cls,
+        documents: list[list[YamlNode]],
+        vocab: Vocabulary,
+        config: YamlBertConfig | None = None,
+    ) -> "YamlDataset":
+        """Build v4 dataset from pre-linearized cached documents."""
+        config = config or YamlBertConfig()
+        instance = cls.__new__(cls)
+        instance.vocab = vocab
+        instance.linearizer = None
+        instance.annotator = None
+        instance.mask_prob = config.mask_prob
+        instance.max_seq_len = config.max_seq_len
+        instance.skip_unk_targets = config.skip_unk_targets
+        instance.documents = documents
+        instance.document_kinds = [_extract_kind(doc) for doc in documents]
+        instance._v4_mode = True  # Flag for __getitem__
+        return instance
+
+    def __len__(self) -> int:
+        return len(self.documents)
+
+    def _getitem_v4(self, idx: int) -> dict[str, torch.Tensor]:
+        """v4 item encoding: hybrid simple/kind labels, no parent_key_ids or kind_ids."""
+        nodes: list[YamlNode] = self.documents[idx]
+
+        # Truncate if needed
+        if len(nodes) > self.max_seq_len:
+            nodes = nodes[: self.max_seq_len]
+
+        seq_len: int = len(nodes)
+        kind: str = self.document_kinds[idx]
+
+        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))
+
+        # Apply masking
+        simple_labels: list[int] = [-100] * seq_len
+        kind_labels: list[int] = [-100] * seq_len
+        mask_token_id: int = self.vocab.special_tokens["[MASK]"]
+        unk_id: int = self.vocab.special_tokens["[UNK]"]
+        skip_unk: bool = getattr(self, "skip_unk_targets", True)
+
+        for i in range(seq_len):
+            if nodes[i].node_type not in _MASKABLE_TYPES:
+                continue
+            if random.random() >= self.mask_prob:
+                continue
+
+            # Compute hybrid target FIRST so we can skip positions whose
+            # target is [UNK] (Lever 1, v6.1 — don't train the model to
+            # confidently predict [UNK])
+            target, head_type = compute_target(nodes[i], kind)
+            if head_type == "simple":
+                target_id = self.vocab.encode_simple_target(target)
+            else:
+                target_id = self.vocab.encode_kind_target(target)
+
+            if skip_unk and target_id == unk_id:
+                continue
+
+            if head_type == "simple":
+                simple_labels[i] = target_id
+                kind_labels[i] = -100
+            else:
+                kind_labels[i] = target_id
+                simple_labels[i] = -100
+
+            # Apply token replacement (80/10/10 rule)
+            rand: float = random.random()
+            if rand < 0.8:
+                token_ids[i] = mask_token_id
+            elif rand < 0.9:
+                random_key_id: int = random.randint(
+                    len(self.vocab.special_tokens),
+                    len(self.vocab.key_vocab) + len(self.vocab.special_tokens) - 1,
+                )
+                token_ids[i] = random_key_id
+            # else 10%: keep unchanged
+
+        return {
+            "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),
+            "simple_labels": torch.tensor(simple_labels, dtype=torch.long),
+            "kind_labels": torch.tensor(kind_labels, dtype=torch.long),
+        }
+
+    def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
+        if getattr(self, "_v4_mode", False):
+            return self._getitem_v4(idx)
+
+        nodes: list[YamlNode] = self.documents[idx]
+
+        # Truncate if needed
+        if len(nodes) > self.max_seq_len:
+            nodes = nodes[: self.max_seq_len]
+
+        seq_len: int = len(nodes)
+
+        # Encode nodes to integer arrays
+        token_ids: list[int] = []
+        node_types: list[int] = []
+        depths: list[int] = []
+        sibling_indices: list[int] = []
+        parent_key_ids: list[int] = []
+
+        for node in nodes:
+            # Token ID — route to correct vocab
+            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))  # clamp to max_depth - 1
+            sibling_indices.append(min(node.sibling_index, 31))  # clamp to max_sibling - 1
+
+            parent_key: str = Vocabulary.extract_parent_key(node.parent_path)
+            parent_key_ids.append(self.vocab.encode_key(parent_key))
+
+        # Apply masking (only to KEY and LIST_KEY nodes)
+        labels: list[int] = [-100] * seq_len
+        mask_token_id: int = self.vocab.special_tokens["[MASK]"]
+
+        for i in range(seq_len):
+            if nodes[i].node_type not in _MASKABLE_TYPES:
+                continue
+            if random.random() >= self.mask_prob:
+                continue
+
+            # This position is selected for masking
+            labels[i] = token_ids[i]  # save original token ID as label
+
+            rand: float = random.random()
+            if rand < 0.8:
+                # 80%: replace with [MASK]
+                token_ids[i] = mask_token_id
+            elif rand < 0.9:
+                # 10%: replace with random key
+                random_key_id: int = random.randint(
+                    len(self.vocab.special_tokens),
+                    len(self.vocab.key_vocab) + len(self.vocab.special_tokens) - 1,
+                )
+                token_ids[i] = random_key_id
+            # else 10%: keep unchanged
+
+        kind: str = self.document_kinds[idx]
+        kind_id: int = self.vocab.encode_kind(kind)
+        kind_ids: list[int] = [kind_id] * seq_len
+
+        return {
+            "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),
+            "parent_key_ids": torch.tensor(parent_key_ids, dtype=torch.long),
+            "labels": torch.tensor(labels, dtype=torch.long),
+            "kind_ids": torch.tensor(kind_ids, dtype=torch.long),
+        }
+
+
+def collate_fn(batch: list[dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]:
+    """Pad a batch of variable-length sequences and create padding mask."""
+    max_len: int = max(item["token_ids"].size(0) for item in batch)
+
+    padded: dict[str, list[torch.Tensor]] = {key: [] for key in batch[0].keys()}
+    padding_masks: list[torch.Tensor] = []
+
+    for item in batch:
+        seq_len: int = item["token_ids"].size(0)
+        pad_len: int = max_len - seq_len
+
+        for key in item:
+            if pad_len > 0:
+                pad_value: int = -100 if (key == "labels" or key.endswith("_labels")) else 0
+                padding: torch.Tensor = torch.full(
+                    (pad_len,), pad_value, dtype=torch.long
+                )
+                padded[key].append(torch.cat([item[key], padding]))
+            else:
+                padded[key].append(item[key])
+
+        mask: torch.Tensor = 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)
+        padding_masks.append(mask)
+
+    result: dict[str, torch.Tensor] = {
+        key: torch.stack(tensors) for key, tensors in padded.items()
+    }
+    result["padding_mask"] = torch.stack(padding_masks)
+
+    return result

yaml_bert/embedding.py

@@ -0,0 +1,80 @@
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from yaml_bert.config import TreePosVariant, YamlBertConfig
+
+
+class YamlBertEmbedding(nn.Module):
+    """Embedding layer with tree positional encoding for YAML-BERT.
+
+    Produces input vectors by summing:
+    - Token embedding (key_embedding or value_embedding, routed by node_type)
+    - 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,
+    ) -> 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.node_type_embedding: nn.Embedding = nn.Embedding(4, d)
+
+        use_depth: bool = variant in (TreePosVariant.FULL, TreePosVariant.NO_SIBLING)
+        use_sibling: bool = variant in (TreePosVariant.FULL, TreePosVariant.NO_DEPTH)
+        use_seq_pos: bool = variant == TreePosVariant.SEQUENTIAL
+
+        self.depth_embedding: nn.Embedding | None = (
+            nn.Embedding(config.max_depth, d) if use_depth else None
+        )
+        self.sibling_embedding: nn.Embedding | None = (
+            nn.Embedding(config.max_sibling, d) if use_sibling else None
+        )
+        self.pos_embedding: nn.Embedding | None = (
+            nn.Embedding(config.max_seq_len, d) if use_seq_pos else None
+        )
+
+        self.layer_norm: nn.LayerNorm = nn.LayerNorm(d)
+
+    def forward(
+        self,
+        token_ids: torch.Tensor,
+        node_types: torch.Tensor,
+        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)
+
+        tree_pos: torch.Tensor = 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:
+            tree_pos = tree_pos + self.sibling_embedding(sibling_indices)
+        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 = (
+                torch.arange(seq_len, device=token_ids.device)
+                .clamp(max=max_pos - 1)
+                .unsqueeze(0)
+                .expand(token_ids.size(0), seq_len)
+            )
+            tree_pos = tree_pos + self.pos_embedding(positions)
+
+        return self.layer_norm(token_emb + tree_pos)

yaml_bert/evaluate.py

@@ -0,0 +1,240 @@
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+
+from yaml_bert.dataset import YamlDataset
+from yaml_bert.model import YamlBertModel
+from yaml_bert.vocab import Vocabulary
+
+
+class YamlBertEvaluator:
+    """Post-training evaluation for YAML-BERT."""
+
+    def __init__(
+        self,
+        model: YamlBertModel,
+        dataset: YamlDataset,
+        vocab: Vocabulary,
+    ) -> None:
+        self.model: YamlBertModel = model
+        self.dataset: YamlDataset = dataset
+        self.vocab: Vocabulary = vocab
+        self.device: torch.device = next(model.parameters()).device
+        self._id_to_simple: dict[int, str] = {
+            v: k for k, v in vocab.simple_target_vocab.items()
+        }
+        self._id_to_kind: dict[int, str] = {
+            v: k for k, v in vocab.kind_target_vocab.items()
+        }
+
+    def _decode_simple(self, id: int) -> str:
+        if id in self.vocab._id_to_special:
+            return self.vocab._id_to_special[id]
+        return self._id_to_simple.get(id, "[UNK]")
+
+    def _decode_kind(self, id: int) -> str:
+        if id in self.vocab._id_to_special:
+            return self.vocab._id_to_special[id]
+        return self._id_to_kind.get(id, "[UNK]")
+
+    @torch.no_grad()
+    def evaluate_prediction_accuracy(self) -> dict[str, float]:
+        """Compute top-1 and top-5 masked key prediction accuracy over the dataset.
+
+        Evaluates both simple and kind-specific heads independently.
+        """
+        self.model.eval()
+
+        simple_total: int = 0
+        simple_top1: int = 0
+        simple_top5: int = 0
+        kind_total: int = 0
+        kind_top1: int = 0
+        kind_top5: int = 0
+
+        for idx in range(len(self.dataset)):
+            item: dict[str, torch.Tensor] = self.dataset[idx]
+            simple_labels: torch.Tensor = item["simple_labels"]
+            kind_labels: torch.Tensor = item["kind_labels"]
+
+            simple_masked: torch.Tensor = simple_labels != -100
+            kind_masked: torch.Tensor = kind_labels != -100
+            if not simple_masked.any() and not kind_masked.any():
+                continue
+
+            simple_logits, kind_logits = self.model(
+                token_ids=item["token_ids"].unsqueeze(0).to(self.device),
+                node_types=item["node_types"].unsqueeze(0).to(self.device),
+                depths=item["depths"].unsqueeze(0).to(self.device),
+                sibling_indices=item["sibling_indices"].unsqueeze(0).to(self.device),
+            )
+
+            s_logits: torch.Tensor = simple_logits[0]
+            for pos in simple_masked.nonzero(as_tuple=True)[0]:
+                true_id: int = simple_labels[pos].item()
+                pos_logits: torch.Tensor = s_logits[pos]
+                top5_ids: torch.Tensor = pos_logits.topk(5).indices
+
+                if top5_ids[0].item() == true_id:
+                    simple_top1 += 1
+                if true_id in top5_ids.tolist():
+                    simple_top5 += 1
+                simple_total += 1
+
+            k_logits: torch.Tensor = kind_logits[0]
+            for pos in kind_masked.nonzero(as_tuple=True)[0]:
+                true_id = kind_labels[pos].item()
+                pos_logits = k_logits[pos]
+                top5_ids = pos_logits.topk(5).indices
+
+                if top5_ids[0].item() == true_id:
+                    kind_top1 += 1
+                if true_id in top5_ids.tolist():
+                    kind_top5 += 1
+                kind_total += 1
+
+        total_masked: int = simple_total + kind_total
+        total_top1: int = simple_top1 + kind_top1
+        total_top5: int = simple_top5 + kind_top5
+
+        return {
+            "top1_accuracy": total_top1 / max(total_masked, 1),
+            "top5_accuracy": total_top5 / max(total_masked, 1),
+            "total_masked": total_masked,
+            "simple_top1_accuracy": simple_top1 / max(simple_total, 1),
+            "kind_top1_accuracy": kind_top1 / max(kind_total, 1),
+        }
+
+    @torch.no_grad()
+    def analyze_embeddings(self) -> list[dict[str, Any]]:
+        """Compare embeddings of the same key at different tree positions."""
+        self.model.eval()
+        results: list[dict[str, Any]] = []
+
+        test_pairs: list[dict[str, Any]] = [
+            {
+                "key": "spec",
+                "position_a": {"depth": 0},
+                "position_b": {"depth": 2},
+            },
+            {
+                "key": "name",
+                "position_a": {"depth": 1},
+                "position_b": {"depth": 1},
+            },
+        ]
+
+        for pair in test_pairs:
+            key_id: int = self.vocab.encode_key(pair["key"])
+
+            token_ids: torch.Tensor = torch.tensor(
+                [[key_id, key_id]], device=self.device
+            )
+            node_types: torch.Tensor = torch.tensor(
+                [[0, 0]], device=self.device
+            )
+            depths: torch.Tensor = torch.tensor(
+                [[pair["position_a"]["depth"], pair["position_b"]["depth"]]],
+                device=self.device,
+            )
+            siblings: torch.Tensor = torch.tensor(
+                [[0, 0]], device=self.device
+            )
+
+            embeddings: torch.Tensor = self.model.embedding(
+                token_ids, node_types, depths, siblings
+            )
+
+            cosine_sim: float = F.cosine_similarity(
+                embeddings[0, 0].unsqueeze(0),
+                embeddings[0, 1].unsqueeze(0),
+            ).item()
+
+            results.append({
+                "key": pair["key"],
+                "position_a": pair["position_a"],
+                "position_b": pair["position_b"],
+                "cosine_similarity": cosine_sim,
+            })
+
+        return results
+
+    @torch.no_grad()
+    def top_k_predictions(
+        self, doc_idx: int, k: int = 5
+    ) -> list[dict[str, Any]]:
+        """Show top-k predicted keys for each masked position in a document.
+
+        Reports predictions from both the simple and kind-specific heads.
+        """
+        self.model.eval()
+
+        item: dict[str, torch.Tensor] = self.dataset[doc_idx]
+        simple_labels: torch.Tensor = item["simple_labels"]
+        kind_labels: torch.Tensor = item["kind_labels"]
+
+        simple_masked: torch.Tensor = simple_labels != -100
+        kind_masked: torch.Tensor = kind_labels != -100
+        if not simple_masked.any() and not kind_masked.any():
+            return []
+
+        simple_logits, kind_logits = self.model(
+            token_ids=item["token_ids"].unsqueeze(0).to(self.device),
+            node_types=item["node_types"].unsqueeze(0).to(self.device),
+            depths=item["depths"].unsqueeze(0).to(self.device),
+            sibling_indices=item["sibling_indices"].unsqueeze(0).to(self.device),
+        )
+
+        predictions: list[dict[str, Any]] = []
+
+        # Simple head predictions
+        s_logits: torch.Tensor = simple_logits[0]
+        for pos in simple_masked.nonzero(as_tuple=True)[0]:
+            true_id: int = simple_labels[pos].item()
+            pos_logits: torch.Tensor = s_logits[pos]
+            probs: torch.Tensor = F.softmax(pos_logits, dim=-1)
+            topk: torch.return_types.topk = probs.topk(k)
+
+            predicted_keys: list[dict[str, Any]] = [
+                {
+                    "key": self._decode_simple(topk.indices[i].item()),
+                    "probability": topk.values[i].item(),
+                }
+                for i in range(k)
+            ]
+
+            predictions.append({
+                "position": pos.item(),
+                "head": "simple",
+                "true_key": self._decode_simple(true_id),
+                "predicted_keys": predicted_keys,
+            })
+
+        # Kind-specific head predictions
+        k_logits: torch.Tensor = kind_logits[0]
+        for pos in kind_masked.nonzero(as_tuple=True)[0]:
+            true_id = kind_labels[pos].item()
+            pos_logits = k_logits[pos]
+            probs = F.softmax(pos_logits, dim=-1)
+            topk = probs.topk(k)
+
+            predicted_keys = [
+                {
+                    "key": self._decode_kind(topk.indices[i].item()),
+                    "probability": topk.values[i].item(),
+                }
+                for i in range(k)
+            ]
+
+            predictions.append({
+                "position": pos.item(),
+                "head": "kind",
+                "true_key": self._decode_kind(true_id),
+                "predicted_keys": predicted_keys,
+            })
+
+        predictions.sort(key=lambda p: p["position"])
+        return predictions

yaml_bert/linearizer.py

@@ -0,0 +1,99 @@
+from __future__ import annotations
+
+import yaml
+
+from yaml_bert.types import NodeType, YamlNode
+
+
+class YamlLinearizer:
+    def linearize(self, yaml_string: str) -> list[YamlNode]:
+        data = yaml.safe_load(yaml_string)
+        if data is None:
+            return []
+        nodes: list[YamlNode] = []
+        self._walk(data, depth=0, parent_path="", nodes=nodes, in_list=False)
+        return nodes
+
+    def _walk(
+        self,
+        data,
+        depth: int,
+        parent_path: str,
+        nodes: list[YamlNode],
+        in_list: bool,
+    ) -> None:
+        if isinstance(data, dict):
+            for sibling_index, (key, value) in enumerate(data.items()):
+                key_str = str(key)
+                key_type = NodeType.LIST_KEY if in_list else NodeType.KEY
+                nodes.append(
+                    YamlNode(
+                        token=key_str,
+                        node_type=key_type,
+                        depth=depth,
+                        sibling_index=sibling_index,
+                        parent_path=parent_path,
+                    )
+                )
+                if isinstance(value, dict):
+                    child_path = f"{parent_path}.{key_str}" if parent_path else key_str
+                    self._walk(value, depth + 1, child_path, nodes, in_list=False)
+                elif isinstance(value, list):
+                    child_path = f"{parent_path}.{key_str}" if parent_path else key_str
+                    self._walk_list(value, depth + 1, child_path, nodes)
+                else:
+                    value_path = f"{parent_path}.{key_str}" if parent_path else key_str
+                    value_type = NodeType.LIST_VALUE if in_list else NodeType.VALUE
+                    nodes.append(
+                        YamlNode(
+                            token=str(value),
+                            node_type=value_type,
+                            depth=depth,
+                            sibling_index=sibling_index,
+                            parent_path=value_path,
+                        )
+                    )
+
+    def linearize_file(self, path: str) -> list[YamlNode]:
+        with open(path) as f:
+            content = f.read()
+        nodes: list[YamlNode] = []
+        for doc in yaml.safe_load_all(content):
+            if doc is None:
+                continue
+            self._walk(doc, depth=0, parent_path="", nodes=nodes, in_list=False)
+        return nodes
+
+    def linearize_multi_doc(self, yaml_string: str) -> list[list[YamlNode]]:
+        result = []
+        for doc in yaml.safe_load_all(yaml_string):
+            if doc is None:
+                continue
+            nodes: list[YamlNode] = []
+            self._walk(doc, depth=0, parent_path="", nodes=nodes, in_list=False)
+            result.append(nodes)
+        return result
+
+    def _walk_list(
+        self,
+        data: list,
+        depth: int,
+        parent_path: str,
+        nodes: list[YamlNode],
+    ) -> None:
+        for item_index, item in enumerate(data):
+            item_path = f"{parent_path}.{item_index}"
+            if isinstance(item, dict):
+                self._walk(item, depth, item_path, nodes, in_list=True)
+            elif isinstance(item, list):
+                self._walk_list(item, depth, item_path, nodes)
+            else:
+                nodes.append(
+                    YamlNode(
+                        token=str(item),
+                        node_type=NodeType.LIST_VALUE,
+                        depth=depth,
+                        sibling_index=item_index,
+                        parent_path=item_path,
+                    )
+                )

yaml_bert/model.py

@@ -0,0 +1,77 @@
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from yaml_bert.config import YamlBertConfig
+from yaml_bert.embedding import YamlBertEmbedding
+
+
+class YamlBertModel(nn.Module):
+    """YAML-BERT: Transformer encoder with tree positional encoding.
+
+    Two prediction heads: simple (bigram) + kind-specific (trigram).
+    """
+
+    def __init__(
+        self,
+        config: YamlBertConfig,
+        embedding: YamlBertEmbedding,
+        simple_vocab_size: int,
+        kind_vocab_size: int,
+    ) -> None:
+        super().__init__()
+        self.embedding: YamlBertEmbedding = embedding
+
+        encoder_layer: nn.TransformerEncoderLayer = nn.TransformerEncoderLayer(
+            d_model=config.d_model,
+            nhead=config.num_heads,
+            dim_feedforward=config.d_ff,
+            batch_first=True,
+        )
+        self.encoder: nn.TransformerEncoder = nn.TransformerEncoder(
+            encoder_layer, num_layers=config.num_layers
+        )
+
+        self.simple_head: nn.Linear = nn.Linear(config.d_model, simple_vocab_size)
+        self.kind_head: nn.Linear = nn.Linear(config.d_model, kind_vocab_size)
+        self.loss_fn: nn.CrossEntropyLoss = nn.CrossEntropyLoss(ignore_index=-100)
+
+    def forward(
+        self,
+        token_ids: torch.Tensor,
+        node_types: torch.Tensor,
+        depths: torch.Tensor,
+        sibling_indices: torch.Tensor,
+        padding_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        x: torch.Tensor = self.embedding(token_ids, node_types, depths, sibling_indices)
+        x = self.encoder(x, src_key_padding_mask=padding_mask)
+        simple_logits: torch.Tensor = self.simple_head(x)
+        kind_logits: torch.Tensor = self.kind_head(x)
+        return simple_logits, kind_logits
+
+    def compute_loss(
+        self,
+        simple_logits: torch.Tensor,
+        simple_labels: torch.Tensor,
+        kind_logits: torch.Tensor,
+        kind_labels: torch.Tensor,
+    ) -> tuple[torch.Tensor, dict[str, float]]:
+        simple_loss: torch.Tensor = self.loss_fn(
+            simple_logits.view(-1, simple_logits.size(-1)),
+            simple_labels.view(-1),
+        )
+
+        # Kind loss: only compute if there are kind-specific labels in this batch
+        has_kind: bool = (kind_labels != -100).any().item()
+        if has_kind:
+            kind_loss: torch.Tensor = self.loss_fn(
+                kind_logits.view(-1, kind_logits.size(-1)),
+                kind_labels.view(-1),
+            )
+        else:
+            kind_loss = torch.tensor(0.0, device=simple_logits.device)
+
+        total: torch.Tensor = simple_loss + kind_loss
+        return total, {"simple": simple_loss.item(), "kind": kind_loss.item()}

yaml_bert/pooling.py

@@ -0,0 +1,57 @@
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class DocumentPooling(nn.Module):
+    """Pooling by Multi-head Attention.
+
+    The kind node queries all other nodes via cross-attention
+    to produce a single document embedding.
+    """
+
+    def __init__(self, d_model: int, num_heads: int = 4) -> None:
+        super().__init__()
+        self.query_proj: nn.Linear = nn.Linear(d_model, d_model)
+        self.cross_attn: nn.MultiheadAttention = nn.MultiheadAttention(
+            d_model, num_heads, batch_first=True,
+        )
+        self.layer_norm: nn.LayerNorm = nn.LayerNorm(d_model)
+
+    def forward(
+        self,
+        kind_hidden: torch.Tensor,
+        all_hidden: torch.Tensor,
+    ) -> torch.Tensor:
+        query: torch.Tensor = self.query_proj(kind_hidden)
+        doc_emb, _ = self.cross_attn(query, all_hidden, all_hidden)
+        doc_emb = self.layer_norm(doc_emb)
+        return doc_emb.squeeze(1)
+
+
+def supervised_contrastive_loss(
+    embeddings: torch.Tensor,
+    labels: torch.Tensor,
+    temperature: float = 0.1,
+) -> torch.Tensor:
+    embeddings = F.normalize(embeddings, dim=1)
+    batch_size: int = embeddings.shape[0]
+
+    sim: torch.Tensor = embeddings @ embeddings.T / temperature
+
+    label_mask: torch.Tensor = labels.unsqueeze(0) == labels.unsqueeze(1)
+    self_mask: torch.Tensor = torch.eye(batch_size, dtype=torch.bool, device=embeddings.device)
+    label_mask = label_mask & ~self_mask
+
+    sim_max, _ = sim.max(dim=1, keepdim=True)
+    sim = sim - sim_max.detach()
+
+    exp_sim: torch.Tensor = torch.exp(sim) * (~self_mask).float()
+    log_prob: torch.Tensor = sim - torch.log(exp_sim.sum(dim=1, keepdim=True) + 1e-8)
+
+    pos_count: torch.Tensor = label_mask.float().sum(dim=1).clamp(min=1)
+    loss: torch.Tensor = -(label_mask.float() * log_prob).sum(dim=1) / pos_count
+
+    return loss.mean()

yaml_bert/similarity.py

@@ -0,0 +1,86 @@
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+
+from yaml_bert.annotator import DomainAnnotator
+from yaml_bert.linearizer import YamlLinearizer
+from yaml_bert.model import YamlBertModel
+from yaml_bert.pooling import DocumentPooling
+from yaml_bert.types import NodeType, YamlNode
+from yaml_bert.vocab import Vocabulary
+
+
+_NODE_TYPE_INDEX: dict[NodeType, int] = {
+    NodeType.KEY: 0, NodeType.VALUE: 1,
+    NodeType.LIST_KEY: 2, NodeType.LIST_VALUE: 3,
+}
+
+
+@torch.no_grad()
+def extract_hidden_states(
+    model: YamlBertModel,
+    vocab: Vocabulary,
+    yaml_text: str,
+) -> tuple[torch.Tensor, int]:
+    linearizer: YamlLinearizer = YamlLinearizer()
+    annotator: DomainAnnotator = DomainAnnotator()
+    # Use linearize_multi_doc to handle --- separators, take the first doc
+    docs: list[list[YamlNode]] = linearizer.linearize_multi_doc(yaml_text)
+    nodes: list[YamlNode] = docs[0] if docs else []
+    if not nodes:
+        return torch.empty(0), -1
+    annotator.annotate(nodes)
+
+    token_ids: list[int] = []
+    node_types: list[int] = []
+    depths: list[int] = []
+    siblings: list[int] = []
+    kind_pos: int = -1
+
+    for i, node in enumerate(nodes):
+        if node.node_type in (NodeType.KEY, NodeType.LIST_KEY):
+            token_ids.append(vocab.encode_key(node.token))
+        else:
+            token_ids.append(vocab.encode_value(node.token))
+        node_types.append(_NODE_TYPE_INDEX[node.node_type])
+        depths.append(min(node.depth, 15))
+        siblings.append(min(node.sibling_index, 31))
+
+        if node.token == "kind" and node.depth == 0 and node.node_type == NodeType.KEY and kind_pos == -1:
+            kind_pos = i
+
+    t = lambda x: torch.tensor([x])
+    model.eval()
+
+    x: torch.Tensor = model.embedding(
+        t(token_ids), t(node_types), t(depths), t(siblings),
+    )
+    for layer in model.encoder.layers:
+        x = layer(x)
+
+    return x.squeeze(0), kind_pos
+
+
+@torch.no_grad()
+def get_document_embedding(
+    model: YamlBertModel,
+    pooling: DocumentPooling,
+    vocab: Vocabulary,
+    yaml_text: str,
+) -> torch.Tensor:
+    hidden, kind_pos = extract_hidden_states(model, vocab, yaml_text)
+    if hidden.shape[0] == 0 or kind_pos < 0:
+        return torch.zeros(hidden.shape[1] if hidden.dim() > 1 else 1)
+
+    pooling.eval()
+    kind_hidden: torch.Tensor = hidden[kind_pos].unsqueeze(0).unsqueeze(0)
+    all_hidden: torch.Tensor = hidden.unsqueeze(0)
+    return pooling(kind_hidden, all_hidden).squeeze(0)
+
+
+def cosine_similarity_matrix(embeddings: torch.Tensor) -> torch.Tensor:
+    normed: torch.Tensor = F.normalize(embeddings, dim=1)
+    return normed @ normed.T

yaml_bert/suggest.py

@@ -0,0 +1,236 @@
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+
+from yaml_bert.annotator import DomainAnnotator
+from yaml_bert.dataset import _extract_kind
+from yaml_bert.linearizer import YamlLinearizer
+from yaml_bert.model import YamlBertModel
+from yaml_bert.types import NodeType, YamlNode
+from yaml_bert.vocab import Vocabulary, compute_target
+
+
+# Keys managed by the cluster, not written by users
+_CLUSTER_MANAGED_KEYS: set[str] = {
+    "status", "creationTimestamp", "generation", "resourceVersion",
+    "selfLink", "uid", "managedFields",
+}
+
+_NODE_TYPE_INDEX: dict[NodeType, int] = {
+    NodeType.KEY: 0,
+    NodeType.VALUE: 1,
+    NodeType.LIST_KEY: 2,
+    NodeType.LIST_VALUE: 3,
+}
+
+
+def suggest_missing_fields(
+    model: YamlBertModel,
+    vocab: Vocabulary,
+    yaml_text: str,
+    threshold: float = 0.3,
+    top_k: int = 10,
+) -> list[dict[str, Any]]:
+    """Suggest missing fields in a YAML document based on model conventions.
+
+    Uses the dual-head (simple + kind-specific) prediction approach.
+    For each parent level, a fake [MASK] node is appended and both heads
+    are consulted based on which head would apply at that tree position.
+
+    Args:
+        model: Trained YAML-BERT model
+        vocab: Vocabulary
+        yaml_text: Raw YAML text
+        threshold: Minimum confidence to report a missing field
+        top_k: Number of predictions per masked position
+
+    Returns:
+        List of suggestions sorted by confidence (highest first).
+        Each suggestion: {"parent_path": str, "missing_key": str, "confidence": float}
+    """
+    linearizer: YamlLinearizer = YamlLinearizer()
+    annotator: DomainAnnotator = DomainAnnotator()
+    nodes: list[YamlNode] = linearizer.linearize(yaml_text)
+    if not nodes:
+        return []
+    annotator.annotate(nodes)
+
+    token_ids, node_types, depths, siblings = _encode_nodes(nodes, vocab)
+
+    kind: str = _extract_kind(nodes)
+    mask_id: int = vocab.special_tokens["[MASK]"]
+
+    # Build reverse lookups for decoding predictions
+    id_to_simple: dict[int, str] = {v: k for k, v in vocab.simple_target_vocab.items()}
+    id_to_kind: dict[int, str] = {v: k for k, v in vocab.kind_target_vocab.items()}
+    id_to_special: dict[int, str] = {v: k for k, v in vocab.special_tokens.items()}
+
+    # Group key nodes by parent_path
+    keys_by_parent: dict[str, set[str]] = {}
+    key_positions_by_parent: dict[str, list[int]] = {}
+
+    for i, node in enumerate(nodes):
+        if node.node_type in (NodeType.KEY, NodeType.LIST_KEY):
+            keys_by_parent.setdefault(node.parent_path, set()).add(node.token)
+            key_positions_by_parent.setdefault(node.parent_path, []).append(i)
+
+    # For each parent level, append a fake masked node as the "next sibling"
+    # to discover missing keys. Filter out noise: self-references, keys that
+    # already exist at root level, and special tokens.
+    model.eval()
+    predicted_keys_by_parent: dict[str, dict[str, float]] = {}
+    skipped_by_parent: dict[str, list[dict[str, Any]]] = {}
+
+    # Collect all existing keys across the entire document for cross-reference filtering
+    all_root_keys: set[str] = {
+        n.token for n in nodes
+        if n.node_type in (NodeType.KEY, NodeType.LIST_KEY) and n.depth == 0
+    }
+
+    t = lambda x: torch.tensor([x])
+
+    for parent_path, positions in key_positions_by_parent.items():
+        predicted: dict[str, float] = {}
+        # Extract parent key name, skipping numeric list indices
+        # e.g., "spec.containers.0" → "containers", not "0"
+        parent_key_name: str = ""
+        if parent_path:
+            for part in reversed(parent_path.split(".")):
+                if not part.isdigit():
+                    parent_key_name = part
+                    break
+
+        # Use the last key node at this level as a template for the fake node
+        last_pos: int = positions[-1]
+        last_node: YamlNode = nodes[last_pos]
+        next_sibling: int = min(last_node.sibling_index + 1, 31)
+
+        # Find the insertion point: right after the last sibling's subtree.
+        # Walk forward from last_pos until we find a node at the same or
+        # shallower depth (i.e., the subtree under last_pos has ended).
+        insert_pos: int = last_pos + 1
+        while insert_pos < len(token_ids):
+            if nodes[insert_pos].depth <= last_node.depth:
+                break
+            insert_pos += 1
+
+        # Insert a fake [MASK] node at the correct position in the sequence
+        fake_token_ids: list[int] = token_ids[:insert_pos] + [mask_id] + token_ids[insert_pos:]
+        fake_node_types: list[int] = node_types[:insert_pos] + [_NODE_TYPE_INDEX[last_node.node_type]] + node_types[insert_pos:]
+        fake_depths: list[int] = depths[:insert_pos] + [last_node.depth] + depths[insert_pos:]
+        fake_siblings: list[int] = siblings[:insert_pos] + [next_sibling] + siblings[insert_pos:]
+        fake_pos: int = insert_pos
+
+        # Determine which head applies for this position using compute_target logic.
+        # Create a fake YamlNode to probe which head to use.
+        fake_node = YamlNode(
+            token="__probe__",
+            node_type=last_node.node_type,
+            depth=last_node.depth,
+            sibling_index=next_sibling,
+            parent_path=last_node.parent_path,
+        )
+        _, head_type = compute_target(fake_node, kind)
+
+        with torch.no_grad():
+            simple_logits, kind_logits = model(
+                t(fake_token_ids), t(fake_node_types), t(fake_depths),
+                t(fake_siblings),
+            )
+
+        if head_type == "simple":
+            probs: torch.Tensor = F.softmax(simple_logits[0, fake_pos], dim=-1)
+            id_to_target = id_to_simple
+        else:
+            probs = F.softmax(kind_logits[0, fake_pos], dim=-1)
+            id_to_target = id_to_kind
+
+        topk = probs.topk(top_k + 5)
+
+        for j in range(topk.indices.shape[0]):
+            target_id: int = topk.indices[j].item()
+            prob: float = topk.values[j].item()
+
+            # Decode target string
+            if target_id in id_to_special:
+                target_name: str = id_to_special[target_id]
+            else:
+                target_name = id_to_target.get(target_id, "[UNK]")
+
+            # Extract the raw key name and validate parent from composite targets.
+            # Simple targets: "parent::key" or just "key" (root level)
+            # Kind targets: "Kind::parent::key"
+            parts: list[str] = target_name.split("::")
+            key_name: str = parts[-1]
+
+            # Validate that the predicted target's parent matches where we're probing.
+            # E.g., "spec::imagePullSecrets" should only be accepted when probing under spec,
+            # not under matchLabels or metadata.
+            if len(parts) >= 2:
+                predicted_parent: str = parts[-2]
+                if predicted_parent != parent_key_name:
+                    skipped_by_parent.setdefault(parent_path, []).append({
+                        "target": target_name,
+                        "key": key_name,
+                        "predicted_parent": predicted_parent,
+                        "actual_parent": parent_key_name,
+                        "confidence": prob,
+                    })
+                    continue
+
+            # Skip special tokens
+            if key_name in ("[PAD]", "[UNK]", "[MASK]"):
+                continue
+            # Skip self-referencing (key same as parent key name)
+            if key_name == parent_key_name:
+                continue
+            # Skip root-level keys suggested as children (e.g., roleRef under subjects)
+            if parent_path and key_name in all_root_keys and last_node.depth > 0:
+                continue
+
+            predicted[key_name] = prob
+
+        predicted_keys_by_parent[parent_path] = predicted
+
+    # Find missing keys
+    suggestions: list[dict[str, Any]] = []
+
+    for parent_path, predicted in predicted_keys_by_parent.items():
+        existing: set[str] = keys_by_parent.get(parent_path, set())
+        for key_name, confidence in predicted.items():
+            if (key_name not in existing
+                    and key_name not in _CLUSTER_MANAGED_KEYS
+                    and confidence >= threshold):
+                suggestions.append({
+                    "parent_path": parent_path,
+                    "missing_key": key_name,
+                    "confidence": confidence,
+                })
+
+    suggestions.sort(key=lambda s: -s["confidence"])
+    return suggestions, skipped_by_parent
+
+
+def _encode_nodes(
+    nodes: list[YamlNode],
+    vocab: Vocabulary,
+) -> tuple[list[int], list[int], list[int], list[int]]:
+    """Encode nodes to integer lists for model input."""
+    token_ids: list[int] = []
+    node_types: list[int] = []
+    depths: list[int] = []
+    siblings: list[int] = []
+
+    for node in nodes:
+        if node.node_type in (NodeType.KEY, NodeType.LIST_KEY):
+            token_ids.append(vocab.encode_key(node.token))
+        else:
+            token_ids.append(vocab.encode_value(node.token))
+        node_types.append(_NODE_TYPE_INDEX[node.node_type])
+        depths.append(min(node.depth, 15))
+        siblings.append(min(node.sibling_index, 31))
+
+    return token_ids, node_types, depths, siblings

yaml_bert/trainer.py

@@ -0,0 +1,127 @@
+from __future__ import annotations
+
+import os
+
+import torch
+from torch.optim import AdamW
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+from yaml_bert.config import YamlBertConfig
+from yaml_bert.dataset import YamlDataset, collate_fn
+from yaml_bert.model import YamlBertModel
+
+
+class YamlBertTrainer:
+    """Training loop with hybrid loss from two prediction heads."""
+
+    def __init__(
+        self,
+        config: YamlBertConfig,
+        model: YamlBertModel,
+        dataset: YamlDataset,
+        checkpoint_dir: str | None = None,
+        checkpoint_every: int = 1,
+        resume_from: str | None = None,
+    ) -> None:
+        self.config = config
+        self.model = model
+        self.dataset = dataset
+        self.checkpoint_dir = checkpoint_dir
+        self.checkpoint_every = checkpoint_every
+        self.resume_from = resume_from
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    def train(self) -> list[float]:
+        from datetime import datetime
+        self.model.to(self.device)
+        self.model.train()
+
+        print(f"Training started: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
+        num_params = sum(p.numel() for p in self.model.parameters())
+        print(f"Model parameters: {num_params:,}")
+        print(f"Config: d_model={self.config.d_model}, layers={self.config.num_layers}, heads={self.config.num_heads}")
+        print(f"Device: {self.device}")
+
+        optimizer = AdamW(self.model.parameters(), lr=self.config.lr, weight_decay=0.01)
+        start_epoch = 0
+
+        if self.resume_from:
+            checkpoint = torch.load(self.resume_from, map_location=self.device)
+            self.model.load_state_dict(checkpoint["model_state_dict"])
+            optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
+            start_epoch = checkpoint["epoch"]
+            print(f"Resumed from epoch {start_epoch}")
+
+        dataloader = DataLoader(
+            self.dataset, batch_size=self.config.batch_size,
+            shuffle=True, collate_fn=collate_fn,
+        )
+
+        epoch_losses: list[float] = []
+        for epoch in range(start_epoch, self.config.num_epochs):
+            total_loss: float = 0.0
+            num_batches: int = 0
+            running_breakdown: dict[str, float] = {}
+
+            pbar = tqdm(dataloader, desc=f"Epoch {epoch+1}/{self.config.num_epochs}", leave=True)
+            for batch in pbar:
+                batch = {k: v.to(self.device) for k, v in batch.items()}
+                optimizer.zero_grad()
+
+                simple_logits, kind_logits = self.model(
+                    token_ids=batch["token_ids"],
+                    node_types=batch["node_types"],
+                    depths=batch["depths"],
+                    sibling_indices=batch["sibling_indices"],
+                    padding_mask=batch["padding_mask"],
+                )
+
+                loss, breakdown = self.model.compute_loss(
+                    simple_logits, batch["simple_labels"],
+                    kind_logits, batch["kind_labels"],
+                )
+
+                if torch.isnan(loss):
+                    continue  # skip bad batches
+
+                loss.backward()
+                torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
+                optimizer.step()
+
+                total_loss += loss.item()
+                for k, v in breakdown.items():
+                    running_breakdown[k] = running_breakdown.get(k, 0.0) + v
+                num_batches += 1
+
+                postfix = {"loss": f"{total_loss/num_batches:.4f}"}
+                for k in ["simple", "kind"]:
+                    if k in running_breakdown:
+                        postfix[k] = f"{running_breakdown[k]/num_batches:.4f}"
+                pbar.set_postfix(**postfix)
+
+            avg_loss: float = total_loss / max(num_batches, 1)
+            epoch_losses.append(avg_loss)
+            breakdown_str: str = " | ".join(
+                f"{k}: {v/max(num_batches,1):.4f}" for k, v in sorted(running_breakdown.items())
+            )
+            print(f"Epoch {epoch+1}/{self.config.num_epochs} — loss: {avg_loss:.4f} ({breakdown_str})")
+
+            if self.checkpoint_dir and (epoch+1) % self.checkpoint_every == 0:
+                self._save_checkpoint(epoch+1, optimizer)
+
+        if self.checkpoint_dir:
+            self._save_checkpoint(self.config.num_epochs, optimizer)
+
+        return epoch_losses
+
+    def _save_checkpoint(self, epoch: int, optimizer: AdamW) -> None:
+        os.makedirs(self.checkpoint_dir, exist_ok=True)
+        path = os.path.join(self.checkpoint_dir, f"yaml_bert_v4_epoch_{epoch}.pt")
+        torch.save({
+            "epoch": epoch,
+            "model_state_dict": self.model.state_dict(),
+            "optimizer_state_dict": optimizer.state_dict(),
+            "tree_pos_variant": self.config.tree_pos_variant.value,
+        }, path)
+        print(f"Checkpoint saved: {path}")

yaml_bert/types.py

@@ -0,0 +1,29 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any
+
+
+class NodeType(Enum):
+    KEY = "KEY"
+    VALUE = "VALUE"
+    LIST_KEY = "LIST_KEY"
+    LIST_VALUE = "LIST_VALUE"
+
+
+@dataclass
+class YamlNode:
+    token: str
+    node_type: NodeType
+    depth: int
+    sibling_index: int
+    parent_path: str
+    annotations: dict[str, Any] = field(default_factory=dict)
+
+    def __repr__(self) -> str:
+        return (
+            f"YamlNode({self.token!r}, {self.node_type.value}, "
+            f"depth={self.depth}, sibling={self.sibling_index}, "
+            f"path={self.parent_path!r})"
+        )

yaml_bert/visualize.py

@@ -0,0 +1,145 @@
+from __future__ import annotations
+
+from typing import Any
+
+import matplotlib
+matplotlib.use("Agg")  # non-interactive backend
+import matplotlib.pyplot as plt
+import torch
+
+
+def plot_training_loss(
+    losses: list[float],
+    output_path: str = "training_loss.png",
+    title: str = "YAML-BERT Training Loss",
+) -> None:
+    """Plot training loss curve over epochs."""
+    fig, ax = plt.subplots(figsize=(10, 6))
+    ax.plot(range(1, len(losses) + 1), losses, marker="o", linewidth=2)
+    ax.set_xlabel("Epoch")
+    ax.set_ylabel("Loss")
+    ax.set_title(title)
+    ax.grid(True, alpha=0.3)
+    fig.tight_layout()
+    fig.savefig(output_path, dpi=150)
+    plt.close(fig)
+    print(f"Training loss plot saved: {output_path}")
+
+
+def plot_accuracy(
+    results: dict[str, float],
+    output_path: str = "accuracy.png",
+    title: str = "YAML-BERT Key Prediction Accuracy",
+) -> None:
+    """Plot top-1 and top-5 prediction accuracy as a bar chart."""
+    labels: list[str] = ["Top-1", "Top-5"]
+    values: list[float] = [results["top1_accuracy"], results["top5_accuracy"]]
+
+    fig, ax = plt.subplots(figsize=(6, 5))
+    bars = ax.bar(labels, values, color=["steelblue", "coral"], width=0.5)
+    ax.set_ylabel("Accuracy")
+    ax.set_title(title)
+    ax.set_ylim(0, 1)
+
+    for bar, val in zip(bars, values):
+        ax.text(
+            bar.get_x() + bar.get_width() / 2,
+            bar.get_height() + 0.02,
+            f"{val:.1%}",
+            ha="center",
+            fontsize=14,
+            fontweight="bold",
+        )
+
+    total: float = results.get("total_masked", 0)
+    ax.text(
+        0.5, -0.1,
+        f"Total masked positions: {int(total)}",
+        ha="center",
+        transform=ax.transAxes,
+        fontsize=10,
+        color="gray",
+    )
+
+    fig.tight_layout()
+    fig.savefig(output_path, dpi=150)
+    plt.close(fig)
+    print(f"Accuracy plot saved: {output_path}")
+
+
+def plot_embedding_similarity(
+    results: list[dict[str, Any]],
+    output_path: str = "embedding_similarity.png",
+    title: str = "Tree Position Embedding Similarity",
+) -> None:
+    """Plot cosine similarity between same-key embeddings at different tree positions."""
+    labels: list[str] = []
+    similarities: list[float] = []
+
+    for r in results:
+        pa: dict[str, Any] = r["position_a"]
+        pb: dict[str, Any] = r["position_b"]
+        label: str = (
+            f"{r['key']}\n"
+            f"d={pa['depth']},p={pa['parent_key']}\n"
+            f"vs d={pb['depth']},p={pb['parent_key']}"
+        )
+        labels.append(label)
+        similarities.append(r["cosine_similarity"])
+
+    fig, ax = plt.subplots(figsize=(max(8, len(results) * 3), 6))
+    bars = ax.bar(range(len(results)), similarities, color="steelblue")
+    ax.set_xticks(range(len(results)))
+    ax.set_xticklabels(labels, fontsize=9)
+    ax.set_ylabel("Cosine Similarity")
+    ax.set_title(title)
+    ax.set_ylim(-1, 1)
+    ax.axhline(y=0, color="gray", linestyle="--", alpha=0.5)
+
+    for bar, sim in zip(bars, similarities):
+        ax.text(
+            bar.get_x() + bar.get_width() / 2,
+            bar.get_height() + 0.02,
+            f"{sim:.3f}",
+            ha="center",
+            fontsize=10,
+        )
+
+    fig.tight_layout()
+    fig.savefig(output_path, dpi=150)
+    plt.close(fig)
+    print(f"Embedding similarity plot saved: {output_path}")
+
+
+def plot_attention_patterns(
+    attention_weights: torch.Tensor,
+    token_labels: list[str],
+    output_path: str = "attention_patterns.png",
+    title: str = "Attention Patterns",
+) -> None:
+    """Plot attention heatmaps for each head.
+
+    Args:
+        attention_weights: (num_heads, seq_len, seq_len)
+        token_labels: labels for each position in the sequence
+    """
+    num_heads: int = attention_weights.shape[0]
+    fig, axes = plt.subplots(1, num_heads, figsize=(6 * num_heads, 5))
+
+    if num_heads == 1:
+        axes = [axes]
+
+    for head_idx, ax in enumerate(axes):
+        weights: torch.Tensor = attention_weights[head_idx].cpu()
+        im = ax.imshow(weights.numpy(), cmap="Blues", vmin=0, vmax=1)
+        ax.set_title(f"Head {head_idx}")
+        ax.set_xticks(range(len(token_labels)))
+        ax.set_yticks(range(len(token_labels)))
+        ax.set_xticklabels(token_labels, rotation=45, ha="right", fontsize=7)
+        ax.set_yticklabels(token_labels, fontsize=7)
+
+    fig.suptitle(title)
+    fig.tight_layout()
+    fig.savefig(output_path, dpi=150)
+    plt.close(fig)
+    print(f"Attention pattern plot saved: {output_path}")

yaml_bert/vocab.py

@@ -0,0 +1,345 @@
+from __future__ import annotations
+
+import json
+from yaml_bert.types import NodeType, YamlNode
+
+
+SPECIAL_TOKENS = ["[PAD]", "[UNK]", "[MASK]"]
+
+UNIVERSAL_ROOT_KEYS: set[str] = {"apiVersion", "kind", "metadata"}
+
+# Canonical casing for known Kubernetes kinds.
+# Maps lowercase → canonical form.
+_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)
+
+
+def compute_target(node: YamlNode, kind: str) -> tuple[str, str]:
+    """Compute hybrid prediction target.
+
+    Returns (target_string, head_type) where head_type is "simple" or "kind_specific".
+    """
+    if node.depth == 0:
+        return node.token, "simple"
+
+    parent_key: str = Vocabulary.extract_parent_key(node.parent_path)
+
+    if node.depth == 1 and parent_key not in UNIVERSAL_ROOT_KEYS and parent_key != "":
+        return f"{kind}::{parent_key}::{node.token}", "kind_specific"
+
+    return (f"{parent_key}::{node.token}" if parent_key else node.token), "simple"
+
+
+class Vocabulary:
+    def __init__(
+        self,
+        key_vocab: dict[str, int],
+        value_vocab: dict[str, int],
+        special_tokens: dict[str, int],
+        kind_vocab: dict[str, int] | None = None,
+        simple_target_vocab: dict[str, int] | None = None,
+        kind_target_vocab: dict[str, int] | None = None,
+    ) -> None:
+        self.key_vocab = key_vocab
+        self.value_vocab = value_vocab
+        self.special_tokens = special_tokens
+        self.kind_vocab = kind_vocab or {"[NO_KIND]": 0}
+        self.simple_target_vocab = simple_target_vocab or {}
+        self.kind_target_vocab = kind_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()}
+
+    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 encode_kind(self, kind: str) -> int:
+        if not kind:
+            return self.kind_vocab["[NO_KIND]"]
+        return self.kind_vocab.get(kind, self.kind_vocab["[NO_KIND]"])
+
+    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_simple_target(self, target: str) -> int:
+        return self.simple_target_vocab.get(target, self.special_tokens["[UNK]"])
+
+    def encode_kind_target(self, target: str) -> int:
+        return self.kind_target_vocab.get(target, self.special_tokens["[UNK]"])
+
+    @staticmethod
+    def extract_parent_key(parent_path: str) -> str:
+        """Extract the last non-numeric component from a parent_path.
+
+        Examples:
+            "spec.template.spec.containers.0" -> "containers"
+            "metadata" -> "metadata"
+            "" -> ""
+        """
+        if not parent_path:
+            return ""
+        parts = parent_path.split(".")
+        for part in reversed(parts):
+            if not part.isdigit():
+                return part
+        return ""
+
+    @property
+    def key_vocab_size(self) -> int:
+        return len(self.key_vocab) + len(self.special_tokens)
+
+    @property
+    def value_vocab_size(self) -> int:
+        return len(self.value_vocab) + len(self.special_tokens)
+
+    @property
+    def kind_vocab_size(self) -> int:
+        return len(self.kind_vocab)
+
+    @property
+    def simple_target_vocab_size(self) -> int:
+        return len(self.simple_target_vocab) + len(self.special_tokens)
+
+    @property
+    def kind_target_vocab_size(self) -> int:
+        return len(self.kind_target_vocab) + len(self.special_tokens)
+
+    def save(self, path: str) -> None:
+        data = {
+            "key_vocab": self.key_vocab,
+            "value_vocab": self.value_vocab,
+            "special_tokens": self.special_tokens,
+            "kind_vocab": self.kind_vocab,
+            "simple_target_vocab": self.simple_target_vocab,
+            "kind_target_vocab": self.kind_target_vocab,
+        }
+        with open(path, "w") as f:
+            json.dump(data, f, indent=2)
+
+    @classmethod
+    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"],
+            kind_vocab=data.get("kind_vocab"),
+            simple_target_vocab=data.get("simple_target_vocab", {}),
+            kind_target_vocab=data.get("kind_target_vocab", {}),
+        )
+
+
+class VocabBuilder:
+    def build(self, nodes: list[YamlNode], min_freq: int = 1) -> Vocabulary:
+        key_counts: dict[str, int] = {}
+        value_counts: dict[str, int] = {}
+        kind_set: set[str] = set()
+        simple_target_counts: dict[str, int] = {}
+        kind_target_counts: dict[str, int] = {}
+
+        # First pass: collect all kind values to build canonical casing map
+        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
+        current_kind: str = ""
+        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)
+                target, head_type = compute_target(node, current_kind)
+                if head_type == "kind_specific":
+                    kind_target_counts[target] = kind_target_counts.get(target, 0) + 1
+                else:
+                    simple_target_counts[target] = simple_target_counts.get(target, 0) + 1
+            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:
+                    normalized = normalize_kind(node.token)
+                    kind_set.add(normalized)
+                    current_kind = normalized
+                prev_was_kind_key = False
+
+        # Filter target vocabs by min_freq
+        simple_target_set: set[str] = {t for t, c in simple_target_counts.items() if c >= min_freq}
+        kind_target_set: set[str] = {t for t, c in kind_target_counts.items() if c >= min_freq}
+
+        return self.build_from_counts(
+            key_counts, value_counts, min_freq, kind_set,
+            simple_target_set, kind_target_set,
+        )
+
+    @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,
+        simple_target_set: set[str] | None = None,
+        kind_target_set: set[str] | None = None,
+    ) -> Vocabulary:
+        """Build vocabulary from pre-computed token counts."""
+        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 >= 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))
+        }
+
+        kind_vocab: dict[str, int] = {"[NO_KIND]": 0}
+        for i, kind in enumerate(sorted(kind_set or [])):
+            kind_vocab[kind] = i + 1
+
+        simple_target_vocab = {
+            target: i + offset
+            for i, target in enumerate(sorted(simple_target_set or []))
+        }
+
+        kind_target_vocab = {
+            target: i + offset
+            for i, target in enumerate(sorted(kind_target_set or []))
+        }
+
+        return Vocabulary(key_vocab, value_vocab, special_tokens, kind_vocab,
+                          simple_target_vocab, kind_target_vocab)
+
+    @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)")
+
+    @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)
+        """
+        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)