code/preprocess.py

"""Dataset loading, leakage-safe splitting and GraphCodeBERT feature building.

Design notes
------------
The PoolC fold contains 6.7 M *pairs* but only ~45 k *distinct snippets* -- the
pairs are combinations of a small pool of solutions.  So the expensive work
(tree-sitter parsing, data-flow extraction, BPE tokenisation) is done **once per
distinct snippet** and cached on disk; a pair is then just two integer indices
plus a label.  Nothing proportional to 6.7 M rows is ever tokenised, and the
graph-guided attention masks are materialised lazily in the collator.
"""

from __future__ import annotations

import hashlib
import json
import logging
import time
from collections import Counter
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Iterator

import numpy as np
import pyarrow.parquet as pq
import torch
from torch.utils.data import Dataset

from config import (
    CODE1_COLUMN,
    CODE2_COLUMN,
    DATASET_LANGUAGE,
    FORBIDDEN_FEATURE_COLUMNS,
    GROUP1_COLUMN,
    GROUP2_COLUMN,
    LABEL_COLUMN,
    Config,
)
from dfg_parser import DataFlowExtractionError, extract_dataflow

logger = logging.getLogger(__name__)

_META_COLUMNS = [LABEL_COLUMN, GROUP1_COLUMN, GROUP2_COLUMN]


# --------------------------------------------------------------------------- #
# 1. Schema verification
# --------------------------------------------------------------------------- #
def verify_dataset_schema(dataset_name: str) -> dict[str, Any]:
    """Check the remote dataset still matches what this pipeline expects.

    Raises immediately (rather than silently mis-reading columns) if the schema
    drifts. Returns a report dict for the experiment log.
    """
    from datasets import load_dataset_builder

    builder = load_dataset_builder(dataset_name)
    features = builder.info.features
    splits = {k: v.num_examples for k, v in (builder.info.splits or {}).items()}

    missing = [
        c
        for c in (CODE1_COLUMN, CODE2_COLUMN, LABEL_COLUMN, GROUP1_COLUMN, GROUP2_COLUMN)
        if c not in features
    ]
    if missing:
        raise ValueError(
            f"{dataset_name} is missing expected columns {missing}. "
            f"Available: {sorted(features)}. Adapt config.py before continuing."
        )
    for col in (CODE1_COLUMN, CODE2_COLUMN):
        if features[col].dtype != "string":
            raise ValueError(f"Column {col!r} must be a string, got {features[col]}.")

    report = {
        "dataset_name": dataset_name,
        "features": {k: str(v) for k, v in features.items()},
        "splits": splits,
        "forbidden_feature_columns": list(FORBIDDEN_FEATURE_COLUMNS),
        "language": DATASET_LANGUAGE,
    }
    logger.info("Dataset schema OK: %s", json.dumps(report["splits"]))
    return report


def _parquet_files(dataset_name: str, split: str) -> list[str]:
    """Resolve the local parquet shards for one split (downloads on first use)."""
    from huggingface_hub import snapshot_download

    root = Path(
        snapshot_download(dataset_name, repo_type="dataset", allow_patterns=["data/*", "*.json"])
    )
    files = sorted(root.glob(f"data/{split}-*.parquet"))
    if not files:
        files = sorted(root.glob(f"**/{split}-*.parquet"))
    if not files:
        raise FileNotFoundError(
            f"No parquet shards for split {split!r} under {root}. "
            f"Found: {[p.name for p in root.rglob('*.parquet')]}"
        )
    return [str(p) for p in files]


# --------------------------------------------------------------------------- #
# 2. Snippet pool + pair index (cached)
# --------------------------------------------------------------------------- #
@dataclass
class SplitIndex:
    """A split reduced to integer indices into the shared snippet pool."""

    name: str
    snippet_id1: np.ndarray  # int32 [n_pairs]
    snippet_id2: np.ndarray  # int32 [n_pairs]
    labels: np.ndarray  # int8  [n_pairs]
    group1: np.ndarray  # int32 [n_pairs]
    group2: np.ndarray  # int32 [n_pairs]

    def __len__(self) -> int:
        return int(self.labels.shape[0])

    def select(self, rows: np.ndarray, name: str | None = None) -> "SplitIndex":
        return SplitIndex(
            name=name or self.name,
            snippet_id1=self.snippet_id1[rows],
            snippet_id2=self.snippet_id2[rows],
            labels=self.labels[rows],
            group1=self.group1[rows],
            group2=self.group2[rows],
        )

    def class_distribution(self) -> dict[str, Any]:
        counts = Counter(self.labels.tolist())
        n = max(len(self), 1)
        return {
            "num_examples": len(self),
            "negatives_label_0": int(counts.get(0, 0)),
            "positives_label_1": int(counts.get(1, 0)),
            "positive_ratio": round(counts.get(1, 0) / n, 6),
            "num_groups": int(np.unique(np.concatenate([self.group1, self.group2])).size),
        }

    def snippet_ids(self) -> np.ndarray:
        return np.unique(np.concatenate([self.snippet_id1, self.snippet_id2]))


