diff --git a/skills/prep/cls.md b/skills/prep/cls.md index 9f46fda..6f8a962 100644 --- a/skills/prep/cls.md +++ b/skills/prep/cls.md @@ -1,6 +1,6 @@ --- name: cls -description: Convert datasets to the cls format with the prep CLI — image-label classification records. Covers --labels and the image transform registry (--transforms or TRANSFORMS env var, e.g. crop_black_border). Use for classification data preparation with prep. +description: Convert datasets to the cls format with the prep CLI — image-label classification records. Covers --labels, label anonymization via --label-resolve, and the image transform registry (--transforms or TRANSFORMS env var, e.g. crop_black_border). Use for classification data preparation with prep. --- # prep `cls` format @@ -26,7 +26,16 @@ See the shared `prep` skill for source loading, conversion controls, and save op } ``` -`--labels` (or the `LABELS` env var) is required unless the label column is already a Hugging Face `ClassLabel` feature; it provides the class names used for casting. +`--labels` (or the `LABELS` env var) can override the `'label'` column; it should be exactly the list of class names. + +## Label resolution and anonymization + +`--label-resolve .json` (or the `LABEL_RESOLVE` env var) points to a JSON file keyed by `'label'` (as provided to `--labels`). During `cls` conversion, case names are anonymized in favor of case IDs: the `label` column stores only integer case IDs (the index of each case name in `--labels`), and class names are dropped from the saved feature. + +Two sidecar JSON files are written next to the resolve path: + +- `.resolved.json`: `{caseid: resolved label}` — agents can read this file directly. +- `.original.json`: `{caseid: casename}` — the caseid-to-casename mapping. Agents should avoid reading this file directly, since it holds the original, un-anonymized case names. ## Image transforms diff --git a/src/prep/api/formatter.py b/src/prep/api/formatter.py index 8f7ddfc..bad7e7f 100644 --- a/src/prep/api/formatter.py +++ b/src/prep/api/formatter.py @@ -1,5 +1,6 @@ from collections.abc import Callable from dataclasses import dataclass +from pathlib import Path from typing import TYPE_CHECKING from datasets import ClassLabel, Features, Image, List, Value @@ -65,6 +66,16 @@ def decorator(function: LoadFn) -> LoadFn: return decorator +def dump_mapping(mapping: list[str] | dict, fp: Path) -> None: + import json + + if isinstance(mapping, list): + mapping = dict(enumerate(mapping)) + + with open(fp, "w", encoding="utf-8") as f: + json.dump(mapping, f, ensure_ascii=False, indent=4) + + @dataclass(frozen=True) class FormatterPipeline: id_: str @@ -181,21 +192,39 @@ def cast_eval(d: "Dataset", args: ProcArgs) -> "Dataset": @staticmethod def cast_cls(d: "Dataset", args: ProcArgs) -> "Dataset": - if isinstance(d.features.get("label"), ClassLabel): - label_feature = d.features["label"] - elif args.labels: - label_feature = ClassLabel(names=args.labels) - else: + original_feat = d.features.get("label") + if not isinstance(original_feat, (ClassLabel, Value)): raise ValueError( - "Please provide labels via ENV `LABELS='cls1 cls2'`" - " or CLI option `--labels cls1 --labels cls2`." + f"Expected label feature to be ClassLabel or Value, got {type(original_feat)}" + ) + + candidate_feat = args.classlabel + if isinstance(original_feat, Value) and candidate_feat is None: + unique_labels = d.unique("label") + candidate_feat = ClassLabel(names=sorted(map(str, unique_labels))) + label_feat: ClassLabel = candidate_feat or original_feat + assert isinstance(label_feat, ClassLabel) + + if args.label_resolve is not None and args.classlabel_resolve is not None: + # if user provides a classlabel_resolve, we need to resolve it according to the label_feat + # and anonymize the label names to avoid leaking information about the original labels. + after_resolve = { + label_feat.str2int(k): v for k, v in args.classlabel_resolve.items() + } + # save this mapping to a file + dump_mapping( + after_resolve, args.label_resolve.with_suffix(".resolved.json") + ) + dump_mapping( + label_feat.names, args.label_resolve.with_suffix(".original.json") ) + label_feat = ClassLabel(num_classes=len(label_feat.names)) return d.cast( Features( id=Value("string"), image=Image(decode=True), - label=label_feature, + label=label_feat, extra_info=Value("large_string"), ), num_proc=args.num_proc, diff --git a/src/prep/api/types.py b/src/prep/api/types.py index 61cb30c..c8b9f77 100644 --- a/src/prep/api/types.py +++ b/src/prep/api/types.py @@ -1,7 +1,11 @@ +import json import logging from dataclasses import dataclass, field +from pathlib import Path from typing import Any, Literal, get_args +from datasets.features import ClassLabel + from ..constants import ( DEFAULT_ACOLS, DEFAULT_ATEMP, @@ -53,6 +57,7 @@ class ProcArgs: # image-to-image transforms applied in order during conversion labels: list[str] = field(default_factory=lambda: []) + label_resolve: Path | None = None transforms: list[str] = field(default_factory=lambda: []) # to fill in verl fields, this does not affect verl training @@ -82,6 +87,21 @@ def __post_init__(self): "question_template does not contain '{question}' placeholder." ) + @property + def classlabel(self) -> "ClassLabel | None": + if len(self.labels) == 0: + return None + + return ClassLabel(names=self.labels) + + @property + def classlabel_resolve(self) -> "dict[str, str] | None": + if self.label_resolve is None: + return None + + with open(self.label_resolve, encoding="utf-8") as f: + return json.load(f) + def peek(self, d: Any, level: int = logging.INFO): logger.log(level, d) diff --git a/src/prep/cli.py b/src/prep/cli.py index 79f8a73..f11a74f 100644 --- a/src/prep/cli.py +++ b/src/prep/cli.py @@ -57,34 +57,35 @@ def prep( hf_nproc: int | None = typer.Option(None, envvar="HF_NPROC", rich_help_panel=_HF), # auto conversion q_cols: list[str] = typer.Option( - default=DEFAULT_QCOLS, envvar="Q_COLS", rich_help_panel=_AUTO + DEFAULT_QCOLS, envvar="Q_COLS", rich_help_panel=_AUTO ), q_template: str = typer.Option( - default=DEFAULT_QTEMP, envvar="Q_TEMP", rich_help_panel=_AUTO + DEFAULT_QTEMP, envvar="Q_TEMP", rich_help_panel=_AUTO ), op_cols: list[str] = typer.Option( - default=DEFAULT_OPCOLS, envvar="OP_COLS", rich_help_panel=_AUTO + DEFAULT_OPCOLS, envvar="OP_COLS", rich_help_panel=_AUTO ), a_cols: list[str] = typer.Option( - default=DEFAULT_ACOLS, envvar="A_COLS", rich_help_panel=_AUTO + DEFAULT_ACOLS, envvar="A_COLS", rich_help_panel=_AUTO ), a_template: str = typer.Option( - default=DEFAULT_ATEMP, envvar="A_TEMP", rich_help_panel=_AUTO + DEFAULT_ATEMP, envvar="A_TEMP", rich_help_panel=_AUTO ), extra_info: str | None = typer.Option( - default=None, envvar="EXTRA_INFO", rich_help_panel=_AUTO + None, envvar="EXTRA_INFO", rich_help_panel=_AUTO ), # auto conversion (verl) verl_ability: str = typer.Option( - default="math", envvar="VERL_ABILITY", rich_help_panel=_VERL - ), - verl_style: str = typer.Option( - default="rule", envvar="VERL_STYLE", rich_help_panel=_VERL + "math", envvar="VERL_ABILITY", rich_help_panel=_VERL ), + verl_style: str = typer.Option("rule", envvar="VERL_STYLE", rich_help_panel=_VERL), # auto conversion (cls) - labels: list[str] = typer.Option(default=[], envvar="LABELS", rich_help_panel=_CLS), + labels: list[str] = typer.Option([], envvar="LABELS", rich_help_panel=_CLS), + label_resolve: Path | None = typer.Option( + None, envvar="LABEL_RESOLVE", rich_help_panel=_CLS + ), transforms: list[str] = typer.Option( - default=[], + [], envvar="TRANSFORMS", rich_help_panel=_CLS, help=f"Image transforms applied in order. Available: {list_transform_names()}", @@ -108,6 +109,7 @@ def prep( answer_template=a_template, extra_info=extra_info, labels=labels, + label_resolve=label_resolve, transforms=transforms, verl_ability=verl_ability, verl_style=verl_style, @@ -135,7 +137,7 @@ def prep( @app.command() def ppls( - save_root: Path = typer.Argument(default=Path("out"), envvar="SAVE_DIR"), + save_root: Path = typer.Argument(Path("out"), envvar="SAVE_DIR"), as_json: bool = False, filter_format: str = "*", list_format: bool = False, diff --git a/tests/api/test_formatter.py b/tests/api/test_formatter.py index e5891ba..a55a33c 100644 --- a/tests/api/test_formatter.py +++ b/tests/api/test_formatter.py @@ -268,8 +268,8 @@ def test_cast_cls_fails_without_a_label_schema_or_configured_labels(self): show_first_n=0, ) - with pytest.raises(ValueError, match="labels"): - FormatterPipeline.cast_cls(dataset, args_without_labels) + result = FormatterPipeline.cast_cls(dataset, args_without_labels) + assert result.features["label"] == ClassLabel(names=["airplane"]) def test_cast_dataset_returns_show_and_rejects_unknown_format(self): dataset = Dataset.from_dict({"value": [1]})