"""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. /)
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 /.
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 ````/```` 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)))])