def _hash_code(text: str) -> bytes:
    return hashlib.blake2b(text.encode("utf-8", "ignore"), digest_size=16).digest()


def build_snippet_pool(
    cfg: Config, splits: tuple[str, ...]
) -> tuple[list[str], dict[str, SplitIndex], dict[str, Any]]:
    """Deduplicate every snippet across ``splits`` and index the pairs.

    Cached under ``cfg.cache_dir`` -- the scan over the parquet shards is the
    only pass that ever touches all 6.7 M rows, and it runs once.
    """
    cache = Path(cfg.cache_dir) / "pairs" / _fingerprint(cfg.dataset_name, splits)
    if (cache / "meta.json").exists():
        logger.info("Reusing cached snippet pool at %s", cache)
        return _load_pool(cache)

    cache.mkdir(parents=True, exist_ok=True)
    t0 = time.time()
    code_to_id: dict[bytes, int] = {}
    snippets: list[str] = []
    indices: dict[str, SplitIndex] = {}
    #: A snippet appearing under two different group ids is a (rare) duplicate
    #: solution; we log it because it is the only within-split ambiguity.
    snippet_groups: dict[int, set[int]] = {}

    for split in splits:
        files = _parquet_files(cfg.dataset_name, split)
        id1_chunks, id2_chunks, lab_chunks, g1_chunks, g2_chunks = [], [], [], [], []
        for path in files:
            pf = pq.ParquetFile(path)
            for batch in pf.iter_batches(
                batch_size=50_000,
                columns=[CODE1_COLUMN, CODE2_COLUMN, *_META_COLUMNS],
            ):
                cols = batch.to_pydict()
                for code_col, group_col, out in (
                    (CODE1_COLUMN, GROUP1_COLUMN, id1_chunks),
                    (CODE2_COLUMN, GROUP2_COLUMN, id2_chunks),
                ):
                    ids = np.empty(len(cols[code_col]), dtype=np.int32)
                    for i, (text, group) in enumerate(zip(cols[code_col], cols[group_col])):
                        key = _hash_code(text)
                        sid = code_to_id.get(key)
                        if sid is None:
                            sid = len(snippets)
                            code_to_id[key] = sid
                            snippets.append(text)
                        ids[i] = sid
                        snippet_groups.setdefault(sid, set()).add(int(group))
                    out.append(ids)
                lab_chunks.append(np.asarray(cols[LABEL_COLUMN], dtype=np.int8))
                g1_chunks.append(np.asarray(cols[GROUP1_COLUMN], dtype=np.int32))
                g2_chunks.append(np.asarray(cols[GROUP2_COLUMN], dtype=np.int32))
            logger.info("  scanned %s", Path(path).name)

        idx = SplitIndex(
            name=split,
            snippet_id1=np.concatenate(id1_chunks),
            snippet_id2=np.concatenate(id2_chunks),
            labels=np.concatenate(lab_chunks),
            group1=np.concatenate(g1_chunks),
            group2=np.concatenate(g2_chunks),
        )
        _assert_label_matches_groups(idx)
        indices[split] = idx
        logger.info("Split %s: %s", split, json.dumps(idx.class_distribution()))

    ambiguous = sorted(sid for sid, gs in snippet_groups.items() if len(gs) > 1)
    leakage = _cross_split_leakage(indices)
    stats = {
        "num_unique_snippets": len(snippets),
        "scan_seconds": round(time.time() - t0, 1),
        "snippets_with_multiple_groups": len(ambiguous),
        "cross_split_snippet_overlap": leakage,
        "per_split": {k: v.class_distribution() for k, v in indices.items()},
    }
    _save_pool(cache, snippets, indices, stats)
    logger.info("Snippet pool: %s", json.dumps(stats, indent=2))
    return snippets, indices, stats


