[feat] support extra-info columns

src/prep/api/types.py

@@ -48,9 +48,11 @@ class ProcArgs:
 
     answer_cols: list[str] = field(default_factory=lambda: DEFAULT_ACOLS)
     answer_template: str = DEFAULT_ATEMP
-    labels: list[str] = field(default_factory=lambda: [])
+
+    extra_info: str | None = None
 
     # image-to-image transforms applied in order during conversion
+    labels: list[str] = field(default_factory=lambda: [])
     transforms: list[str] = field(default_factory=lambda: [])
 
     # to fill in verl fields, this does not affect verl training

src/prep/cli.py

@@ -79,6 +79,9 @@ def prep(
     a_template: str = typer.Option(
         default=DEFAULT_ATEMP, envvar="A_TEMP", rich_help_panel=_AUTO
     ),
+    extra_info: str | None = typer.Option(
+        default=None, envvar="EXTRA_INFO", rich_help_panel=_AUTO
+    ),
     verl_ability: str = typer.Option(
         default="math", envvar="VERL_ABILITY", rich_help_panel=_AUTO
     ),
@@ -113,6 +116,7 @@ def prep(
             option_cols=op_cols,
             answer_cols=a_cols,
             answer_template=a_template,
+            extra_info=extra_info,
             labels=labels,
             transforms=transforms,
             verl_ability=verl_ability,

src/prep/formatter/auto_cls.py

@@ -16,6 +16,7 @@ def auto_cls(
     split: str,
     a_cols: list[str],
     transforms: list[str],
+    extra_info: str | None = None,
 ):
     images = extract_images(e)
     if len(images) != 1:
@@ -30,7 +31,7 @@ def auto_cls(
         "id": f"{data_name}/{split}{idx:08d}",
         "image": image,
         "label": label,
-        "extra_info": e.get("extra_info", ""),
+        "extra_info": e.get(extra_info, extra_info) or "",
     }
 
 
@@ -47,6 +48,7 @@ def load_cls(path: str, split: str, args: ProcArgs):
             split=split,
             a_cols=args.answer_cols,
             transforms=args.transforms,
+            extra_info=args.extra_info,
         ),
         remove_columns=d.column_names,
         num_proc=args.num_proc,

src/prep/formatter/auto_sft.py

@@ -12,6 +12,7 @@ def auto_sft(
     option_cols: list[str],
     q_template: str,
     a_template: str,
+    extra_info: str | None = None,
 ) -> dict:
     images = extract_images(e)
     question, answer = extract_qa(
@@ -25,7 +26,7 @@ def auto_sft(
             {"role": "assistant", "content": answer},
         ],
         "id": f"{data_name}/{idx:08d}",
-        "extra_info": e.get("extra_info", ""),
+        "extra_info": e.get(extra_info, extra_info) or "",
     }
 
 
@@ -44,6 +45,7 @@ def load_sft(path: str, split: str, args: ProcArgs):
             option_cols=args.option_cols,
             q_template=args.question_template,
             a_template=args.answer_template,
+            extra_info=args.extra_info,
         ),
         remove_columns=d.column_names,
         num_proc=args.num_proc,

src/prep/formatter/auto_verl.py

@@ -15,6 +15,7 @@ def auto_verl(
     a_template: str,
     ability: str,
     style: str,
+    extra_info: str | None = None,
 ):
     images = extract_images(e)
     question, answer = extract_qa(
@@ -35,7 +36,7 @@ def auto_verl(
             "split": split,
             "index": f"{idx:08d}",
             "explanation": e.get("explanation", ""),
-            "misc": e.get("extra_info", ""),
+            "misc": e.get(extra_info, extra_info) or "",
         },
     }
 
@@ -58,6 +59,7 @@ def load_verl(path: str, split: str, args: ProcArgs):
             a_template=args.answer_template,
             ability=args.verl_ability,
             style=args.verl_style,
+            extra_info=args.extra_info,
         ),
         remove_columns=d.column_names,
         num_proc=args.num_proc,

tests/formatter/test_auto_cls.py

@@ -37,7 +37,13 @@ def test_auto_cls_basic(self):
         }
 
         result = auto_cls(
-            example, 0, data_name="test", split="train", a_cols=["label"], transforms=[]
+            example,
+            0,
+            data_name="test",
+            split="train",
+            a_cols=["label"],
+            transforms=[],
+            extra_info="extra_info",
         )
 
         assert result["id"] == "test/train00000000"

tests/formatter/test_auto_sft.py

@@ -24,6 +24,7 @@ def test_auto_sft_basic_text_only(self):
             option_cols=[],
             q_template="Question: {question}",
             a_template="{answer}",
+            extra_info="extra_info",
         )
 
         assert result["id"] == "test/00000000"

tests/formatter/test_auto_verl.py

@@ -109,6 +109,7 @@ def test_auto_verl_extra_info_fields(self):
             a_template="{answer}",
             ability="math",
             style="rule",
+            extra_info="extra_info",
         )
 
         assert result["extra_info"]["split"] == "train"