Build provenance, typed rendering, safety-task leaks, answer-first MC fix, privacy-200k

scripts/build_jev_dataset.py

@@ -44,10 +44,12 @@
 }
 
 
-def load_jev_task(row, max_rows, max_rows_eval):
+def load_jev_task(row, max_rows, max_rows_eval, revision=None):
+    """``revision`` pins the task's Hub dataset to the commit recorded in the build report."""
     return load_task(
         row.id, recast="jev", multilingual=row.multilingual,
         max_rows=max_rows, max_rows_eval=max_rows_eval,
+        **({"revision": revision} if revision else {}),
     )
 
 
@@ -64,12 +66,12 @@ def load_jev_task(row, max_rows, max_rows_eval):
 )
 
 
-def load_native_task(source_id, max_rows, max_rows_eval):
+def load_native_task(source_id, max_rows, max_rows_eval, revision=None):
     """Sources authored as typed Jev questions, grouped by state (no recast)."""
     if source_id.startswith(graded.SOURCE_PREFIX):
-        rows = graded.load_family(source_id[len(graded.SOURCE_PREFIX):], max_rows, max_rows_eval)
+        rows = graded.load_family(source_id[len(graded.SOURCE_PREFIX):], max_rows, max_rows_eval, revision)
     else:
-        dataset = load_dataset(procedural.REPO_ID, source_id[len(procedural.SOURCE_PREFIX):])
+        dataset = load_dataset(procedural.REPO_ID, source_id[len(procedural.SOURCE_PREFIX):], revision=revision)
         dataset = sample_dataset(dataset, max_rows, max_rows_eval)
         rows = {split: [row for index, example in enumerate(examples)
                         for row in procedural.jev_rows(example, source_id, normalized_split(split), index)]
@@ -90,13 +92,42 @@ def slug(task_id):
     return f"{readable}-{digest}"
 
 
-def read_completed(report_path, data_dir=None):
-    completed = set()
-    if not report_path.exists():
-        return completed
-    for line in report_path.read_text().splitlines():
-        record = json.loads(line)
-        if record["status"] == "ok":
+# Arguments that change what a task's shards contain; with the code state they form the
+# build fingerprint, so a resumed build never mixes shards from different settings or code.
+SHARD_PARAMETERS = ("max_rows", "max_rows_eval", "noul_rate", "score_rate", "permutation_rate", "prompt_rate",
+                    "paired_format_rate", "pack_rate", "pack_max_tokens", "pack_tokenizer", "pack_max_items")
+
+
+def _git(*arguments):
+    root = Path(__file__).resolve().parents[1]
+    return subprocess.run(["git", *arguments], cwd=root, capture_output=True, check=False).stdout
+
+
+def code_state():
+    """Commit plus a hash of uncommitted changes under src/ and scripts/, untracked files included."""
+    root = Path(__file__).resolve().parents[1]
+    digest = hashlib.sha256(_git("diff", "--binary", "HEAD", "--", "src", "scripts"))
+    for name in sorted(_git("ls-files", "--others", "--exclude-standard", "--", "src", "scripts").decode().split()):
+        if not name.endswith((".pyc", ".pyo")):
+            digest.update(name.encode() + b"\0" + (root / name).read_bytes())
+    return {"git_commit": _git("rev-parse", "HEAD").decode().strip(), "uncommitted_sha256": digest.hexdigest()}
+
+
+def build_fingerprint(args, state=None):
+    shard_args = {name: getattr(args, name, None) for name in SHARD_PARAMETERS}
+    payload = json.dumps({**(state or code_state()), **shard_args}, sort_keys=True, default=str)
+    return hashlib.sha256(payload.encode()).hexdigest()[:16]
+
+
+def read_completed(report_path, data_dir=None, fingerprint=None):
+    """Tasks whose latest successful shards exist; with ``fingerprint``, return
+    ``(completed, stale)``: ``stale`` tasks have shards from another fingerprint."""
+    completed, stale = set(), set()
+    if report_path.exists():
+        for line in report_path.read_text().splitlines():
+            record = json.loads(line)
+            if record["status"] != "ok":
+                continue
             task = record["task"]
             splits = record.get("rows", {})
             if data_dir is None or (
@@ -105,8 +136,24 @@ def read_completed(report_path, data_dir=None):
                     for split in splits
                 )
             ):