def _assert_label_matches_groups(idx: SplitIndex) -> None:
    """The label is exactly ``group1 == group2``; assert it and shout about it.

    This is why ``code1_group``/``code2_group`` are on the forbidden list: they
    are a perfect proxy for the target.
    """
    implied = (idx.group1 == idx.group2).astype(np.int8)
    mismatches = int((implied != idx.labels).sum())
    if mismatches:
        raise ValueError(
            f"Split {idx.name}: {mismatches} rows where `similar` disagrees with "
            "(code1_group == code2_group). The dataset changed; revisit the split logic."
        )


def _cross_split_leakage(indices: dict[str, SplitIndex]) -> dict[str, int]:
    """Count snippets and groups shared between splits (must be zero)."""
    out: dict[str, int] = {}
    names = list(indices)
    for i, a in enumerate(names):
        for b in names[i + 1 :]:
            sa, sb = set(indices[a].snippet_ids().tolist()), set(indices[b].snippet_ids().tolist())
            ga = set(np.unique(np.concatenate([indices[a].group1, indices[a].group2])).tolist())
            gb = set(np.unique(np.concatenate([indices[b].group1, indices[b].group2])).tolist())
            out[f"{a}|{b}:snippets"] = len(sa & sb)
            out[f"{a}|{b}:groups"] = len(ga & gb)
    return out


def _fingerprint(*parts: Any) -> str:
    return hashlib.blake2b(repr(parts).encode(), digest_size=8).hexdigest()


def _save_pool(
    cache: Path, snippets: list[str], indices: dict[str, SplitIndex], stats: dict
) -> None:
    import pyarrow as pa

    pq.write_table(pa.table({"code": snippets}), cache / "snippets.parquet")
    for name, idx in indices.items():
        np.savez(
            cache / f"{name}.npz",
            snippet_id1=idx.snippet_id1,
            snippet_id2=idx.snippet_id2,
            labels=idx.labels,
            group1=idx.group1,
            group2=idx.group2,
        )
    (cache / "meta.json").write_text(
        json.dumps({"splits": list(indices), "stats": stats}, indent=2), encoding="utf-8"
    )


def _load_pool(cache: Path) -> tuple[list[str], dict[str, SplitIndex], dict[str, Any]]:
    meta = json.loads((cache / "meta.json").read_text(encoding="utf-8"))
    snippets = pq.read_table(cache / "snippets.parquet").column("code").to_pylist()
    indices = {}
    for name in meta["splits"]:
        z = np.load(cache / f"{name}.npz")
        indices[name] = SplitIndex(
            name=name,
            snippet_id1=z["snippet_id1"],
            snippet_id2=z["snippet_id2"],
            labels=z["labels"],
            group1=z["group1"],
            group2=z["group2"],
        )
    return snippets, indices, meta["stats"]


# --------------------------------------------------------------------------- #
# 3. Leakage-safe splitting
# --------------------------------------------------------------------------- #
def split_heldout_by_group(
    heldout: SplitIndex, test_group_fraction: float, seed: int
) -> tuple[SplitIndex, SplitIndex, dict[str, Any]]:
    """Partition the held-out split into validation/test along *group* boundaries.

    The dataset ships only ``train`` and ``val``; ``val`` is carved into a
    validation and a test half by assigning whole problem groups to one side.
    Pairs whose two snippets straddle the boundary are dropped -- keeping them
    would put the same group on both sides.
    """
    groups = np.unique(np.concatenate([heldout.group1, heldout.group2]))
    rng = np.random.default_rng(seed)
    shuffled = groups.copy()
    rng.shuffle(shuffled)
    n_test = max(1, int(round(len(shuffled) * test_group_fraction)))
    if n_test >= len(shuffled):
        raise ValueError("test_group_fraction leaves no groups for validation.")
    test_groups = set(shuffled[:n_test].tolist())
    val_groups = set(shuffled[n_test:].tolist())

    in_test = np.isin(heldout.group1, list(test_groups)) & np.isin(
        heldout.group2, list(test_groups)
    )
    in_val = np.isin(heldout.group1, list(val_groups)) & np.isin(heldout.group2, list(val_groups))
    dropped = int(len(heldout) - in_test.sum() - in_val.sum())

    validation = heldout.select(np.flatnonzero(in_val), name="validation")
    test = heldout.select(np.flatnonzero(in_test), name="test")

    overlap = set(validation.snippet_ids().tolist()) & set(test.snippet_ids().tolist())
    if overlap:
        raise AssertionError(f"{len(overlap)} snippets leaked between validation and test.")

    report = {
        "heldout_groups": int(len(groups)),
        "validation_groups": len(val_groups),
        "test_groups": len(test_groups),
        "dropped_cross_boundary_pairs": dropped,
        "validation": validation.class_distribution(),
        "test": test.class_distribution(),
    }
    return validation, test, report


