Add GraphCodeBERT clone-detection model

README.md

@@ -0,0 +1,184 @@
+---
+license: mit
+library_name: transformers
+pipeline_tag: text-classification
+tags:
+  - code
+  - clone-detection
+  - graphcodebert
+  - code-similarity
+base_model: microsoft/graphcodebert-base
+datasets:
+  - PoolC/1-fold-clone-detection-600k-5fold
+language:
+  - code
+metrics:
+  - accuracy
+  - precision
+  - recall
+  - f1
+model-index:
+  - name: graphcodebert-code-clone-detection
+    results:
+      - task:
+          type: text-classification
+          name: Binary code clone detection
+        dataset:
+          type: PoolC/1-fold-clone-detection-600k-5fold
+          name: PoolC/1-fold-clone-detection-600k-5fold
+          split: test (group-disjoint half of the `val` fold)
+        metrics:
+          - type: f1
+            value: 0.8805
+          - type: accuracy
+            value: 0.8747
+          - type: precision
+            value: 0.841
+          - type: recall
+            value: 0.924
+---
+
+# graphcodebert-code-clone-detection
+
+Binary code-clone detection. Full fine-tune of
+[`microsoft/graphcodebert-base`](https://huggingface.co/microsoft/graphcodebert-base) on
+[`PoolC/1-fold-clone-detection-600k-5fold`](https://huggingface.co/datasets/PoolC/1-fold-clone-detection-600k-5fold),
+using GraphCodeBERT's data-flow-aware pairwise architecture.
+
+Output labels: `0 = not clone`, `1 = clone`.
+
+## Architecture
+
+Not a generic sequence-pair classifier. The two snippets are encoded
+**separately** by one shared GraphCodeBERT encoder, each with its own
+graph-guided masked attention, and the two `<s>` vectors are concatenated for
+classification:
+
+```
+Linear(2 x 768 -> 768) -> tanh -> Linear(768 -> 2)
+```
+
+Per-snippet input layout (length 640):
+
+| segment | length | content | `position_idx` |
+|---|---|---|---|
+| code tokens | 512 | `<s>` + BPE code tokens + `</s>` | `2 .. n+1` |
+| data-flow nodes | 128 | one slot per DFG variable node (`<unk>` id) | `0` |
+| padding | remainder | `<pad>` | `1` |
+
+A data-flow node's input embedding is the **average of the embeddings of the
+code tokens it was identified from**. Graph-guided attention allows: code to
+code; `<s>`/`</s>` to everything; node to the code tokens it comes from (and
+back); node to adjacent nodes.
+
+## Preprocessing
+
+The dataset contains **Python** snippets, so data flow is extracted with the
+`tree-sitter-python` grammar via a port of GraphCodeBERT's `DFG_python`
+extractor (comment/docstring stripping -> AST -> variable states ->
+`comesFrom` / `computedFrom` edges).
+
+| | |
+|---|---|
+| `code_length` | 512 |
+| `data_flow_length` | 128 |
+| total sequence length | 640 |
+| distinct snippets featurised | 44,950 |
+| mean data-flow nodes / snippet | 44.22 |
+| snippets with empty data flow | 263 |
+| total data-flow edges | 2,481,388 |
+| extraction status counts | `{"ok": 44930, "comment_strip_failed": 13, "dfg_failed": 7}` |
+
+No example was dropped: a snippet whose data flow could not be extracted is
+kept with an empty graph and counted above.
+
+## Data splits
+
+The repository provides one of 5 predefined folds as `train` + `val`; those groups are disjoint and are kept as-is. `val` is partitioned further into validation/test along problem-group boundaries.
+
+`similar` equals `(code1_group == code2_group)` for every row, so the group
+columns are a perfect label proxy and are never used as features.
+
+| split | source | pairs | positives | negatives | groups |
+|---|---|---:|---:|---:|---:|
+| train | `train` fold | 50,000 | 25,000 | 25,000 | 240 |
+| validation | half of `val` by group | 20,000 | 10,000 | 10,000 | 29 |
+| test | other half of `val` by group | 20,000 | 10,000 | 10,000 | 30 |
+
+Train/validation/test share **no problem group and no code snippet**; this is
+asserted at runtime before training starts. 337,398 pairs of the
+held-out fold were dropped because their two snippets fell on opposite sides of
+the validation/test group boundary.
+
+Class weighting: Measured majority-class share 0.5000 is within the 0.6 threshold, so weighted cross entropy is NOT used.
+
+## Training
+
+| | |
+|---|---|
+| optimizer | adamw_torch |
+| learning rate | 2e-05 |
+| scheduler | linear with 0.1 warmup ratio (938 steps) |
+| epochs | 3.0 |
+| per-device batch size | 16 |
+| gradient accumulation | 1 |
+| effective batch size | 16 |
+| weight decay | 0.01 |
+| gradient clipping | 1.0 |
+| mixed precision | fp16 |
+| gradient checkpointing | False |
+| seed | 42 |
+| trainable parameters | 125,236,994 |
+| training time | 1.923 h |
+| GPU | NVIDIA GeForce RTX 5060 Ti (15.9 GB) |
+| torch / transformers | 2.11.0+cu128 / 5.17.0 |
+
+Checkpoint selection: best validation **F1** (`load_best_model_at_end=True`,
+`metric_for_best_model="f1"`). Best validation F1 = **0.8672**.
+The test split was scored once, after selection.
+
+## Results
+
+| split | accuracy | precision | recall | F1 | TP | TN | FP | FN |
+|---|---:|---:|---:|---:|---:|---:|---:|---:|
+| validation | 0.8557 | 0.8032 | 0.9422 | 0.8672 | 9,422 | 7,692 | 2,308 | 578 |
+| test | 0.8747 | 0.8410 | 0.9240 | 0.8805 | 9,240 | 8,253 | 1,747 | 760 |
+
+Test confusion matrix (`[[TN, FP], [FN, TP]]`): `[[8253, 1747], [760, 9240]]`
+
+## Usage
+
+This checkpoint uses a **custom pairwise head and a graph-guided attention
+mask**, so `AutoModelForSequenceClassification` will not reproduce these
+results. Use the repository's own model class and preprocessing:
+
+```python
+import torch
+from transformers import AutoTokenizer
+from modeling import load_model                      # from this project
+from preprocess import build_snippet_features, CloneCollator
+from config import Config
+
+cfg = Config()
+tokenizer = AutoTokenizer.from_pretrained("thealper2/graphcodebert-code-clone-detection")
+model = load_model("thealper2/graphcodebert-code-clone-detection").eval()
+
+features = build_snippet_features(cfg, [code_a, code_b], tokenizer, num_proc=1)
+collator = CloneCollator(features)
+batch = collator([(0, 1, 0)])                        # (snippet_a, snippet_b, dummy label)
+with torch.no_grad():
+    logits = model(**{k: v for k, v in batch.items() if k != "labels"}).logits
+label = int(logits.argmax(-1))                       # 0 = not clone, 1 = clone
+```
+
+## Limitations
+
+- Trained on competitive-programming Python solutions grouped by problem;
+  "clone" therefore means *solves the same problem*, which is closer to
+  semantic (Type-4) similarity than to syntactic copy-paste detection.
+- Data flow is extracted with the Python grammar only. Other languages need the
+  matching `tree_sitter_<lang>` grammar and `DFG_<lang>` function.
+- Snippets longer than 512 BPE tokens are truncated; 885 of
+  44,950 distinct snippets hit that limit.
+- Both directions of a pair are not explicitly symmetrised; the head sees
+  `concat(<s>_1, <s>_2)` in the given order.

code/config.py

@@ -0,0 +1,275 @@
+"""Central configuration for GraphCodeBERT binary code-clone detection.
+
+Every tunable lives in :class:`Config`.  ``Config.from_cli`` turns the dataclass
+fields into ``argparse`` flags automatically, so ``train.py`` / ``evaluate.py``
+never drift apart from this file.
+"""
+
+from __future__ import annotations
+
+import argparse
+import dataclasses
+import json
+from dataclasses import dataclass, field, fields
+from pathlib import Path
+from typing import Any
+
+# --------------------------------------------------------------------------- #
+# Facts discovered by inspecting the dataset (see README "Dataset structure").
+# They are kept here as documented defaults, *not* as blind assumptions:
+# preprocess.py re-verifies every one of them at runtime and raises if the
+# remote dataset ever changes.
+# --------------------------------------------------------------------------- #
+CODE1_COLUMN = "code1"
+CODE2_COLUMN = "code2"
+LABEL_COLUMN = "similar"
+GROUP1_COLUMN = "code1_group"
+GROUP2_COLUMN = "code2_group"
+#: Metadata columns that MUST NOT reach the model. ``code1_group``/``code2_group``
+#: determine the label exactly (``similar == (code1_group == code2_group)``), so
+#: feeding them in any form would be a 100 % label leak.
+FORBIDDEN_FEATURE_COLUMNS = (
+    GROUP1_COLUMN,
+    GROUP2_COLUMN,
+    "pair_id",
+    "question_pair_id",
+)
+#: The source language of the snippets, needed to pick the tree-sitter grammar.
+DATASET_LANGUAGE = "python"
+
+
+@dataclass
+class Config:
+    """All knobs for preprocessing, training and evaluation."""
+
+    # ---------------- dataset ---------------- #
+    dataset_name: str = "PoolC/1-fold-clone-detection-600k-5fold"
+    #: HF split that becomes the training set (group-disjoint from ``val``).
+    train_split: str = "train"
+    #: HF split that is partitioned by *group* into validation and test.
+    heldout_split: str = "val"
+    #: Fraction of the held-out split's groups reserved for the test set.
+    test_group_fraction: float = 0.5
+
+    #: Cap on the number of pairs per split. ``-1`` = use everything.
+    #: The full training split has 5.39 M pairs; see README for why the default
+    #: is a subsample and how to raise it.
+    max_train_samples: int = 50_000
+    max_eval_samples: int = 20_000
+    max_test_samples: int = 20_000
+    #: Keep the 50/50 label balance exactly when subsampling.
+    balance_subsamples: bool = True
+
+    # ---------------- model ---------------- #
+    model_name_or_path: str = "microsoft/graphcodebert-base"
+    #: GraphCodeBERT clone-detection defaults from the original paper/repo.
+    code_length: int = 512
+    data_flow_length: int = 128
+    attn_implementation: str = "sdpa"
+
+    # ---------------- training ---------------- #
+    learning_rate: float = 2e-5
+    num_train_epochs: float = 3.0
+    per_device_train_batch_size: int = 4
+    per_device_eval_batch_size: int = 4
+    gradient_accumulation_steps: int = 4
+    weight_decay: float = 0.01
+    warmup_ratio: float = 0.1
+    max_grad_norm: float = 1.0
+    fp16: bool = True
+    bf16: bool = False
+    gradient_checkpointing: bool = False
+    optim: str = "adamw_torch"
+    lr_scheduler_type: str = "linear"
+
+    #: Class-weighted cross entropy. ``"auto"`` enables it only when the
+    #: measured training distribution is more skewed than
+    #: ``class_weight_threshold``; ``"off"`` never, ``"on"`` always.
+    class_weighting: str = "auto"
+    class_weight_threshold: float = 0.6
+
+    # ---------------- evaluation / checkpointing ---------------- #
+    eval_strategy: str = "steps"
+    eval_steps: int = 1000
+    save_strategy: str = "steps"
+    save_steps: int = 1000
+    save_total_limit: int = 2
+    logging_steps: int = 100
+    metric_for_best_model: str = "f1"
+    greater_is_better: bool = True
+    load_best_model_at_end: bool = True
+
+    # ---------------- runtime ---------------- #
+    seed: int = 42
+    full_determinism: bool = False
+    dataloader_num_workers: int = 4
+    #: Processes used for tree-sitter data-flow extraction.
+    preprocessing_num_workers: int = 8
+    output_dir: str = "./outputs"
+    model_dir: str = "./models/graphcodebert-clone-detection"
+    logging_dir: str = "./logs"
+    cache_dir: str = "./outputs/feature_cache"
+    report_to: str = "none"
+    run_sanity_check: bool = True
+    sanity_check_samples: int = 64
+
+    # ---------------- derived ---------------- #
+    @property
+    def total_sequence_length(self) -> int:
+        """Length of one encoded snippet: code tokens + data-flow nodes."""
+        return self.code_length + self.data_flow_length
+
+    @property
+    def effective_batch_size(self) -> int:
+        return self.per_device_train_batch_size * self.gradient_accumulation_steps
+
+    # ---------------- (de)serialisation ---------------- #
+    def to_dict(self) -> dict[str, Any]:
+        d = dataclasses.asdict(self)
+        d["total_sequence_length"] = self.total_sequence_length
+        d["effective_batch_size"] = self.effective_batch_size
+        return d
+
+    def save(self, path: str | Path) -> None:
+        path = Path(path)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.write_text(json.dumps(self.to_dict(), indent=2), encoding="utf-8")
+
+    @classmethod
+    def from_json(cls, path: str | Path) -> "Config":
+        raw = json.loads(Path(path).read_text(encoding="utf-8"))
+        known = {f.name for f in fields(cls)}
+        return cls(**{k: v for k, v in raw.items() if k in known})
+
+    # ---------------- CLI ---------------- #
+    @classmethod
+    def build_parser(cls, description: str = "") -> argparse.ArgumentParser:
+        parser = argparse.ArgumentParser(
+            description=description,
+            formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+        )
+        parser.add_argument(
+            "--config_json",
+            type=str,
+            default=None,
+            help="Load defaults from a saved training_config.json, then apply CLI overrides.",
+        )
+        for f in fields(cls):
+            flag = f"--{f.name}"
+            if f.type is bool or f.type == "bool":
+                # Accept --flag / --flag true / --flag false.
+                parser.add_argument(
+                    flag,
+                    type=_str2bool,
+                    nargs="?",
+                    const=True,
+                    default=None,
+                    help=f"(bool) default: {f.default}",
+                )
+            else:
+                parser.add_argument(flag, type=type(f.default), default=None)
+        return parser
+
+    @classmethod
+    def from_cli(cls, argv: list[str] | None = None, description: str = "") -> "Config":
+        parser = cls.build_parser(description)
+        args, unknown = parser.parse_known_args(argv)
+        if unknown:
+            raise SystemExit(f"Unrecognised arguments: {unknown}")
+        cfg = cls.from_json(args.config_json) if args.config_json else cls()
+        for f in fields(cls):
+            value = getattr(args, f.name, None)
+            if value is not None:
+                setattr(cfg, f.name, value)
+        cfg.validate()
+        return cfg
+
+    def validate(self) -> None:
+        """Fail fast on impossible combinations instead of dying mid-training."""
+        if self.fp16 and self.bf16:
+            raise ValueError("Enable at most one of fp16 / bf16.")
+        if not 0.0 < self.test_group_fraction < 1.0:
+            raise ValueError("test_group_fraction must lie strictly between 0 and 1.")
+        if self.code_length <= 3:
+            raise ValueError("code_length must leave room for <s> and </s>.")
+        if self.data_flow_length < 0:
+            raise ValueError("data_flow_length must be >= 0.")
+        # GraphCodeBERT position ids run up to code_length + 1; RoBERTa's
+        # embedding table holds 514 slots (512 + <pad> + offset).
+        if self.code_length > 512:
+            raise ValueError(
+                "code_length > 512 exceeds GraphCodeBERT's position embeddings (514 slots)."
+            )
+        if self.class_weighting not in {"auto", "on", "off"}:
+            raise ValueError("class_weighting must be one of: auto, on, off.")
+        if self.metric_for_best_model not in {
+            "f1",
+            "accuracy",
+            "precision",
+            "recall",
+            "loss",
+        }:
+            raise ValueError(f"Unsupported metric_for_best_model: {self.metric_for_best_model}")
+        if self.load_best_model_at_end and self.eval_strategy != self.save_strategy:
+            raise ValueError("load_best_model_at_end requires eval_strategy == save_strategy.")
+        if (
+            self.load_best_model_at_end
+            and self.eval_strategy == "steps"
+            and self.save_steps % self.eval_steps != 0
+        ):
+            raise ValueError("save_steps must be a multiple of eval_steps.")
+
+
+#: Placeholder namespace in the Makefile default; replaced by the logged-in user.
+PLACEHOLDER_HUB_NAMESPACE = "your-username"
+DEFAULT_HUB_MODEL_NAME = "graphcodebert-clone-detection"
+
+
+def resolve_hub_repo_id(repo_id: str | None, token: str | None = None) -> str:
+    """Expand a bare model name into ``<namespace>/<name>`` for the Hub.
+
+    Accepts ``None``, a bare name, or a full ``user/name``. The namespace is
+    taken from the caller's Hugging Face credential (``huggingface-cli login``,
+    ``HF_TOKEN``, or an explicit ``token``), so the token never has to be typed
+    on the command line.
+
+    Raises:
+        RuntimeError: if no namespace is given and no credential is available.
+    """
+    from huggingface_hub import HfApi, get_token
+
+    name = (repo_id or DEFAULT_HUB_MODEL_NAME).strip().strip("/")
+    if "/" in name:
+        namespace, _, model_name = name.partition("/")
+        if namespace != PLACEHOLDER_HUB_NAMESPACE:
+            return f"{namespace}/{model_name}"
+        name = model_name or DEFAULT_HUB_MODEL_NAME
+
+    effective = token or get_token()
+    if not effective:
+        raise RuntimeError(
+            "No Hugging Face credential found, so the namespace for "
+            f"{name!r} cannot be resolved. Run `huggingface-cli login`, export "
+            "HF_TOKEN, or pass the full repo id as HF_REPO=user/name."
+        )
+    try:
+        who = HfApi().whoami(token=effective)
+    except Exception as exc:
+        raise RuntimeError(
+            f"Could not identify the logged-in Hugging Face account: {exc}. "
+            "Pass the full repo id as HF_REPO=user/name."
+        ) from exc
+    namespace = who.get("name")
+    if not namespace:
+        raise RuntimeError("Hugging Face account has no username; pass HF_REPO=user/name.")
+    return f"{namespace}/{name}"
+
+
+def _str2bool(value: str | bool) -> bool:
+    if isinstance(value, bool):
+        return value
+    if value.lower() in {"true", "t", "yes", "y", "1"}:
+        return True
+    if value.lower() in {"false", "f", "no", "n", "0"}:
+        return False
+    raise argparse.ArgumentTypeError(f"Expected a boolean, got {value!r}")

code/dfg_parser.py

@@ -0,0 +1,421 @@
+"""Data-flow graph (DFG) extraction for GraphCodeBERT.
+
+This is a faithful port of Microsoft's GraphCodeBERT ``parser/`` package
+(``utils.py`` + the ``DFG_python`` extractor from ``DFG.py``), adapted to the
+modern ``py-tree-sitter`` API (>= 0.22, ``Language(tree_sitter_python.language())``)
+instead of the original hand-compiled ``my-languages.so``.
+
+The dataset used in this project contains **Python** snippets (verified in
+``preprocess.py``), so only the Python extractor is ported; adding another
+language means adding its ``DFG_<lang>`` function and grammar package here.
+
+A DFG entry is the 5-tuple used throughout GraphCodeBERT::
+
+    (variable_name, token_index, edge_type, source_variable_names, source_token_indices)
+
+``edge_type`` is ``"comesFrom"`` (value flows from a previous definition) or
+``"computedFrom"`` (value is computed from the right-hand side of an assignment).
+"""
+
+from __future__ import annotations
+
+import io
+import re
+import sys
+import tokenize
+from typing import Any
+
+from tree_sitter import Language, Node, Parser
+
+__all__ = [
+    "get_parser",
+    "extract_dataflow",
+    "remove_comments_and_docstrings",
+    "DataFlowExtractionError",
+]
+
+#: tree-sitter recursion is mirrored by the recursive Python walkers below.
+#: Competitive-programming snippets can nest deeply, so raise the ceiling but
+#: keep it bounded so a pathological file raises RecursionError instead of
+#: segfaulting the worker.
+_RECURSION_LIMIT = 10_000
+
+
+class DataFlowExtractionError(RuntimeError):
+    """Raised when a snippet cannot be turned into code tokens at all."""
+
+
+_PARSER_CACHE: dict[str, Parser] = {}
+
+
+def get_parser(language: str = "python") -> Parser:
+    """Return a cached tree-sitter parser for ``language``.
+
+    Cached per process so that ``datasets.map(num_proc=...)`` workers each build
+    the parser once rather than once per snippet.
+    """
+    if language in _PARSER_CACHE:
+        return _PARSER_CACHE[language]
+    if language != "python":
+        raise ValueError(
+            f"Only the Python grammar is wired up (requested {language!r}). "
+            "Add the matching tree_sitter_<lang> package and DFG_<lang> function."
+        )
+    try:
+        import tree_sitter_python
+    except ImportError as exc:  # pragma: no cover - environment problem
+        raise ImportError(
+            "tree_sitter_python is required for GraphCodeBERT data-flow extraction. "
+            "Install it with `pip install tree-sitter tree-sitter-python`."
+        ) from exc
+    parser = Parser(Language(tree_sitter_python.language()))
+    _PARSER_CACHE[language] = parser
+    return parser
+
+
+# --------------------------------------------------------------------------- #
+# parser/utils.py
+# --------------------------------------------------------------------------- #
+def remove_comments_and_docstrings(source: str, lang: str = "python") -> str:
+    """Strip comments and docstrings, preserving token columns.
+
+    Column positions are preserved because the DFG indices are ``(row, column)``
+    points into the *cleaned* source.
+    """
+    if lang == "python":
+        io_obj = io.StringIO(source)
+        out = ""
+        prev_toktype = tokenize.INDENT
+        last_lineno = -1
+        last_col = 0
+        for tok in tokenize.generate_tokens(io_obj.readline):
+            token_type, token_string = tok[0], tok[1]
+            start_line, start_col = tok[2]
+            end_line, end_col = tok[3]
+            if start_line > last_lineno:
+                last_col = 0
+            if start_col > last_col:
+                out += " " * (start_col - last_col)
+            if token_type == tokenize.COMMENT:
+                pass
+            elif token_type == tokenize.STRING:
+                # A string that starts a logical line is a docstring -> drop it.
+                if prev_toktype != tokenize.INDENT and prev_toktype != tokenize.NEWLINE:
+                    if start_col > 0:
+                        out += token_string
+            else:
+                out += token_string
+            prev_toktype = token_type
+            last_col = end_col
+            last_lineno = end_line
+        return "\n".join(x for x in out.split("\n") if x.strip() != "")
+
+    def _replacer(match: re.Match[str]) -> str:
+        s = match.group(0)
+        return " " if s.startswith("/") else s
+
+    pattern = re.compile(
+        r"//.*?$|/\*.*?\*/|\'(?:\\.|[^\\\'])*\'|\"(?:\\.|[^\\\"])*\"",
+        re.DOTALL | re.MULTILINE,
+    )
+    cleaned = re.sub(pattern, _replacer, source)
+    return "\n".join(x for x in cleaned.split("\n") if x.strip() != "")
+
+
+def tree_to_token_index(root_node: Node) -> list[tuple[Any, Any]]:
+    """Collect ``(start_point, end_point)`` spans of every leaf token."""
+    if (len(root_node.children) == 0 or root_node.type == "string") and root_node.type != "comment":
+        return [(root_node.start_point, root_node.end_point)]
+    spans: list[tuple[Any, Any]] = []
+    for child in root_node.children:
+        spans += tree_to_token_index(child)
+    return spans
+
+
+def tree_to_variable_index(root_node: Node, index_to_code: dict) -> list[tuple[Any, Any]]:
+    """Collect spans of leaves that are *variables* (token text != node type)."""
+    if (len(root_node.children) == 0 or root_node.type == "string") and root_node.type != "comment":
+        index = (root_node.start_point, root_node.end_point)
+        _, code = index_to_code[index]
+        return [] if root_node.type == code else [index]
+    spans: list[tuple[Any, Any]] = []
+    for child in root_node.children:
+        spans += tree_to_variable_index(child, index_to_code)
+    return spans
+
+
+def index_to_code_token(index: tuple[Any, Any], code: list[str]) -> str:
+    """Slice the source text covered by a ``(start_point, end_point)`` span."""
+    start_point, end_point = index
+    if start_point[0] == end_point[0]:
+        return code[start_point[0]][start_point[1] : end_point[1]]
+    s = code[start_point[0]][start_point[1] :]
+    for i in range(start_point[0] + 1, end_point[0]):
+        s += code[i]
+    s += code[end_point[0]][: end_point[1]]
+    return s
+
+
+# --------------------------------------------------------------------------- #
+# parser/DFG.py :: DFG_python
+# --------------------------------------------------------------------------- #
+_ASSIGNMENT = ("assignment", "augmented_assignment", "for_in_clause")
+_IF_STATEMENT = ("if_statement",)
+_FOR_STATEMENT = ("for_statement",)
+_WHILE_STATEMENT = ("while_statement",)
+_DO_FIRST_STATEMENT = ("for_in_clause",)
+_DEF_STATEMENT = ("default_parameter",)
+
+
+def DFG_python(root_node: Node, index_to_code: dict, states: dict) -> tuple[list, dict]:
+    """Build the data-flow graph of a Python AST subtree.
+
+    Returns ``(dfg_edges, variable_states)`` where ``variable_states`` maps a
+    variable name to the token indices that currently define it.
+    """
+    states = states.copy()
+
+    if (len(root_node.children) == 0 or root_node.type == "string") and root_node.type != "comment":
+        idx, code = index_to_code[(root_node.start_point, root_node.end_point)]
+        if root_node.type == code:  # a keyword/operator, not a variable
+            return [], states
+        if code in states:
+            return [(code, idx, "comesFrom", [code], states[code].copy())], states
+        if root_node.type == "identifier":
+            states[code] = [idx]
+        return [(code, idx, "comesFrom", [], [])], states
+
+    if root_node.type in _DEF_STATEMENT:
+        name = root_node.child_by_field_name("name")
+        value = root_node.child_by_field_name("value")
+        dfg: list = []
+        if value is None:
+            for index in tree_to_variable_index(name, index_to_code):
+                idx, code = index_to_code[index]
+                dfg.append((code, idx, "comesFrom", [], []))
+                states[code] = [idx]
+            return sorted(dfg, key=lambda x: x[1]), states
+        name_indexs = tree_to_variable_index(name, index_to_code)
+        value_indexs = tree_to_variable_index(value, index_to_code)
+        temp, states = DFG_python(value, index_to_code, states)
+        dfg += temp
+        for index1 in name_indexs:
+            idx1, code1 = index_to_code[index1]
+            for index2 in value_indexs:
+                idx2, code2 = index_to_code[index2]
+                dfg.append((code1, idx1, "comesFrom", [code2], [idx2]))
+            states[code1] = [idx1]
+        return sorted(dfg, key=lambda x: x[1]), states
+
+    if root_node.type in _ASSIGNMENT:
+        if root_node.type == "for_in_clause":
+            right_nodes = [root_node.children[-1]]
+            left_nodes = [root_node.child_by_field_name("left")]
+        else:
+            if root_node.child_by_field_name("right") is None:
+                return [], states
+            left_nodes = [x for x in root_node.child_by_field_name("left").children if x.type != ","]
+            right_nodes = [
+                x for x in root_node.child_by_field_name("right").children if x.type != ","
+            ]
+            if len(right_nodes) != len(left_nodes):
+                left_nodes = [root_node.child_by_field_name("left")]
+                right_nodes = [root_node.child_by_field_name("right")]
+            if len(left_nodes) == 0:
+                left_nodes = [root_node.child_by_field_name("left")]
+            if len(right_nodes) == 0:
+                right_nodes = [root_node.child_by_field_name("right")]
+        dfg = []
+        for node in right_nodes:
+            temp, states = DFG_python(node, index_to_code, states)
+            dfg += temp
+        for left_node, right_node in zip(left_nodes, right_nodes):
+            left_tokens_index = tree_to_variable_index(left_node, index_to_code)
+            right_tokens_index = tree_to_variable_index(right_node, index_to_code)
+            for token1_index in left_tokens_index:
+                idx1, code1 = index_to_code[token1_index]
+                dfg.append(
+                    (
+                        code1,
+                        idx1,
+                        "computedFrom",
+                        [index_to_code[x][1] for x in right_tokens_index],
+                        [index_to_code[x][0] for x in right_tokens_index],
+                    )
+                )
+                states[code1] = [idx1]
+        return sorted(dfg, key=lambda x: x[1]), states
+
+    if root_node.type in _IF_STATEMENT:
+        dfg = []
+        current_states = states.copy()
+        others_states = []
+        tag = "else" in root_node.type
+        for child in root_node.children:
+            if "else" in child.type:
+                tag = True
+            if child.type not in ("elif_clause", "else_clause"):
+                temp, current_states = DFG_python(child, index_to_code, current_states)
+                dfg += temp
+            else:
+                temp, new_states = DFG_python(child, index_to_code, states)
+                dfg += temp
+                others_states.append(new_states)
+        others_states.append(current_states)
+        if tag is False:
+            others_states.append(states)
+        merged: dict = {}
+        for dic in others_states:
+            for key in dic:
+                merged.setdefault(key, [])
+                merged[key] += dic[key]
+        for key in merged:
+            merged[key] = sorted(set(merged[key]))
+        return sorted(dfg, key=lambda x: x[1]), merged
+
+    if root_node.type in _FOR_STATEMENT:
+        dfg = []
+        # Two passes: loop bodies can consume values defined later in the loop.
+        for _ in range(2):
+            right_nodes = [x for x in root_node.child_by_field_name("right").children if x.type != ","]
+            left_nodes = [x for x in root_node.child_by_field_name("left").children if x.type != ","]
+            if len(right_nodes) != len(left_nodes):
+                left_nodes = [root_node.child_by_field_name("left")]
+                right_nodes = [root_node.child_by_field_name("right")]
+            if len(left_nodes) == 0:
+                left_nodes = [root_node.child_by_field_name("left")]
+            if len(right_nodes) == 0:
+                right_nodes = [root_node.child_by_field_name("right")]
+            for node in right_nodes:
+                temp, states = DFG_python(node, index_to_code, states)
+                dfg += temp
+            for left_node, right_node in zip(left_nodes, right_nodes):
+                left_tokens_index = tree_to_variable_index(left_node, index_to_code)
+                right_tokens_index = tree_to_variable_index(right_node, index_to_code)
+                for token1_index in left_tokens_index:
+                    idx1, code1 = index_to_code[token1_index]
+                    dfg.append(
+                        (
+                            code1,
+                            idx1,
+                            "computedFrom",
+                            [index_to_code[x][1] for x in right_tokens_index],
+                            [index_to_code[x][0] for x in right_tokens_index],
+                        )
+                    )
+                    states[code1] = [idx1]
+            if root_node.children[-1].type == "block":
+                temp, states = DFG_python(root_node.children[-1], index_to_code, states)
+                dfg += temp
+        return _merge_duplicate_edges(dfg), states
+
+    if root_node.type in _WHILE_STATEMENT:
+        dfg = []
+        for _ in range(2):
+            for child in root_node.children:
+                temp, states = DFG_python(child, index_to_code, states)
+                dfg += temp
+        return _merge_duplicate_edges(dfg), states
+
+    dfg = []
+    for child in root_node.children:
+        if child.type in _DO_FIRST_STATEMENT:
+            temp, states = DFG_python(child, index_to_code, states)
+            dfg += temp
+    for child in root_node.children:
+        if child.type not in _DO_FIRST_STATEMENT:
+            temp, states = DFG_python(child, index_to_code, states)
+            dfg += temp
+    return sorted(dfg, key=lambda x: x[1]), states
+
+
+def _merge_duplicate_edges(dfg: list) -> list:
+    """Collapse the duplicate edges produced by the two-pass loop handling."""
+    dic: dict = {}
+    for x in dfg:
+        key = (x[0], x[1], x[2])
+        if key not in dic:
+            dic[key] = [x[3], x[4]]
+        else:
+            dic[key][0] = list(set(dic[key][0] + x[3]))
+            dic[key][1] = sorted(set(dic[key][1] + x[4]))
+    merged = [(k[0], k[1], k[2], v[0], v[1]) for k, v in sorted(dic.items(), key=lambda t: t[0][1])]
+    return sorted(merged, key=lambda x: x[1])
+
+
+# --------------------------------------------------------------------------- #
+# Public entry point (GraphCodeBERT's `extract_dataflow`)
+# --------------------------------------------------------------------------- #
+def extract_dataflow(code: str, language: str = "python") -> tuple[list[str], list, dict]:
+    """Tokenise ``code`` and extract its data-flow graph.
+
+    Returns ``(code_tokens, dfg, status)``. ``status`` records *why* a stage
+    degraded so callers can report it instead of hiding it:
+
+    ``comment_strip`` : ``"ok"`` | ``"failed"``
+    ``parse``         : ``"ok"`` | ``"failed"``
+    ``dfg``           : ``"ok"`` | ``"failed"`` | ``"recursion_limit"``
+    ``error``         : ``None`` or ``"<ExcType>: <message>"``
+
+    A degraded DFG yields an **empty** data-flow component -- the snippet is
+    still trained on (GraphCodeBERT tolerates zero nodes), it is never dropped.
+    """
+    status: dict[str, Any] = {"comment_strip": "ok", "parse": "ok", "dfg": "ok", "error": None}
+
+    try:
+        cleaned = remove_comments_and_docstrings(code, language)
+    except Exception as exc:
+        # Syntactically broken snippets are common in the wild; fall back to the
+        # raw source rather than discarding the example.
+        status["comment_strip"] = "failed"
+        status["error"] = f"{type(exc).__name__}: {exc}"
+        cleaned = code
+
+    parser = get_parser(language)
+    try:
+        tree = parser.parse(bytes(cleaned, "utf8"))
+        root_node = tree.root_node
+    except Exception as exc:
+        raise DataFlowExtractionError(f"tree-sitter failed to parse snippet: {exc}") from exc
+
+    old_limit = sys.getrecursionlimit()
+    sys.setrecursionlimit(_RECURSION_LIMIT)
+    try:
+        try:
+            tokens_index = tree_to_token_index(root_node)
+        except RecursionError as exc:
+            status["parse"] = "failed"
+            status["dfg"] = "recursion_limit"
+            status["error"] = f"{type(exc).__name__}: token index recursion limit"
+            raise DataFlowExtractionError("snippet nests deeper than the recursion limit") from exc
+
+        lines = cleaned.split("\n")
+        code_tokens = [index_to_code_token(x, lines) for x in tokens_index]
+        index_to_code = {
+            index: (idx, token) for idx, (index, token) in enumerate(zip(tokens_index, code_tokens))
+        }
+
+        try:
+            dfg, _ = DFG_python(root_node, index_to_code, {})
+        except RecursionError as exc:
+            status["dfg"] = "recursion_limit"
+            status["error"] = f"{type(exc).__name__}: DFG recursion limit"
+            dfg = []
+        except Exception as exc:
+            status["dfg"] = "failed"
+            status["error"] = f"{type(exc).__name__}: {exc}"
+            dfg = []
+    finally:
+        sys.setrecursionlimit(old_limit)
+
+    # Keep only nodes that participate in at least one edge (GraphCodeBERT does
+    # the same: isolated nodes carry no data-flow signal).
+    dfg = sorted(dfg, key=lambda x: x[1])
+    keep: set[int] = set()
+    for d in dfg:
+        if len(d[-1]) != 0:
+            keep.add(d[1])
+        keep.update(d[-1])
+    dfg = [d for d in dfg if d[1] in keep]
+
+    return code_tokens, dfg, status

code/modeling.py

@@ -0,0 +1,193 @@
+"""GraphCodeBERT clone-detection model.
+
+Reimplements the architecture from Microsoft's ``GraphCodeBERT/clonedetection``
+on top of ``transformers`` v5:
+
+* the two snippets are encoded **separately** by one shared GraphCodeBERT
+  encoder, each with its own graph-guided masked attention;
+* a data-flow node's input embedding is the average of the embeddings of the
+  code tokens it was identified from;
+* the two ``<s>`` representations are concatenated and fed to a
+  ``Linear(2H -> H) -> tanh -> Linear(H -> 2)`` head.
+
+The only real adaptation is the attention mask: ``transformers`` v5 builds masks
+through ``masking_utils`` and only forwards a mask untouched when it is already
+4-D, so the boolean ``[B, L, L]`` graph mask is expanded to an additive
+``[B, 1, L, L]`` mask here.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+from transformers import RobertaConfig, RobertaModel, RobertaPreTrainedModel
+from transformers.modeling_outputs import SequenceClassifierOutput
+
+__all__ = ["GraphCodeBERTForCloneDetection", "CloneClassificationHead"]
+
+
+def _autocast_dtype(device: torch.device, fallback: torch.dtype) -> torch.dtype:
+    """Dtype the attention scores will actually have, honouring autocast.
+
+    SDPA requires an additive ``attn_mask`` whose dtype matches the query, so a
+    hard-coded float32 mask would break under ``fp16=True``.
+    """
+    try:
+        if torch.is_autocast_enabled(device.type):
+            return torch.get_autocast_dtype(device.type)
+    except TypeError:  # older signature without a device argument
+        if device.type == "cuda" and torch.is_autocast_enabled():
+            return torch.get_autocast_gpu_dtype()
+    return fallback
+
+
+def _to_additive_mask(bool_mask: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
+    """``[B, L, L]`` boolean -> ``[B, 1, L, L]`` additive mask (0 / -inf)."""
+    additive = torch.zeros(bool_mask.shape, dtype=dtype, device=bool_mask.device)
+    additive.masked_fill_(~bool_mask, torch.finfo(dtype).min)
+    return additive.unsqueeze(1)
+
+
+class CloneClassificationHead(nn.Module):
+    """Pairwise head over the two ``<s>`` vectors (GraphCodeBERT's own head)."""
+
+    def __init__(self, config: RobertaConfig) -> None:
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size * 2, config.hidden_size)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+        self.out_proj = nn.Linear(config.hidden_size, 2)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """``hidden_states``: ``[B*2, L, H]`` -> logits ``[B, 2]``."""
+        x = hidden_states[:, 0, :]  # <s> of each snippet
+        x = x.reshape(-1, x.size(-1) * 2)  # pair the two snippets back up
+        x = self.dropout(x)
+        x = torch.tanh(self.dense(x))
+        x = self.dropout(x)
+        return self.out_proj(x)
+
+
+class GraphCodeBERTForCloneDetection(RobertaPreTrainedModel):
+    """Binary clone classifier: ``0 = not clone``, ``1 = clone``."""
+
+    config_class = RobertaConfig
+    base_model_prefix = "roberta"
+    supports_gradient_checkpointing = True
+
+    def __init__(self, config: RobertaConfig, class_weights: list[float] | None = None) -> None:
+        super().__init__(config)
+        config.num_labels = 2
+        self.roberta = RobertaModel(config, add_pooling_layer=False)
+        self.classifier = CloneClassificationHead(config)
+        self.register_buffer(
+            "class_weights",
+            torch.tensor(class_weights, dtype=torch.float32) if class_weights else None,
+            persistent=False,
+        )
+        self.post_init()
+
+    # ------------------------------------------------------------------ #
+    def _embed_with_dataflow(
+        self, input_ids: torch.Tensor, position_idx: torch.Tensor, attn_mask: torch.Tensor
+    ) -> torch.Tensor:
+        """Word embeddings where each data-flow node averages its code tokens.
+
+        ``position_idx`` encodes the role of every slot: ``0`` = data-flow node,
+        ``1`` (= ``<pad>``) = padding, ``>= 2`` = real code token.
+        """
+        nodes_mask = position_idx.eq(0)
+        token_mask = position_idx.ge(2)
+
+        embeddings = self.roberta.embeddings.word_embeddings(input_ids)
+        # For every node row, the code-token columns it may look at.
+        nodes_to_token = nodes_mask[:, :, None] & token_mask[:, None, :] & attn_mask
+        nodes_to_token = nodes_to_token.to(embeddings.dtype)
+        nodes_to_token = nodes_to_token / (nodes_to_token.sum(-1) + 1e-10)[:, :, None]
+        averaged = torch.einsum("abc,acd->abd", nodes_to_token, embeddings)
+        return embeddings * (~nodes_mask)[:, :, None] + averaged * nodes_mask[:, :, None]
+
+    def _encode(
+        self, input_ids: torch.Tensor, position_idx: torch.Tensor, attn_mask: torch.Tensor
+    ) -> torch.Tensor:
+        embeddings = self._embed_with_dataflow(input_ids, position_idx, attn_mask)
+        dtype = _autocast_dtype(input_ids.device, embeddings.dtype)
+        outputs = self.roberta(
+            inputs_embeds=embeddings,
+            attention_mask=_to_additive_mask(attn_mask, dtype),
+            position_ids=position_idx,
+            token_type_ids=torch.zeros_like(position_idx),
+        )
+        return outputs.last_hidden_state
+
+    # ------------------------------------------------------------------ #
+    def forward(
+        self,
+        input_ids_1: torch.Tensor,
+        position_idx_1: torch.Tensor,
+        attn_mask_1: torch.Tensor,
+        input_ids_2: torch.Tensor,
+        position_idx_2: torch.Tensor,
+        attn_mask_2: torch.Tensor,
+        labels: torch.Tensor | None = None,
+    ) -> SequenceClassifierOutput:
+        """Encode both snippets with the shared encoder and classify the pair.
+
+        Args:
+            input_ids_*: ``[B, L]`` token ids; data-flow slots hold ``<unk>``.
+            position_idx_*: ``[B, L]`` role/position ids (see ``_embed_with_dataflow``).
+            attn_mask_*: ``[B, L, L]`` boolean graph-guided attention mask.
+            labels: ``[B]`` with values in ``{0, 1}``.
+        """
+        batch_size, seq_len = input_ids_1.shape
+        # Stack both snippets into one encoder call: [B, L] x2 -> [B*2, L].
+        input_ids = torch.cat((input_ids_1[:, None], input_ids_2[:, None]), 1).view(-1, seq_len)
+        position_idx = torch.cat((position_idx_1[:, None], position_idx_2[:, None]), 1).view(
+            -1, seq_len
+        )
+        attn_mask = torch.cat((attn_mask_1[:, None], attn_mask_2[:, None]), 1).view(
+            -1, seq_len, seq_len
+        )
+
+        hidden = self._encode(input_ids, position_idx, attn_mask)
+        logits = self.classifier(hidden)
+
+        loss = None
+        if labels is not None:
+            weight = None
+            if self.class_weights is not None:
+                weight = self.class_weights.to(device=logits.device, dtype=logits.dtype)
+            loss = nn.functional.cross_entropy(logits, labels.view(-1), weight=weight)
+
+        return SequenceClassifierOutput(loss=loss, logits=logits)
+
+
+def load_model(
+    model_name_or_path: str,
+    attn_implementation: str = "sdpa",
+    class_weights: list[float] | None = None,
+    gradient_checkpointing: bool = False,
+) -> GraphCodeBERTForCloneDetection:
+    """Load GraphCodeBERT weights into the pairwise clone-detection head."""
+    model = GraphCodeBERTForCloneDetection.from_pretrained(
+        model_name_or_path,
+        attn_implementation=attn_implementation,
+    )
+    # Set after loading: `from_pretrained` should not have to carry runtime-only
+    # arguments, and the weights are a training artefact, not part of the config.
+    model.class_weights = (
+        torch.tensor(class_weights, dtype=torch.float32) if class_weights else None
+    )
+    if model.config.model_type != "roberta":
+        raise ValueError(
+            f"Expected a RoBERTa-architecture checkpoint (GraphCodeBERT), "
+            f"got model_type={model.config.model_type!r}."
+        )
+    if gradient_checkpointing:
+        model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
+    return model
+
+
+def count_parameters(model: nn.Module) -> dict[str, int]:
+    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    total = sum(p.numel() for p in model.parameters())
+    return {"trainable_parameters": trainable, "total_parameters": total}

code/preprocess.py

@@ -0,0 +1,839 @@
+"""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)))])

code/requirements.txt

@@ -0,0 +1,17 @@
+# Core
+torch>=2.4
+transformers>=5.0
+datasets>=3.0
+accelerate>=1.0
+tokenizers>=0.20
+
+# GraphCodeBERT data-flow extraction (tree-sitter is required by the original
+# GraphCodeBERT preprocessing pipeline -- it is not an optional extra).
+tree-sitter>=0.21
+tree-sitter-python>=0.21
+
+# Data / metrics
+numpy>=1.24
+pyarrow>=14
+scikit-learn>=1.3
+huggingface-hub>=0.26

config.json

@@ -0,0 +1,30 @@
+{
+  "add_cross_attention": false,
+  "architectures": [
+    "GraphCodeBERTForCloneDetection"
+  ],
+  "attention_probs_dropout_prob": 0.1,
+  "bos_token_id": 0,
+  "classifier_dropout": null,
+  "dtype": "float32",
+  "eos_token_id": 2,
+  "gradient_checkpointing": false,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 768,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "is_decoder": false,
+  "layer_norm_eps": 1e-05,
+  "max_position_embeddings": 514,
+  "model_type": "roberta",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "output_past": true,
+  "pad_token_id": 1,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.17.0",
+  "type_vocab_size": 1,
+  "use_cache": false,
+  "vocab_size": 50265
+}

experiment_record.json

@@ -0,0 +1,298 @@
+{
+  "model_name": "microsoft/graphcodebert-base",
+  "model_dir": "/mnt/d/work2/graphcodebert-code-clone-detection/models/graphcodebert-clone-detection",
+  "dataset_name": "PoolC/1-fold-clone-detection-600k-5fold",
+  "dataset_split_strategy": {
+    "train_split": "train",
+    "heldout_split": "val",
+    "note": "The repository provides one of 5 predefined folds as `train` + `val`; those groups are disjoint and are kept as-is. `val` is partitioned further into validation/test along problem-group boundaries.",
+    "heldout_groups": 59,
+    "validation_groups": 29,
+    "test_groups": 30,
+    "dropped_cross_boundary_pairs": 337398,
+    "validation": {
+      "num_examples": 483738,
+      "negatives_label_0": 157558,
+      "positives_label_1": 326180,
+      "positive_ratio": 0.674291,
+      "num_groups": 29
+    },
+    "test": {
+      "num_examples": 503224,
+      "negatives_label_0": 167224,
+      "positives_label_1": 336000,
+      "positive_ratio": 0.667695,
+      "num_groups": 30
+    }
+  },
+  "num_train_examples": 50000,
+  "num_validation_examples": 20000,
+  "num_test_examples": 20000,
+  "class_distributions": {
+    "train": {
+      "num_examples": 50000,
+      "negatives_label_0": 25000,
+      "positives_label_1": 25000,
+      "positive_ratio": 0.5,
+      "num_groups": 240
+    },
+    "validation": {
+      "num_examples": 20000,
+      "negatives_label_0": 10000,
+      "positives_label_1": 10000,
+      "positive_ratio": 0.5,
+      "num_groups": 29
+    },
+    "test": {
+      "num_examples": 20000,
+      "negatives_label_0": 10000,
+      "positives_label_1": 10000,
+      "positive_ratio": 0.5,
+      "num_groups": 30
+    }
+  },
+  "class_weighting": {
+    "mode": "auto",
+    "threshold": 0.6,
+    "majority_class_share": 0.5,
+    "applied": false,
+    "weights": null,
+    "reason": "Measured majority-class share 0.5000 is within the 0.6 threshold, so weighted cross entropy is NOT used."
+  },
+  "sequence_length": 512,
+  "data_flow_length": 128,
+  "total_sequence_length": 640,
+  "per_device_train_batch_size": 16,
+  "gradient_accumulation_steps": 1,
+  "effective_batch_size": 16,
+  "learning_rate": 2e-05,
+  "num_train_epochs": 3.0,
+  "optimizer": "adamw_torch",
+  "scheduler": "linear",
+  "warmup_ratio": 0.1,
+  "warmup_steps": 938,
+  "weight_decay": 0.01,
+  "max_grad_norm": 1.0,
+  "mixed_precision": "fp16",
+  "gradient_checkpointing": false,
+  "seed": 42,
+  "training_time_seconds": 6921.9,
+  "training_time_hours": 1.923,
+  "train_runtime_metrics": {
+    "train_runtime": 6918.8074,
+    "train_samples_per_second": 21.68,
+    "train_steps_per_second": 1.355,
+    "total_flos": 0.0,
+    "train_loss": 0.43145505716959637,
+    "epoch": 3.0
+  },
+  "gpu": {
+    "cuda_available": true,
+    "torch_version": "2.11.0+cu128",
+    "transformers_version": "5.17.0",
+    "python_version": "3.12.3",
+    "platform": "Linux-6.18.33.2-microsoft-standard-WSL2-x86_64-with-glibc2.39",
+    "gpu_name": "NVIDIA GeForce RTX 5060 Ti",
+    "gpu_count": 1,
+    "gpu_total_memory_gb": 15.9,
+    "gpu_capability": "12.0",
+    "cuda_version": "12.8"
+  },
+  "parameters": {
+    "trainable_parameters": 125236994,
+    "total_parameters": 125236994
+  },
+  "best_validation_f1": 0.8671882190520018,
+  "best_checkpoint": "./outputs/checkpoint-9000",
+  "validation_metrics": {
+    "loss": 0.34459105134010315,
+    "accuracy": 0.8557,
+    "precision": 0.8032395566922421,
+    "recall": 0.9422,
+    "f1": 0.8671882190520018,
+    "macro_f1": 0.8546121719233737,
+    "tp": 9422,
+    "tn": 7692,
+    "fp": 2308,
+    "fn": 578,
+    "confusion_matrix": [
+      [
+        7692,
+        2308
+      ],
+      [
+        578,
+        9422
+      ]
+    ],
+    "confusion_matrix_layout": "[[TN, FP], [FN, TP]]",
+    "support": {
+      "num_examples": 20000,
+      "label_0": 10000,
+      "label_1": 10000
+    },
+    "runtime": 173.9737,
+    "samples_per_second": 114.96,
+    "steps_per_second": 3.592,
+    "num_examples": 20000,
+    "eval_seconds": 174.0
+  },
+  "test_metrics": {
+    "loss": 0.3165612816810608,
+    "model_preparation_time": 0.0022,
+    "accuracy": 0.87465,
+    "precision": 0.8409939018840448,
+    "recall": 0.924,
+    "f1": 0.8805450993472149,
+    "macro_f1": 0.874343974488208,
+    "tp": 9240,
+    "tn": 8253,
+    "fp": 1747,
+    "fn": 760,
+    "confusion_matrix": [
+      [
+        8253,
+        1747
+      ],
+      [
+        760,
+        9240
+      ]
+    ],
+    "confusion_matrix_layout": "[[TN, FP], [FN, TP]]",
+    "support": {
+      "num_examples": 20000,
+      "label_0": 10000,
+      "label_1": 10000
+    },
+    "runtime": 181.6455,
+    "samples_per_second": 110.105,
+    "steps_per_second": 3.441,
+    "num_examples": 20000,
+    "eval_seconds": 181.7
+  },
+  "preprocessing": {
+    "snippet_pool": {
+      "num_unique_snippets": 44950,
+      "scan_seconds": 46.3,
+      "snippets_with_multiple_groups": 1,
+      "cross_split_snippet_overlap": {
+        "train|val:snippets": 0,
+        "train|val:groups": 0
+      },
+      "per_split": {
+        "train": {
+          "num_examples": 5388622,
+          "negatives_label_0": 2694311,
+          "positives_label_1": 2694311,
+          "positive_ratio": 0.5,
+          "num_groups": 240
+        },
+        "val": {
+          "num_examples": 1324360,
+          "negatives_label_0": 662180,
+          "positives_label_1": 662180,
+          "positive_ratio": 0.5,
+          "num_groups": 59
+        }
+      }
+    },
+    "feature_extraction": {
+      "num_snippets": 44950,
+      "extraction_seconds": 64.1,
+      "status_counts": {
+        "ok": 44930,
+        "comment_strip_failed": 13,
+        "dfg_failed": 7
+      },
+      "snippets_with_empty_dataflow": 263,
+      "dataflow_nodes_mean": 44.22,
+      "dataflow_nodes_p50": 33,
+      "dataflow_nodes_p95": 127,
+      "dataflow_nodes_max": 193,
+      "code_tokens_mean": 137.9,
+      "code_tokens_truncated": 885,
+      "total_dataflow_edges": 2481388,
+      "sequence_length": 640
+    },
+    "subsampling": {
+      "train": {
+        "subsampled": true,
+        "kept": 50000,
+        "balanced": true
+      },
+      "validation": {
+        "subsampled": true,
+        "kept": 20000,
+        "balanced": true
+      },
+      "test": {
+        "subsampled": true,
+        "kept": 20000,
+        "balanced": true
+      }
+    }
+  },
+  "sanity_check": {
+    "passed": true,
+    "device": "cuda",
+    "checks": [
+      {
+        "check": "1_dataset_loading",
+        "passed": true,
+        "detail": "train=50000 val=20000 test=20000"
+      },
+      {
+        "check": "2_column_detection",
+        "passed": true,
+        "detail": "code columns=code1/code2, label=similar; excluded from features: ['code1_group', 'code2_group', 'pair_id', 'question_pair_id']"
+      },
+      {
+        "check": "3_label_correctness",
+        "passed": true,
+        "detail": "labels in {0,1}; validation positive ratio=0.5"
+      },
+      {
+        "check": "4_tokenisation",
+        "passed": true,
+        "detail": "<s>...</s> wrapping OK, 93 code tokens, decodes to 'import math a , b , c = map ( int , input ( ) . split ( ) ) '"
+      },
+      {
+        "check": "5_dataflow_extraction",
+        "passed": true,
+        "detail": "44687/44950 snippets have a non-empty data-flow graph, 2481388 edges, mean nodes=44.22"
+      },
+      {
+        "check": "6_attention_mask",
+        "passed": true,
+        "detail": "shape=(640, 640) (code_length 512 + data_flow_length 128), 93 code tokens, 30 nodes, density=0.0215, padding rows empty"
+      },
+      {
+        "check": "7_forward_pass",
+        "passed": true,
+        "detail": "logits shape=(16, 2) for batch of 16 pairs; mask tensor=(16, 640, 640)"
+      },
+      {
+        "check": "8_finite_loss",
+        "passed": true,
+        "detail": "loss=0.6978 (chance level ~0.6931)"
+      },
+      {
+        "check": "9_training_step",
+        "passed": true,
+        "detail": "loss=0.6798, grad_norm=2.4929, weights updated (amp=on, dtype=torch.float16)"
+      }
+    ],
+    "parameters": {
+      "trainable_parameters": 125236994,
+      "total_parameters": 125236994
+    },
+    "sequence_length": 640
+  },
+  "test_metrics_provenance": {
+    "note": "Re-scored after fixing a transformers v5 dataloader-caching bug: Trainer.get_eval_dataloader caches under the key 'eval' when a Dataset object is passed and dataloader_persistent_workers=True, so the original in-training test evaluation silently re-scored the validation split. evaluate_split now uses Trainer.predict and verifies the returned labels.",
+    "superseded_test_f1": 0.8671882190520018,
+    "corrected_test_f1": 0.8805450993472149,
+    "validation_f1_unchanged": 0.8671882190520018
+  }
+}

model.safetensors

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

tokenizer_config.json

@@ -0,0 +1,17 @@
+{
+  "add_prefix_space": false,
+  "backend": "tokenizers",
+  "bos_token": "<s>",
+  "cls_token": "<s>",
+  "eos_token": "</s>",
+  "errors": "replace",
+  "is_local": false,
+  "local_files_only": false,
+  "mask_token": "<mask>",
+  "model_max_length": 512,
+  "pad_token": "<pad>",
+  "sep_token": "</s>",
+  "tokenizer_class": "RobertaTokenizer",
+  "trim_offsets": true,
+  "unk_token": "<unk>"
+}

training_args.bin

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

training_config.json

@@ -0,0 +1,51 @@
+{
+  "dataset_name": "PoolC/1-fold-clone-detection-600k-5fold",
+  "train_split": "train",
+  "heldout_split": "val",
+  "test_group_fraction": 0.5,
+  "max_train_samples": 50000,
+  "max_eval_samples": 20000,
+  "max_test_samples": 20000,
+  "balance_subsamples": true,
+  "model_name_or_path": "microsoft/graphcodebert-base",
+  "code_length": 512,
+  "data_flow_length": 128,
+  "attn_implementation": "sdpa",
+  "learning_rate": 2e-05,
+  "num_train_epochs": 3.0,
+  "per_device_train_batch_size": 16,
+  "per_device_eval_batch_size": 32,
+  "gradient_accumulation_steps": 1,
+  "weight_decay": 0.01,
+  "warmup_ratio": 0.1,
+  "max_grad_norm": 1.0,
+  "fp16": true,
+  "bf16": false,
+  "gradient_checkpointing": false,
+  "optim": "adamw_torch",
+  "lr_scheduler_type": "linear",
+  "class_weighting": "auto",
+  "class_weight_threshold": 0.6,
+  "eval_strategy": "steps",
+  "eval_steps": 1000,
+  "save_strategy": "steps",
+  "save_steps": 1000,
+  "save_total_limit": 2,
+  "logging_steps": 100,
+  "metric_for_best_model": "f1",
+  "greater_is_better": true,
+  "load_best_model_at_end": true,
+  "seed": 42,
+  "full_determinism": false,
+  "dataloader_num_workers": 4,
+  "preprocessing_num_workers": 8,
+  "output_dir": "./outputs",
+  "model_dir": "./models/graphcodebert-clone-detection",
+  "logging_dir": "./logs",
+  "cache_dir": "./outputs/feature_cache",
+  "report_to": "none",
+  "run_sanity_check": true,
+  "sanity_check_samples": 64,
+  "total_sequence_length": 640,
+  "effective_batch_size": 16
+}