-                completed.add(task)
-    return completed
+                current = fingerprint is None or record.get("fingerprint") == fingerprint
+                (completed if current else stale).add(task)
+                (stale if current else completed).discard(task)
+    return completed if fingerprint is None else (completed, stale)
+
+
+def resolve_revisions(provenance):
+    """Commit sha of each Hub repo a source loads, at the revision it requests."""
+    api, pins = HfApi(), {}
+    requested = {**({provenance["dataset"]: provenance.get("revision")} if provenance.get("dataset") else {}),
+                 **{repo: provenance.get("data_file_revisions", {}).get(repo)
+                    for repo in provenance.get("data_files_from", [])}}
+    for repo, revision in requested.items():
+        try:
+            pins[repo] = api.dataset_info(repo, revision=revision).sha
+        except Exception:  # gated or offline: loading will tell; the report records no pin
+            pins[repo] = None
+    return pins
 
 
 def build_manifest(args, tasks):
@@ -134,6 +181,7 @@ def git(*arguments):
     return {
         "git_commit": git("rev-parse", "HEAD").decode().strip(),
         "git_diff_sha256": hashlib.sha256(diff).hexdigest(),
+        "fingerprint": build_fingerprint(args),
         "working_tree_changes": status,
         "python": sys.version.split()[0],
         "packages": versions,
@@ -705,7 +753,7 @@ def publish_dataset(
         json.dumps(release_audit, indent=2, ensure_ascii=False) + "\n",
         encoding="utf-8",
     )
-    write_sources_yaml(output / "sources.yaml", release_audit)
+    write_sources_yaml(output / "sources.yaml", release_audit, latest_records(output / "build-report.jsonl"))
     api = HfApi()
     api.upload_file(
         path_or_fileobj=str(output / "README.md"),
@@ -738,28 +786,26 @@ def publish_dataset(
 def source_provenance(source):
     """Hub provenance of one published source (see tasksource.task_provenance)."""
     if source.startswith(procedural.SOURCE_PREFIX):
-        return {"generated": "procedural (tasksource.jev.procedural)"}
+        return {"dataset": procedural.REPO_ID, "config": source[len(procedural.SOURCE_PREFIX):],
+                "generated": "procedural (tasksource.jev.procedural)"}
     if source.startswith(graded.SOURCE_PREFIX):
         family = graded.FAMILIES[source[len(graded.SOURCE_PREFIX):]]
         info = {"dataset": family.dataset, "config": family.config,
+                "revision": (family.load_kwargs or {}).get("revision"),
                 "originals": sorted(set(ORIGINALS.get(family.dataset, [])) - {family.dataset})}
         return {k: v for k, v in info.items() if v}
     if source.startswith("multilingual/"):
         return task_provenance(source[len("multilingual/"):], multilingual=True)
     return task_provenance(source)
 
 
-def write_sources_yaml(path, release_audit):
-    """Every source in the release with its rows, Hub dataset, revision and originals."""
-    api, revisions = HfApi(), {}
+def write_sources_yaml(path, release_audit, records=None):
+    """Every source in the release with its rows, Hub dataset, revision and originals.
 
-    def revision(repo):
-        if repo not in revisions:
-            try:
-                revisions[repo] = api.dataset_info(repo).sha
-            except Exception:  # gated, renamed or offline: provenance without a pin
-                revisions[repo] = None
-        return revisions[repo]
+    Revisions are the commits each task was loaded from, recorded in the build report
+    (``records``, the latest report record per task); a task built before revisions were
+    recorded is marked ``revisions_unrecorded`` rather than given today's commit."""
+    records = records or {}
 
     splits = release_audit["splits"]
     sources = {}
@@ -769,9 +815,11 @@ def revision(repo):
             info.update(source_provenance(source))
         except KeyError as error:
             info["provenance_error"] = str(error)
-        for repo in [info.get("dataset"), *info.get("data_files_from", [])]:
-            if repo and revision(repo):
-                info.setdefault("revisions", {})[repo] = revision(repo)
+        recorded = (records.get(source) or {}).get("revisions")
+        if recorded:
+            info["revisions"] = {repo: sha for repo, sha in recorded.items() if sha}
+        else:
+            info["revisions_unrecorded"] = True
         sources[source] = info
     try:
         commit = subprocess.run(["git", "rev-parse", "HEAD"], capture_output=True, text=True,
@@ -877,16 +925,26 @@ def build(args):
             args.prompt_rate, args.paired_format_rate,
         )
     report_path = output / "build-report.jsonl"
-    completed = read_completed(report_path, data_dir)
+    fingerprint = build_fingerprint(args)
+    completed, stale = read_completed(report_path, data_dir, fingerprint)
     tasks = select_tasks(args)
+    stale &= set(tasks.source_id)
+    if stale and not (args.finalize_only or args.reuse_incompatible_shards):
+        raise SystemExit(
+            f"{len(stale)} task shards in {output} come from other code or settings (build fingerprint "
+            f"{fingerprint} differs), e.g. {sorted(stale)[:3]}. Use a fresh --output, or pass "
+            "--reuse-incompatible-shards to keep them (build-manifest.json records their fingerprints).")
+    if args.reuse_incompatible_shards:
+        completed |= stale
     packing_budget = LengthBudget(args.pack_max_tokens, args.pack_tokenizer)
     print(f"Selected {len(tasks)} tasks; {len(completed)} already complete", flush=True)
     manifest_path = output / "build-manifest.json"
-    if not manifest_path.exists():
-        manifest_path.write_text(
-            json.dumps(build_manifest(args, tasks), indent=2, ensure_ascii=False) + "\n",
-            encoding="utf-8",
-        )
+    manifest = build_manifest(args, tasks)
+    if manifest_path.exists():  # keep earlier runs' provenance next to this one's
+        previous = json.loads(manifest_path.read_text(encoding="utf-8"))
+        manifest["previous_runs"] = previous.pop("previous_runs", []) + [
+            {key: previous.get(key) for key in ("git_commit", "git_diff_sha256", "fingerprint", "parameters")}]
+    manifest_path.write_text(json.dumps(manifest, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
 
     for position, row in enumerate(
         () if args.finalize_only else tasks.itertuples(index=False), start=1
@@ -903,10 +961,14 @@ def build(args):
             if row.task_type == "TokenClassification":
                 max_rows = max(1, max_rows // 2)
                 max_rows_eval = max(1, max_rows_eval // 2)
+            # pin every Hub repo to the commit its requested revision points to now, and record it
+            provenance = source_provenance(task_id)
+            pins = resolve_revisions(provenance)
+            loaded = pins.get(provenance.get("dataset"))
             if row.task_type == "NativeJev":
-                dataset = load_native_task(task_id, max_rows, max_rows_eval)
+                dataset = load_native_task(task_id, max_rows, max_rows_eval, loaded)
             else:
-                dataset = load_jev_task(row, max_rows, max_rows_eval)
+                dataset = load_jev_task(row, max_rows, max_rows_eval, loaded)
             split_rows = {}
             for split, split_dataset in dataset.items():
                 if row.task_type == "NativeJev":
@@ -940,6 +1002,8 @@ def build(args):
                 "task_type": row.task_type,
                 "status": "ok",
                 "rows": split_rows,
+                "fingerprint": fingerprint,
+                "revisions": pins,
                 "seconds": round(time.time() - started, 3),
             }
             completed.add(task_id)
@@ -981,6 +1045,12 @@ def build(args):
             "outdated_datasets": len(outdated),
             "parquet_files": len(list(data_dir.glob("*.parquet"))),
             "excluded_source_families": list(PUBLISH_EXCLUDED_PREFIXES),
+            # the code/settings fingerprints the selected shards were built with; more than
+            # one means --reuse-incompatible-shards or --finalize-only kept older shards
+            "build_fingerprint": fingerprint,
+            "shard_fingerprints": dict(Counter(
+                latest[task].get("fingerprint") or "unrecorded" for task in sorted(selected & set(latest))
+                if latest[task]["status"] == "ok")),
         }
         (output / "build-summary.json").write_text(
             json.dumps(summary, indent=2) + "\n", encoding="utf-8"
@@ -1016,6 +1086,10 @@ def parse_args():
         "--english-only", action="store_true",
         help="Exclude the multilingual Tasksource catalog.",
     )
+    parser.add_argument(
+        "--reuse-incompatible-shards", action="store_true",
+        help="Resume even when existing shards come from other code or shard settings.",
+    )
     parser.add_argument("--max-rows", type=int, default=30_000)
     parser.add_argument("--max-rows-eval", type=int, default=3_000)
     parser.add_argument(

src/tasksource/access.py

@@ -1,6 +1,8 @@
 from .preprocess import Preprocessing, MultipleChoiceFields, add_question
 from .jev.options import JEV_MAX_MC_OPTIONS
 import re
+from urllib.parse import unquote
+import numpy as np
 import pandas as pd
 from . import tasks, recast as recast_module
 from .metadata import dataset_rank
@@ -44,8 +46,14 @@ def pretty_name(x):
     tn = x.task_name if x.task_name else ""
     return f"{dn}/{cn}/{tn}".replace('//','/').rstrip('/')
 
+def list_tasks(tasks_path=f'{os.path.dirname(__file__)}/tasks.py', multilingual=False, instruct=False, excluded=()):
+    """The task catalog as a DataFrame; ``excluded`` holds substrings of task ids to leave out.
+
+    Each call returns a fresh copy, so callers may edit it without affecting later calls."""
+    return _list_tasks(tasks_path, multilingual, instruct, tuple(excluded)).copy()
+
 @cache
-def list_tasks(tasks_path=f'{os.path.dirname(__file__)}/tasks.py',multilingual=False,instruct=False, excluded=[]):
+def _list_tasks(tasks_path, multilingual, instruct, excluded):
     if multilingual:
         tasks_path=tasks_path.replace('/tasks.py','/multilingual_tasks.py')
     task_order = open(tasks_path).readlines()
@@ -123,16 +131,18 @@ def task_provenance(task_id, multilingual=False):
     if row.empty:
         raise KeyError(f"unknown task: {task_id}")
     row = next(row.itertuples())
-    files = sorted(set(re.findall(r"hf://datasets/([\w.-]+/[\w.-]+)", str(row.mapping.load_dataset_kwargs))))
+    kwargs = row.mapping.load_dataset_kwargs or {}
+    # hf:// data files may pin a ref: hf://datasets/owner/name@refs%2Fconvert%2Fparquet/...
+    refs = {repo: unquote(ref) or None for repo, ref in
+            re.findall(r"hf://datasets/([\w.-]+/[\w.-]+)(?:@([^/\s'\"]+))?", str(kwargs))}
+    files = sorted(refs)
     dataset = None if row.dataset_name in RAW_BUILDERS else row.dataset_name
     originals = sorted({o for key in {dataset, row.id, *files} if key for o in ORIGINALS.get(key, [])} - {dataset})
-    info = {"dataset": dataset, "config": row.config_name or None, "data_files_from": files, "originals": originals}
+    info = {"dataset": dataset, "config": row.config_name or None, "revision": kwargs.get("revision"),
+            "data_files_from": files, "data_file_revisions": {repo: ref for repo, ref in refs.items() if ref},
+            "originals": originals}
     return {k: v for k, v in info.items() if v}
 
-def dict_to_query(d=dict(), **kwargs):
-    d={**d,**kwargs}
-    return '&'.join([f'`{k}`=="{v}"' for k,v in d.items()])
-
 def _format_loader_kwargs(value, **context):
     """Resolve per-config placeholders in generic-builder data_files settings."""
     if isinstance(value, str):
@@ -146,8 +156,13 @@ def _format_loader_kwargs(value, **context):
     return value
 
 def load_preprocessing(tasks=tasks, **kwargs):
-    _tasks_df = list_tasks(multilingual=tasks==lmtasks)
-    y = _tasks_df.copy().query(dict_to_query(**kwargs)).iloc[0]
+    df = list_tasks(multilingual=tasks==lmtasks)
+    matches = df[np.logical_and.reduce([df[k] == v for k, v in kwargs.items()] + [np.ones(len(df), bool)])]
+    if matches.empty:
+        raise KeyError(f"unknown task: {kwargs}")
+    if len(matches) > 1:
+        raise ValueError(f"{kwargs} matches {len(matches)} tasks ({', '.join(matches.id[:5])}...); pass a task id")
+    y = matches.iloc[0]
     preprocessing= copy.copy(getattr(tasks, y.preprocessing_name))
     for c in 'dataset_name','config_name':
         if not isinstance(getattr(preprocessing,c), str):

src/tasksource/jev/graded.py

@@ -10,7 +10,7 @@
 import math
 from dataclasses import dataclass
 
-from datasets import DatasetDict, load_dataset
+from datasets import Dataset, DatasetDict, load_dataset
 
 from ..preprocess import fix_splits, sample_dataset
 from ..tasks import render_dialogue, render_helpsteer_prompt
@@ -67,6 +67,7 @@ class Family:
     prepare: object = None  # row map adding the state columns
     dedupe: str = None  # column whose repeats are dropped (one row per annotator upstream)
     load_kwargs: dict = None  # e.g. the Hub parquet export of a script-only dataset
+    pre_process: object = None  # DatasetDict -> DatasetDict, e.g. one row per item from per-annotator rows
 
 
 HELPSTEER = dict(
@@ -134,6 +135,15 @@ def _dynasent_votes(example):
     return {"votes": [len(annotators.get(label, [])) for label in DYNASENT]}
 
 
+def _beavertails_votes(dataset):
+    """One row per (prompt, response) with the share of its three annotators who judged it unsafe."""
+    def votes(rows):
+        frame = rows.to_pandas()
+        frame["unsafe_share"] = frame.groupby(["prompt", "response"]).is_safe.transform(lambda safe: (~safe).sum() / len(safe))
+        return Dataset.from_pandas(frame.drop_duplicates(["prompt", "response"]), preserve_index=False)
+    return DatasetDict(train=votes(dataset["330k_train"]), test=votes(dataset["330k_test"]))
+
+
 def _share_of_yes(example):
     return {"share": example["soft_label"][1]}
 
@@ -198,6 +208,9 @@ def _lewidi_noul(config, state, question, prepare=_share_of_yes):
         disagreement=noul("disagreement_rate", f"How much would annotators disagree about {topic}, "
                                                "from 0 (all agree) to 1 (maximal disagreement)?")),
         dedupe="text") for name, topic in DISAGREEMENT.items()},
+    "beavertails": Family("PKU-Alignment/BeaverTails", {"Prompt": "prompt", "Response": "response"}, dict(
+        unsafe=noul("unsafe_share", "What fraction of annotators judged the assistant response unsafe?")),
+        pre_process=_beavertails_votes),
     "unli": Family("Zhengping/UNLI", PREMISE, dict(
         probability=noul("label", "How likely is the hypothesis to be true, given the premise?"))),
     **{f"dynasent_{r}": Family("dynabench/dynasent", {"Sentence": "sentence"}, dict(
@@ -263,10 +276,15 @@ def jev_rows(example, family, source, split, index):
     return rows
 
 
-def load_family(name, max_rows=None, max_rows_eval=None):
-    """Grouped Jev rows per split, split and sampled like any Tasksource source."""
+def load_family(name, max_rows=None, max_rows_eval=None, revision=None):
+    """Grouped Jev rows per split, split and sampled like any Tasksource source.
+
+    ``revision`` pins the source commit, overriding the family's requested ref."""
     family = FAMILIES[name]
-    dataset = DatasetDict(load_dataset(family.dataset, family.config, **(family.load_kwargs or {})))
+    kwargs = {**(family.load_kwargs or {}), **({"revision": revision} if revision else {})}
+    dataset = DatasetDict(load_dataset(family.dataset, family.config, **kwargs))
+    if family.pre_process:
+        dataset = family.pre_process(dataset)
     if family.dedupe:
         dataset = DatasetDict({split: rows.select(
             rows.to_pandas().drop_duplicates(family.dedupe).index.tolist()) for split, rows in dataset.items()})