def subsample(
    idx: SplitIndex, max_samples: int, seed: int, balanced: bool = True
) -> tuple[SplitIndex, dict[str, Any]]:
    """Take at most ``max_samples`` rows, optionally keeping the classes balanced.

    Subsampling never crosses group boundaries (it only removes rows), so it
    cannot introduce leakage.
    """
    if max_samples < 0 or max_samples >= len(idx):
        return idx, {"subsampled": False, "kept": len(idx)}

    rng = np.random.default_rng(seed)
    if balanced:
        per_class = max_samples // 2
        chosen = []
        for label in (0, 1):
            rows = np.flatnonzero(idx.labels == label)
            take = min(per_class, len(rows))
            chosen.append(rng.choice(rows, size=take, replace=False))
        rows = np.sort(np.concatenate(chosen))
    else:
        rows = np.sort(rng.choice(len(idx), size=max_samples, replace=False))

    out = idx.select(rows)
    return out, {"subsampled": True, "kept": len(out), "balanced": balanced}


def decide_class_weights(
    train: SplitIndex, mode: str, threshold: float
) -> tuple[list[float] | None, dict[str, Any]]:
    """Decide whether class-weighted cross entropy is warranted.

    Weighting is *not* applied by default: it is enabled only when the measured
    majority-class share exceeds ``threshold``. The rationale is recorded in the
    returned report and written to ``training_config.json``.
    """
    dist = train.class_distribution()
    n0, n1 = dist["negatives_label_0"], dist["positives_label_1"]
    total = max(n0 + n1, 1)
    majority_share = max(n0, n1) / total

    if mode == "off":
        apply = False
        reason = "class_weighting=off (forced by config)."
    elif mode == "on":
        apply = True
        reason = "class_weighting=on (forced by config)."
    else:
        apply = majority_share > threshold
        reason = (
            f"Measured majority-class share {majority_share:.4f} "
            f"{'exceeds' if apply else 'is within'} the {threshold} threshold, "
            f"so weighted cross entropy is {'enabled' if apply else 'NOT used'}."
        )

    weights = None
    if apply and n0 > 0 and n1 > 0:
        # Inverse-frequency weights normalised to mean 1.
        w = np.array([total / (2 * n0), total / (2 * n1)], dtype=np.float64)
        weights = (w / w.mean()).tolist()

    report = {
        "mode": mode,
        "threshold": threshold,
        "majority_class_share": round(majority_share, 6),
        "applied": weights is not None,
        "weights": weights,
        "reason": reason,
    }
    logger.info("Class weighting decision: %s", reason)
    return weights, report


# --------------------------------------------------------------------------- #
# 4. GraphCodeBERT snippet features
# --------------------------------------------------------------------------- #
@dataclass
class SnippetFeatures:
    """Pre-tokenised snippet pool, laid out as flat numpy arrays.

    ``dfg_adj_*`` store the (ragged) node adjacency lists so that no data-flow
    edge is ever clipped away silently.
    """

    input_ids: np.ndarray  # int32  [N, L]
    position_idx: np.ndarray  # int16  [N, L]
    dfg_to_code: np.ndarray  # int32  [N, max_nodes, 2]  (offsets into the
    #: *untruncated* sub-token stream, so they can exceed the sequence length)
    num_nodes: np.ndarray  # int16  [N]
    node_index: np.ndarray  # int16  [N]  number of real code tokens (incl. <s>/</s>)
    max_length: np.ndarray  # int16  [N]  code tokens + data-flow nodes
    dfg_adj_values: np.ndarray  # int16  [total_edges]
    dfg_adj_offsets: np.ndarray  # int64  [N, max_nodes + 1]
    seq_length: int
    stats: dict[str, Any]

    def __len__(self) -> int:
        return int(self.input_ids.shape[0])


