From 1ea7a016942aed331238cb3a71dcc98daef13bad Mon Sep 17 00:00:00 2001 From: Moenupa <49304833+Moenupa@users.noreply.github.com> Date: Mon, 31 Aug 2026 15:16:41 +0800 Subject: [PATCH] [feat] support extra-info columns --- src/prep/api/types.py | 4 +++- src/prep/cli.py | 4 ++++ src/prep/formatter/auto_cls.py | 4 +++- src/prep/formatter/auto_sft.py | 4 +++- src/prep/formatter/auto_verl.py | 4 +++- tests/formatter/test_auto_cls.py | 8 +++++++- tests/formatter/test_auto_sft.py | 1 + tests/formatter/test_auto_verl.py | 1 + 8 files changed, 25 insertions(+), 5 deletions(-) diff --git a/src/prep/api/types.py b/src/prep/api/types.py index 858c2d7..61cb30c 100644 --- a/src/prep/api/types.py +++ b/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 diff --git a/src/prep/cli.py b/src/prep/cli.py index 8cc031c..463fea8 100644 --- a/src/prep/cli.py +++ b/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, diff --git a/src/prep/formatter/auto_cls.py b/src/prep/formatter/auto_cls.py index 799ab8b..fe477b1 100644 --- a/src/prep/formatter/auto_cls.py +++ b/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, diff --git a/src/prep/formatter/auto_sft.py b/src/prep/formatter/auto_sft.py index 7547cd0..e12acc5 100644 --- a/src/prep/formatter/auto_sft.py +++ b/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, diff --git a/src/prep/formatter/auto_verl.py b/src/prep/formatter/auto_verl.py index 4c40c03..eaba150 100644 --- a/src/prep/formatter/auto_verl.py +++ b/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, diff --git a/tests/formatter/test_auto_cls.py b/tests/formatter/test_auto_cls.py index 09d77fd..630ee0f 100644 --- a/tests/formatter/test_auto_cls.py +++ b/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" diff --git a/tests/formatter/test_auto_sft.py b/tests/formatter/test_auto_sft.py index 4b899d0..7e6e9ea 100644 --- a/tests/formatter/test_auto_sft.py +++ b/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" diff --git a/tests/formatter/test_auto_verl.py b/tests/formatter/test_auto_verl.py index 6463b65..5379eda 100644 --- a/tests/formatter/test_auto_verl.py +++ b/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"