[feat] support label resolve to anonymize labels

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 <to-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:
+
+- `<to-resolve>.resolved.json`: `{caseid: resolved label}` — agents can read this file directly.
+- `<to-resolve>.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
 

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,

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)
 

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,

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]})