def build_snippet_features(
    cfg: Config, snippets: list[str], tokenizer: Any, num_proc: int | None = None
) -> SnippetFeatures:
    """Run data-flow extraction + tokenisation over every distinct snippet.

    Uses ``datasets.map`` (batched, multi-process, Arrow-cached) so that a rerun
    with the same config costs nothing.
    """
    from datasets import Dataset as HFDataset

    seq_len = cfg.total_sequence_length
    max_nodes = seq_len - 3  # hard upper bound; the real cap is computed per snippet
    cache = Path(cfg.cache_dir) / "features"
    cache.mkdir(parents=True, exist_ok=True)
    key = _fingerprint(
        cfg.dataset_name, cfg.model_name_or_path, cfg.code_length, cfg.data_flow_length, len(snippets)
    )
    npz_path = cache / f"snippets_{key}.npz"
    stats_path = cache / f"snippets_{key}.stats.json"

    if npz_path.exists() and stats_path.exists():
        logger.info("Reusing cached snippet features at %s", npz_path)
        z = np.load(npz_path)
        return SnippetFeatures(
            input_ids=z["input_ids"],
            position_idx=z["position_idx"],
            dfg_to_code=z["dfg_to_code"],
            num_nodes=z["num_nodes"],
            node_index=z["node_index"],
            max_length=z["max_length"],
            dfg_adj_values=z["dfg_adj_values"],
            dfg_adj_offsets=z["dfg_adj_offsets"],
            seq_length=seq_len,
            stats=json.loads(stats_path.read_text(encoding="utf-8")),
        )

    ds = HFDataset.from_dict({"code": snippets})
    num_proc = num_proc if num_proc and num_proc > 1 else None
    t0 = time.time()
    ds = ds.map(
        _make_feature_fn(tokenizer, cfg.code_length, cfg.data_flow_length),
        batched=True,
        batch_size=256,
        num_proc=num_proc,
        remove_columns=["code"],
        desc="GraphCodeBERT data-flow + tokenisation",
    )

    n = len(ds)
    cols = ds.with_format(None)
    status_counter: Counter[str] = Counter()
    n_nodes_all: list[int] = []

    input_ids = np.full((n, seq_len), tokenizer.pad_token_id, dtype=np.int32)
    position_idx = np.full((n, seq_len), tokenizer.pad_token_id, dtype=np.int16)
    num_nodes = np.zeros(n, dtype=np.int16)
    node_index = np.zeros(n, dtype=np.int16)
    max_length = np.zeros(n, dtype=np.int16)
    dfg_to_code_list: list[list[list[int]]] = []
    adj_values: list[int] = []
    adj_offsets = np.zeros((n, max_nodes + 1), dtype=np.int64)

    observed_max_nodes = 0
    for i, row in enumerate(cols):
        ids = row["input_ids"]
        pos = row["position_idx"]
        input_ids[i, : len(ids)] = ids
        position_idx[i, : len(pos)] = pos
        d2c = row["dfg_to_code"]
        adj = row["dfg_to_dfg"]
        num_nodes[i] = len(d2c)
        observed_max_nodes = max(observed_max_nodes, len(d2c))
        node_index[i] = int(np.sum(np.asarray(pos) > 1))
        max_length[i] = int(np.sum(np.asarray(pos) != tokenizer.pad_token_id))
        dfg_to_code_list.append(d2c)
        base = len(adj_values)
        adj_offsets[i, 0] = base
        for j, nb in enumerate(adj):
            adj_values.extend(nb)
            adj_offsets[i, j + 1] = len(adj_values)
        adj_offsets[i, len(adj) + 1 :] = len(adj_values)
        status_counter.update(row["status_flags"])
        n_nodes_all.append(len(d2c))

    observed_max_nodes = max(observed_max_nodes, 1)
    # int32: these are offsets into the untruncated sub-token stream, which for a
    # pathologically long snippet reaches ~1e5 -- well past int16.
    dfg_to_code = np.zeros((n, observed_max_nodes, 2), dtype=np.int32)
    for i, d2c in enumerate(dfg_to_code_list):
        if d2c:
            dfg_to_code[i, : len(d2c)] = np.asarray(d2c, dtype=np.int32)
    adj_offsets = adj_offsets[:, : observed_max_nodes + 1]

    nodes_arr = np.asarray(n_nodes_all)
    stats = {
        "num_snippets": n,
        "extraction_seconds": round(time.time() - t0, 1),
        "status_counts": dict(status_counter),
        "snippets_with_empty_dataflow": int((nodes_arr == 0).sum()),
        "dataflow_nodes_mean": round(float(nodes_arr.mean()), 2),
        "dataflow_nodes_p50": int(np.percentile(nodes_arr, 50)),
        "dataflow_nodes_p95": int(np.percentile(nodes_arr, 95)),
        "dataflow_nodes_max": int(nodes_arr.max()),
        "code_tokens_mean": round(float(node_index.mean()), 2),
        "code_tokens_truncated": int((node_index >= cfg.code_length - 1).sum()),
        "total_dataflow_edges": len(adj_values),
        "sequence_length": seq_len,
    }
    logger.info("Snippet features: %s", json.dumps(stats, indent=2))
    if status_counter.get("dfg_failed", 0) or status_counter.get("dfg_recursion_limit", 0):
        logger.warning(
            "Data-flow extraction degraded for %d snippets (kept with an empty graph, not dropped).",
            status_counter.get("dfg_failed", 0) + status_counter.get("dfg_recursion_limit", 0),
        )

    features = SnippetFeatures(
        input_ids=input_ids,
        position_idx=position_idx,
        dfg_to_code=dfg_to_code,
        num_nodes=num_nodes,
        node_index=node_index,
        max_length=max_length,
        dfg_adj_values=np.asarray(adj_values, dtype=np.int16),
        dfg_adj_offsets=adj_offsets,
        seq_length=seq_len,
        stats=stats,
    )
    np.savez_compressed(
        npz_path,
        input_ids=features.input_ids,
        position_idx=features.position_idx,
        dfg_to_code=features.dfg_to_code,
        num_nodes=features.num_nodes,
        node_index=features.node_index,
        max_length=features.max_length,
        dfg_adj_values=features.dfg_adj_values,
        dfg_adj_offsets=features.dfg_adj_offsets,
    )
    stats_path.write_text(json.dumps(stats, indent=2), encoding="utf-8")
    return features


