Datasets

Author: jorge-martinez-gil

README.md

@@ -221,13 +221,14 @@ Currently automated sources:
 |---|---|---|
 | `bcb` | Hugging Face `google/code_x_glue_cc_clone_detection_big_clone_bench` | Stored directly as `pair_jsonl`. |
 | `poj104` | Hugging Face `google/code_x_glue_cc_clone_detection_poj104` | Official task is retrieval; downloader builds deterministic binary pairs for this repository. |
+| `poolc` | Hugging Face `PoolC/5-fold-clone-detection-600k-5fold` | Stored as `pair_jsonl`; the Hugging Face `val` split is deterministically divided into validation and test rows. |
 
 For POJ-104, use `--poj_pairs_per_label all` to materialize exhaustive
 positive pairs. The default samples 1000 positive pairs per label per split and
 the same number of negatives, which is much smaller and suitable for smoke
 tests or constrained runs.
 
-GCJ, Karnalim, PoolC, CodeNet, CLCDSA, SemanticCloneBench, and GPTCloneBench
+GCJ, Karnalim, CodeNet, CLCDSA, SemanticCloneBench, and GPTCloneBench
 still require manual source acquisition or conversion because no stable public
 direct-download endpoint is registered in this repository.
 
@@ -247,7 +248,7 @@ Currently supported local conversions:
 | `gcj` | `train.txt`, `valid.txt`, `test.txt`, plus `googlejam4_src/` files | Adds `data.jsonl`, converts `-1/1` labels to `0/1`, backs up raw splits as `raw_*.txt`. |
 | `karnalim` | `training.json`, `validation.json`, `test.json` | Adds `data.jsonl`, `train.txt`, `valid.txt`, `test.txt`. |
 
-PoolC, CodeNet, CLCDSA, SemanticCloneBench, and GPTCloneBench still need their
+CodeNet, CLCDSA, SemanticCloneBench, and GPTCloneBench still need their
 actual source files or official pair files before they can be normalized.
 
 ## Run Artifacts

docs/REPRODUCIBILITY.md