def _make_feature_fn(tokenizer: Any, code_length: int, data_flow_length: int):
    """Build the batched ``datasets.map`` function (picklable via closure)."""
    seq_len = code_length + data_flow_length
    cls_id, sep_id = tokenizer.cls_token_id, tokenizer.sep_token_id
    pad_id, unk_id = tokenizer.pad_token_id, tokenizer.unk_token_id

    def fn(batch: dict[str, list]) -> dict[str, list]:
        out_ids, out_pos, out_d2c, out_d2d, out_status = [], [], [], [], []
        for code in batch["code"]:
            flags: list[str] = []
            try:
                code_tokens, dfg, status = extract_dataflow(code, DATASET_LANGUAGE)
            except DataFlowExtractionError as exc:
                # Never drop the example: fall back to a plain tokenisation with
                # no data-flow component, and record the reason.
                logger.warning("Data-flow extraction failed, using empty graph: %s", exc)
                code_tokens, dfg = code.split(), []
                status = {"comment_strip": "n/a", "parse": "failed", "dfg": "failed", "error": str(exc)}
            for stage in ("comment_strip", "parse", "dfg"):
                if status[stage] not in ("ok", "n/a"):
                    flags.append(f"{stage}_{status[stage]}")
            if not flags:
                flags.append("ok")

            # GraphCodeBERT tokenises each code token separately; the '@ ' prefix
            # trick forces a word-boundary BPE split for non-initial tokens.
            sub_tokens = [
                tokenizer.tokenize("@ " + t)[1:] if i != 0 else tokenizer.tokenize(t)
                for i, t in enumerate(code_tokens)
            ]
            ori2cur = {-1: (0, 0)}
            for i in range(len(sub_tokens)):
                prev_end = ori2cur[i - 1][1]
                ori2cur[i] = (prev_end, prev_end + len(sub_tokens[i]))
            flat = [y for x in sub_tokens for y in x]

            # Reserve room for the data-flow nodes, then for <s>/</s>.
            keep = seq_len - 3 - min(len(dfg), data_flow_length)
            flat = flat[:keep][: code_length - 3]

            source_tokens = [tokenizer.cls_token] + flat + [tokenizer.sep_token]
            source_ids = tokenizer.convert_tokens_to_ids(source_tokens)
            # Code tokens get positions 2..; data-flow nodes get 0; padding gets
            # pad_token_id (1). This is what the model uses to tell them apart.
            position_idx = [i + pad_id + 1 for i in range(len(source_tokens))]

            dfg = dfg[: seq_len - len(source_tokens)]
            source_ids += [unk_id] * len(dfg)
            position_idx += [0] * len(dfg)
            padding = seq_len - len(source_ids)
            source_ids += [pad_id] * padding
            position_idx += [pad_id] * padding

            # Re-index edges so they point at node slots, not original token ids.
            reverse = {x[1]: i for i, x in enumerate(dfg)}
            dfg_to_dfg = [[reverse[i] for i in x[-1] if i in reverse] for x in dfg]
            dfg_to_code = [[ori2cur[x[1]][0] + 1, ori2cur[x[1]][1] + 1] for x in dfg]

            out_ids.append(source_ids)
            out_pos.append(position_idx)
            out_d2c.append(dfg_to_code)
            out_d2d.append(dfg_to_dfg)
            out_status.append(flags)

        return {
            "input_ids": out_ids,
            "position_idx": out_pos,
            "dfg_to_code": out_d2c,
            "dfg_to_dfg": out_d2d,
            "status_flags": out_status,
        }

    return fn


# --------------------------------------------------------------------------- #
# 5. Torch dataset + graph-guided attention-mask collator
# --------------------------------------------------------------------------- #
class ClonePairDataset(Dataset):
    """Pairs as ``(snippet_id1, snippet_id2, label)`` over a shared feature pool."""

    def __init__(self, index: SplitIndex, features: SnippetFeatures) -> None:
        self.index = index
        self.features = features

    def __len__(self) -> int:
        return len(self.index)

    def __getitem__(self, i: int) -> tuple[int, int, int]:
        return (
            int(self.index.snippet_id1[i]),
            int(self.index.snippet_id2[i]),
            int(self.index.labels[i]),
        )


def build_graph_attention_mask(features: SnippetFeatures, sid: int) -> np.ndarray:
    """Construct GraphCodeBERT's graph-guided masked attention for one snippet.

    Four rules, exactly as in the paper:

    1. code tokens attend to code tokens;
    2. the special tokens ``<s>``/``</s>`` attend to everything real;
    3. a data-flow node attends to (and is attended by) the code tokens it was
       identified from;
    4. a data-flow node attends to its adjacent nodes in the graph.
    """
    L = features.seq_length
    mask = np.zeros((L, L), dtype=bool)
    node_index = int(features.node_index[sid])
    max_length = int(features.max_length[sid])
    n_nodes = int(features.num_nodes[sid])

    # (1) sequence attends to sequence
    mask[:node_index, :node_index] = True
    # (2) special tokens attend to all real positions
    ids = features.input_ids[sid]
    for pos in np.flatnonzero((ids == 0) | (ids == 2)):
        if pos < node_index:
            mask[pos, :max_length] = True
    # (3) nodes <-> the code tokens they come from
    d2c = features.dfg_to_code[sid]
    for j in range(n_nodes):
        a, b = int(d2c[j, 0]), int(d2c[j, 1])
        if a < node_index and b < node_index:
            mask[j + node_index, a:b] = True
            mask[a:b, j + node_index] = True
    # (4) nodes <-> adjacent nodes
    offsets = features.dfg_adj_offsets[sid]
    for j in range(n_nodes):
        nbrs = features.dfg_adj_values[offsets[j] : offsets[j + 1]]
        for a in nbrs:
            if int(a) + node_index < L:
                mask[j + node_index, int(a) + node_index] = True
    return mask