@@ -18,11 +18,14 @@ Automatically retrieve public datasets with stable sources first:
 python scripts/download_datasets.py --dataset all --output_root datasets --skip_existing
 ```
 
-This currently downloads CodeXGLUE BigCloneBench as `datasets/bcb` and
-CodeXGLUE POJ-104 as `datasets/poj104`. POJ-104 is officially a retrieval task;
-the downloader stores the snippets and creates deterministic binary pairs for
-this repository's pair-classification trainers. The source and conversion
-settings are written to `dataset_source.json`.
+This currently downloads CodeXGLUE BigCloneBench as `datasets/bcb`, CodeXGLUE
+POJ-104 as `datasets/poj104`, and PoolC as `datasets/poolc`. POJ-104 is
+officially a retrieval task; the downloader stores the snippets and creates
+deterministic binary pairs for this repository's pair-classification trainers.
+PoolC is downloaded from Hugging Face pair rows; the official `train` split is
+kept intact and the official `val` split is deterministically divided into
+validation and test rows. The source and conversion settings are written to
+`dataset_source.json`.
 
 Run with `--strict_data` when preparing final results. This fails fast on
 malformed rows, missing snippet IDs, and labels outside `{0, 1}`.

scripts/normalize_local_datasets.py

@@ -25,7 +25,6 @@
 
 SUPPORTED_DATASETS = ("gcj", "karnalim")
 MANUAL_DATASETS = {
-    "poolc": "Only data.txt is present locally; download/extract the PoolC files first.",
     "codenet": "Use scripts/prepare_pair_dataset.py on problem-directory source files.",
     "clcdsa": "Use scripts/prepare_pair_dataset.py on problem-directory source files.",
     "semanticclonebench": "Convert the official released pairs to pair_jsonl first.",

small_code_models/dataset_download.py

@@ -32,6 +32,11 @@
             "Code-Code/Clone-detection-POJ-104"
         ),
     },
+    "poolc": {
+        "display_name": "PoolC",
+        "source": "PoolC/5-fold-clone-detection-600k-5fold",
+        "homepage": "https://huggingface.co/datasets/PoolC/5-fold-clone-detection-600k-5fold",
+    },
 }
 
 DATASET_ALIASES = {
@@ -40,12 +45,12 @@
     "poj": "poj104",
     "poj-104": "poj104",
     "codexglue_poj104": "poj104",
+    "pool-c": "poolc",
 }
 
 MANUAL_DATASETS: dict[str, str] = {
     "gcj": "No stable public direct-download endpoint is registered here.",
     "karnalim": "No stable public direct-download endpoint is registered here.",
-    "poolc": "No stable public direct-download endpoint is registered here.",
     "codenet": "Use the official Project CodeNet release, then prepare subsets locally.",
     "clcdsa": "Use the official release, then prepare problem directories locally.",
     "semanticclonebench": "Convert the official released pairs to pair_jsonl locally.",
@@ -419,6 +424,156 @@ def download_poj104(
     return report
 
 
+def download_poolc(
+    output_dir: str | Path,
+    *,
+    overwrite: bool = False,
+    hf_cache_dir: str | Path | None = None,
+    include_diagnostics: bool = False,
+) -> dict[str, Any]:
+    """Download PoolC pair rows and normalize to pair_jsonl layout.
+
+    The Hugging Face release exposes ``train`` and ``val`` splits. This
+    repository requires train/validation/test files, so the converter keeps the
+    official training split intact and deterministically alternates rows from
+    ``val`` into validation and test.
+    """
+    destination = Path(output_dir)
+    _check_can_write(destination, overwrite)
+    load_dataset = _require_hf_datasets()
+    try:
+        from huggingface_hub import hf_hub_url, list_repo_files
+    except ImportError as exc:  # pragma: no cover - installed with datasets
+        raise RuntimeError(
+            "PoolC download requires huggingface_hub, which is installed with "
+            "the Hugging Face datasets package."
+        ) from exc
+
+    spec = AUTO_DATASETS["poolc"]
+    repo_files = list_repo_files(spec["source"], repo_type="dataset")
+    source_files = {
+        "train": sorted(
+            file_name
+            for file_name in repo_files
+            if file_name.startswith("data/train-") and file_name.endswith(".parquet")
+        ),
+        "val": sorted(
+            file_name
+            for file_name in repo_files
+            if file_name.startswith("data/val-") and file_name.endswith(".parquet")
+        ),
+    }
+    if not source_files["train"] or not source_files["val"]:
+        raise ValueError(
+            "PoolC source is missing expected data/train-*.parquet or "
+            "data/val-*.parquet files."
+        )
+    data_files = {
+        split_name: [
+            hf_hub_url(spec["source"], filename=file_name, repo_type="dataset")
+            for file_name in file_names
+        ]
+        for split_name, file_names in source_files.items()
+    }
+    split_rows = {"train": 0, "validation": 0, "test": 0}
+    source_split_rows = {"train": 0, "val": 0}
+    label_counts = {
+        "train": {"0": 0, "1": 0},
+        "validation": {"0": 0, "1": 0},
+        "test": {"0": 0, "1": 0},
+    }
+    snippet_hashes: dict[str, str] = {}
+    conflict_count = 0
+
+    def write_pair(
+        row: Mapping[str, Any],
+        split_name: str,
+        split_handle: Any,
+        data_handle: Any,
+    ) -> None:
+        nonlocal conflict_count
+        left_code = str(row["code1"])
+        right_code = str(row["code2"])
+        label = _label_to_int(row["similar"])
+        left_id = f"poolc:{_text_sha256(left_code)}"
+        right_id = f"poolc:{_text_sha256(right_code)}"
+
+        for snippet_id, code in ((left_id, left_code), (right_id, right_code)):
+            digest = _text_sha256(code)
+            previous_digest = snippet_hashes.get(snippet_id)
+            if previous_digest is None:
+                snippet_hashes[snippet_id] = digest
+                data_handle.write(json.dumps({"idx": snippet_id, "func": code}, sort_keys=True))
+                data_handle.write("\n")
+            elif previous_digest != digest:
+                conflict_count += 1
+
+        split_handle.write(f"{left_id}\t{right_id}\t{label}\n")
+        split_rows[split_name] += 1
+        label_counts[split_name][str(label)] += 1
+
+    load_kwargs = {"cache_dir": None if hf_cache_dir is None else str(hf_cache_dir)}
+    with (destination / "data.jsonl").open("w", encoding="utf-8") as data_handle:
+        train_dataset = load_dataset(
+            "parquet",
+            data_files={"train": data_files["train"]},
+            split="train",
+            streaming=True,
+            **load_kwargs,
+        )
+        with (destination / "train.txt").open("w", encoding="utf-8") as train_handle:
+            for row in train_dataset:
+                write_pair(row, "train", train_handle, data_handle)
+                source_split_rows["train"] += 1
+
+        val_dataset = load_dataset(
+            "parquet",
+            data_files={"val": data_files["val"]},
+            split="val",
+            streaming=True,
+            **load_kwargs,
+        )
+        with (destination / "valid.txt").open(
+            "w",
+            encoding="utf-8",
+        ) as validation_handle, (destination / "test.txt").open(
+            "w",
+            encoding="utf-8",
+        ) as test_handle:
+            for row_index, row in enumerate(val_dataset):
+                split_name = "validation" if row_index % 2 == 0 else "test"
+                split_handle = validation_handle if split_name == "validation" else test_handle
+                write_pair(row, split_name, split_handle, data_handle)
+                source_split_rows["val"] += 1
+
+    if conflict_count:
+        raise ValueError(
+            f"Encountered {conflict_count} conflicting PoolC snippet ids while "
+            "normalizing the dataset."
+        )
+
+    report = {
+        "dataset_key": "poolc",
+        "display_name": spec["display_name"],
+        "source": spec["source"],
+        "homepage": spec["homepage"],
+        "source_files": source_files,
+        "layout": "pair_jsonl",
+        "source_format": "hf_code_pair_rows",
+        "output_dir": str(destination),
+        "snippets": len(snippet_hashes),
+        "split_rows": split_rows,
+        "source_split_rows": source_split_rows,
+        "label_counts": label_counts,
+        "validation_test_source_split": "val",
+        "validation_test_strategy": "alternating_even_odd_rows",
+    }
+    if include_diagnostics:
+        report["diagnostics"] = inspect_dataset_directory(destination)
+    _write_json(destination / "dataset_source.json", report)
+    return report
+
+
 def download_dataset(
     dataset_key: str,
     output_root: str | Path,
@@ -453,4 +608,11 @@ def download_dataset(
             seed=seed,
             include_diagnostics=include_diagnostics,
         )
+    if key == "poolc":
+        return download_poolc(
+            output_dir,
+            overwrite=overwrite,
+            hf_cache_dir=hf_cache_dir,
+            include_diagnostics=include_diagnostics,
+        )
     raise KeyError(f"Dataset {dataset_key!r} is not registered for automatic download.")

tests/test_dataset_download.py

@@ -1,6 +1,9 @@
 """Unit tests for dataset download normalization helpers."""
 
+import json
 import random
+import sys
+import types
 from pathlib import Path
 
 import pytest
@@ -18,6 +21,7 @@
 def test_normalize_dataset_key_accepts_aliases() -> None:
     assert normalize_dataset_key("codexglue_bcb") == "bcb"
     assert normalize_dataset_key("poj-104") == "poj104"
+    assert normalize_dataset_key("pool-c") == "poolc"
 
 
 def test_parse_pairs_per_label() -> None:
@@ -107,3 +111,68 @@ def fake_load_dataset(source: str, *, split: str, cache_dir: str | None = None)
     assert (output_dir / "valid.txt").read_text() == "1\t3\t0\n"
     assert (output_dir / "test.txt").read_text() == "2\t3\t0\n"
     assert (output_dir / "dataset_source.json").exists()
+
+
+def test_download_poolc_writes_normalized_files(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    rows = {
+        "train": [
+            {"code1": "print(1)", "code2": "print(1)", "similar": 1},
+            {"code1": "print(2)", "code2": "print(3)", "similar": 0},
+        ],
+        "val": [
+            {"code1": "a = 1", "code2": "a=1", "similar": 1},
+            {"code1": "b = 1", "code2": "c = 1", "similar": 0},
+            {"code1": "d = 1", "code2": "d=1", "similar": 1},
+        ],
+    }
+
+    def fake_load_dataset(
+        source: str,
+        *,
+        data_files: dict[str, list[str]],
+        split: str,
+        cache_dir: str | None = None,
+        streaming: bool = False,
+    ) -> list:
+        assert source == "parquet"
+        assert split in data_files
+        assert cache_dir is None
+        assert streaming is True
+        return rows[split]
+
+    monkeypatch.setattr(dataset_download, "_require_hf_datasets", lambda: fake_load_dataset)
+    fake_hub = types.ModuleType("huggingface_hub")
+    fake_hub.hf_hub_url = lambda repo_id, filename, repo_type=None: (
+        f"https://example.test/{filename}"
+    )
+    fake_hub.list_repo_files = lambda repo_id, repo_type=None: [
+        "data/train-00000-of-00001.parquet",
+        "data/val-00000-of-00001.parquet",
+    ]
+    monkeypatch.setitem(sys.modules, "huggingface_hub", fake_hub)
+
+    report = download_dataset("poolc", tmp_path)
+
+    output_dir = tmp_path / "poolc"
+    snippets = [
+        json.loads(line)
+        for line in (output_dir / "data.jsonl").read_text().splitlines()
+    ]
+    train_rows = (output_dir / "train.txt").read_text().splitlines()
+    validation_rows = (output_dir / "valid.txt").read_text().splitlines()
+    test_rows = (output_dir / "test.txt").read_text().splitlines()
+
+    assert report["snippets"] == len(snippets)
+    assert len(train_rows) == 2
+    assert len(validation_rows) == 2
+    assert len(test_rows) == 1
+    assert train_rows[0].endswith("\t1")
+    assert train_rows[1].endswith("\t0")
+    assert validation_rows[0].endswith("\t1")
+    assert validation_rows[1].endswith("\t1")
+    assert test_rows[0].endswith("\t0")
+    assert report["validation_test_source_split"] == "val"
+    assert report["validation_test_strategy"] == "alternating_even_odd_rows"