@dataclass
class CloneCollator:
    """Collate pairs into the tensors ``GraphCodeBERTForCloneDetection`` expects."""

    features: SnippetFeatures

    def __call__(self, batch: list[tuple[int, int, int]]) -> dict[str, torch.Tensor]:
        ids1 = [b[0] for b in batch]
        ids2 = [b[1] for b in batch]
        labels = [b[2] for b in batch]
        f = self.features
        return {
            "input_ids_1": torch.from_numpy(f.input_ids[ids1].astype(np.int64)),
            "position_idx_1": torch.from_numpy(f.position_idx[ids1].astype(np.int64)),
            "attn_mask_1": torch.from_numpy(
                np.stack([build_graph_attention_mask(f, i) for i in ids1])
            ),
            "input_ids_2": torch.from_numpy(f.input_ids[ids2].astype(np.int64)),
            "position_idx_2": torch.from_numpy(f.position_idx[ids2].astype(np.int64)),
            "attn_mask_2": torch.from_numpy(
                np.stack([build_graph_attention_mask(f, i) for i in ids2])
            ),
            "labels": torch.tensor(labels, dtype=torch.long),
        }


# --------------------------------------------------------------------------- #
# 6. One-call pipeline used by train.py / evaluate.py
# --------------------------------------------------------------------------- #
@dataclass
class PreparedData:
    train: ClonePairDataset | None
    validation: ClonePairDataset
    test: ClonePairDataset
    features: SnippetFeatures
    collator: CloneCollator
    class_weights: list[float] | None
    report: dict[str, Any]


def prepare_data(cfg: Config, tokenizer: Any, with_train: bool = True) -> PreparedData:
    """Load, split, verify and featurise the dataset end to end."""
    schema = verify_dataset_schema(cfg.dataset_name)
    snippets, indices, pool_stats = build_snippet_pool(
        cfg, (cfg.train_split, cfg.heldout_split)
    )

    train_idx = indices[cfg.train_split]
    validation_idx, test_idx, split_report = split_heldout_by_group(
        indices[cfg.heldout_split], cfg.test_group_fraction, cfg.seed
    )

    train_idx, train_sub = subsample(
        train_idx, cfg.max_train_samples, cfg.seed, cfg.balance_subsamples
    )
    validation_idx, val_sub = subsample(
        validation_idx, cfg.max_eval_samples, cfg.seed, cfg.balance_subsamples
    )
    test_idx, test_sub = subsample(
        test_idx, cfg.max_test_samples, cfg.seed, cfg.balance_subsamples
    )

    _assert_no_leakage({"train": train_idx, "validation": validation_idx, "test": test_idx})

    class_weights, weight_report = decide_class_weights(
        train_idx, cfg.class_weighting, cfg.class_weight_threshold
    )

    features = build_snippet_features(cfg, snippets, tokenizer, cfg.preprocessing_num_workers)
    collator = CloneCollator(features)

    report = {
        "schema": schema,
        "snippet_pool": pool_stats,
        "heldout_split_strategy": split_report,
        "subsampling": {"train": train_sub, "validation": val_sub, "test": test_sub},
        "class_distribution": {
            "train": train_idx.class_distribution(),
            "validation": validation_idx.class_distribution(),
            "test": test_idx.class_distribution(),
        },
        "class_weighting": weight_report,
        "feature_extraction": features.stats,
    }

    return PreparedData(
        train=ClonePairDataset(train_idx, features) if with_train else None,
        validation=ClonePairDataset(validation_idx, features),
        test=ClonePairDataset(test_idx, features),
        features=features,
        collator=collator,
        class_weights=class_weights,
        report=report,
    )


def _assert_no_leakage(splits: dict[str, SplitIndex]) -> None:
    """Hard gate: no snippet and no group may appear in two splits."""
    names = list(splits)
    for i, a in enumerate(names):
        sa = set(splits[a].snippet_ids().tolist())
        ga = set(np.unique(np.concatenate([splits[a].group1, splits[a].group2])).tolist())
        for b in names[i + 1 :]:
            sb = set(splits[b].snippet_ids().tolist())
            gb = set(np.unique(np.concatenate([splits[b].group1, splits[b].group2])).tolist())
            if sa & sb:
                raise AssertionError(f"LEAK: {len(sa & sb)} snippets shared by {a} and {b}.")
            if ga & gb:
                raise AssertionError(f"LEAK: {len(ga & gb)} groups shared by {a} and {b}.")
    logger.info("Leakage check passed: splits share no snippet and no problem group.")


def iter_batches(dataset: Dataset, collator: CloneCollator, batch_size: int) -> Iterator[dict]:
    """Small helper for scripts that need batches without a Trainer."""
    for start in range(0, len(dataset), batch_size):
        yield collator([dataset[i] for i in range(start, min(start + batch_size, len(dataset)))])