diff --git a/.gitignore b/.gitignore index 4ed61d716e3..ab45d9c26a6 100644 --- a/.gitignore +++ b/.gitignore @@ -129,6 +129,7 @@ csrc/third_party/ dataset/ output/ !tests/dataset/ +!paddleformers/datasets_v2/dataset/ # gen codes autogen/ @@ -154,6 +155,7 @@ csrc/third_party/ dataset/ output/ !tests/dataset/ +!paddleformers/datasets_v2/dataset/ # gen codes autogen/ diff --git a/docs/zh/datasets_v2_design.md b/docs/zh/datasets_v2_design.md new file mode 100644 index 00000000000..826836c163d --- /dev/null +++ b/docs/zh/datasets_v2_design.md @@ -0,0 +1,224 @@ +# datasets_v2 架构设计文档 + +## 1. 概述 + +datasets_v2 是 PaddleFormers 的新一代数据处理模块，基于 HuggingFace datasets 库构建，目标是提供统一的、可扩展的数据加载与预处理管线。 + +**核心设计原则：** +- 显式 schema 作为管线契约（所有模块基于同一份字段定义交互） +- 基于 HF Dataset.map() 的批处理框架（兼容 Map 式和 Iterable 流式） +- 多格式归一化：无论输入是什么格式，输出统一为 messages 标准格式 +- 面向多模态 + 纯文本双重场景 + +--- + +## 2. 模块划分与开发顺序 + +``` +datasets_v2/ +├── schema.py # Step 1: 数据契约（字段定义、类型、校验） +├── preprocessors/ # Step 2: 数据预处理（格式归一化） +│ ├── base.py # 批处理框架 +│ ├── response.py # query/response 格式 +│ ├── messages.py # 对话/ShareGPT 格式 +│ ├── extra.py # Alpaca 等薄封装 +│ ├── auto.py # 自动格式检测与分发 +│ └── __init__.py +├── ops.py # Step 3: 数据集操作（sample/split/concat/shuffle/shard） +├── registry.py # Step 4: 数据集注册表 +├── loaders.py # Step 4: 数据加载器（本地/远程/多格式） +├── builder.py # Step 5: load_dataset 统一入口 +├── adapters.py # Step 6: 训练框架适配（Paddle DataLoader 对接） +└── __init__.py # 公共 API 导出 +``` + +--- + +## 3. Step 1: schema.py — 数据契约 + +### 3.1 标准字段（29 个） + +| 类别 | 字段 | +|------|------| +| 核心对话 | messages | +| 多模态 | images, videos, audios | +| 工具调用 | tools | +| Grounding | objects | +| 偏好学习 | rejected_/positive_/negative_ × 上述 6 个 | +| 标量 | rejected_response, label, channel, margin, teacher_prompt | + +### 3.2 消息格式 + +```python +messages = [ + {"role": "system", "content": "...", "loss": False}, # loss 可选 + {"role": "user", "content": "..."}, + {"role": "assistant", "content": "..."}, +] +``` + +合法 role: `system`, `user`, `assistant`, `tool_call`, `tool_response`, `tool` + +### 3.3 多模态数据格式 + +```python +images: List[{"bytes": Optional[bytes], "path": str}] +videos: List[str] # 路径列表 +audios: List[str] # 路径列表 +``` + +### 3.4 关键函数 + +- `check_messages(messages)` — 校验 messages 结构合法性 +- `cast_images(images)` — 将 str/dict/list 归一化为 `List[ImageMedia]` +- `cast_media_list(media)` — 将 str 归一化为 `List[str]` +- `remove_non_standard_keys(row)` — 过滤非标准字段 + +--- + +## 4. Step 2: preprocessors/ — 数据预处理 + +### 4.1 BasePreprocessor（base.py） + +核心机制：将 `HF Dataset.map(batched=True)` 与逐行处理逻辑桥接。 + +``` +Dataset.map(batched=True, batch_size=1000) + │ + ▼ +_batched_preprocess(batched_row) + │ _batched_to_rows: 列式 → 行式 + │ for row: preprocess(row) → dict / list[dict] / None + │ _check_and_cast: 校验 + 类型归一化 + │ _rows_to_batched: 行式 → 列式 + ▼ +返回给 HF 拼接 +``` + +**关键设计：** +- `preprocess(row)` 由子类实现，返回 dict（一行）、list[dict]（展开）、None（跳过） +- 容错：`strict=False` 时单行出错只 warning + 跳过，不中断整体 +- 列重命名：`DEFAULT_COLUMN_ALIASES` + 用户自定义 `columns` 参数 +- `num_proc` 仅对 HfMapDataset 生效（IterableDataset 不支持多进程） + +### 4.2 ResponsePreprocessor（response.py） + +**输入格式：** +```python +{"query": "...", "response": "...", "system": "...", "history": [["q1","r1"], ...]} +``` + +**处理逻辑：** +1. pop response（若为 list 取第一个） +2. pop history（支持字符串格式自动 `ast.literal_eval`） +3. `history.append([query, response])` +4. `history_to_messages(history, system)` → 标准 messages + +**列别名：** query ← prompt/input/instruction/question/problem; response ← answer/output/targets/text/completion/content + +### 4.3 MessagesPreprocessor（messages.py） + +**输入格式（三种）：** + +标准格式： +```python +{"messages": [{"role": "user", "content": "..."}, ...]} +``` + +非标准 key： +```python +{"conversations": [{"from": "human", "value": "..."}, {"from": "gpt", "value": "..."}]} +``` + +ShareGPT paired 格式： +```python +{"conversations": [{"human": "...", "gpt": "..."}, ...]} +``` + +**处理逻辑：** +1. 自动检测 role_key（role/from）、content_key（content/value） +2. 判断是否 ShareGPT 格式（第一条消息无 role/content 字段） +3. 统一角色名：human→user, gpt→assistant, function_call→tool_call 等 +4. 插入 system 消息 + +### 4.4 AlpacaPreprocessor（extra.py） + +继承 ResponsePreprocessor，仅做字段拼接： +```python +query = f"{instruction}\n{input}" if both else instruction or input +row['response'] = output +→ 委托给 ResponsePreprocessor.preprocess() +``` + +### 4.5 AutoPreprocessor（auto.py） + +分发逻辑： +```python +if 'messages'/'conversation'/'conversations' in columns → MessagesPreprocessor +elif 'instruction' + 'input' in columns → AlpacaPreprocessor +else → ResponsePreprocessor +``` + +--- + +## 5. Step 3-7: 待实现模块 + +### 5.1 ops.py — 数据集操作 + +提供： +- `sample_dataset(dataset, n)` — 采样 +- `split_dataset(dataset, ratio)` — 训练/验证拆分 +- `concat_datasets([ds1, ds2, ...])` — 合并 +- `interleave_datasets([ds1, ds2, ...], probs)` — 按比例交错 +- `shuffle_dataset(dataset, seed)` — 打乱 +- `shard_dataset(dataset, num_shards, index)` — 分片 + +### 5.2 registry.py — 数据集注册表 + +- 数据集名称 → 元信息（路径/URL、preprocessor 类型、列映射）的注册机制 +- 支持内置数据集和用户自定义注册 + +### 5.3 loaders.py — 数据加载 + +- 本地文件加载（json/jsonl/csv/parquet） +- HuggingFace Hub 加载 +- 自定义 URL 下载 +- 统一返回 HF Dataset/IterableDataset + +### 5.4 builder.py — load_dataset 入口 + +用户侧唯一入口： +```python +from paddleformers.datasets_v2 import load_dataset +dataset = load_dataset("dataset_name", split="train", streaming=False) +``` +内部串联 registry → loader → preprocessor → ops + +### 5.5 adapters.py — 训练框架适配 + +- 将 HF Dataset 转为 Paddle DataLoader 可消费的格式 +- 处理 collate_fn、padding、动态 batching 等 + +--- + +## 6. 关键设计决策记录 + +1. **schema 独立成文件** — 将字段定义集中到 schema.py 作为单一真相源 +2. **HF Dataset 类型命名** — `HfMapDataset` / `HfIterableDataset`，体现来源（HF）和访问模式（Map/Iterable） +3. **DATASET_TYPE 放在 schema** — 跨模块使用的类型定义统一放在契约层 +4. **preprocessors 分文件** — 基础类各自独立（response、messages），薄封装合并到 extra.py +5. **AlpacaPreprocessor 移除 instruction/input/output 的列映射** — 避免与 ResponsePreprocessor 的 QUERY_KEYS 冲突，这些字段由 preprocess() 手动处理 +6. **_rows_to_batched 使用 set().union()** — 字段顺序不影响 HF Dataset（基于 key 访问），性能优于逐行收集 + +--- + +## 7. 测试 + +测试文件：`tests/datasets_v2/test_preprocessors.py` + +调试配置：最外层 `.vscode/launch.json` → "Test datasets_v2 Preprocessors" + +运行命令： +```bash +python -m pytest tests/datasets_v2/test_preprocessors.py -v +``` diff --git a/examples/config/sft/full.yaml b/examples/config/sft/full.yaml index 7c5907060d7..fc8cf5bf0d8 100644 --- a/examples/config/sft/full.yaml +++ b/examples/config/sft/full.yaml @@ -17,7 +17,7 @@ _attn_implementation: flashmask ### finetuning # base -stage: SFT +stage: SFT-V2 fine_tuning: full seed: 23 do_train: true diff --git a/examples/config/sft/full_function_call.yaml b/examples/config/sft/full_function_call.yaml index ce7d2f58c91..b5552e98cc5 100644 --- a/examples/config/sft/full_function_call.yaml +++ b/examples/config/sft/full_function_call.yaml @@ -20,7 +20,7 @@ loss_subbatch_sequence_length: 8192 ### finetuning # base -stage: SFT +stage: SFT-V2 fine_tuning: full seed: 23 do_train: true diff --git a/examples/config/sft/full_map.yaml b/examples/config/sft/full_map.yaml index 85bb6f7c810..86a5ce40c9f 100644 --- a/examples/config/sft/full_map.yaml +++ b/examples/config/sft/full_map.yaml @@ -18,7 +18,7 @@ _attn_implementation: flashmask ### finetuning # base -stage: SFT +stage: SFT-V2 fine_tuning: full seed: 23 do_train: true diff --git a/examples/config/sft/full_tp_pp.yaml b/examples/config/sft/full_tp_pp.yaml index dfba5e4a420..22ae6d8c56a 100644 --- a/examples/config/sft/full_tp_pp.yaml +++ b/examples/config/sft/full_tp_pp.yaml @@ -17,7 +17,7 @@ _attn_implementation: flashmask ### finetuning # base -stage: SFT +stage: SFT-V2 fine_tuning: full seed: 23 do_train: true diff --git a/examples/config/sft/full_tp_pp_ep.yaml b/examples/config/sft/full_tp_pp_ep.yaml index c85a978c9bf..10ed2ec25fd 100644 --- a/examples/config/sft/full_tp_pp_ep.yaml +++ b/examples/config/sft/full_tp_pp_ep.yaml @@ -18,7 +18,7 @@ _attn_implementation: flashmask ### finetuning # base -stage: SFT +stage: SFT-V2 fine_tuning: full seed: 23 do_train: true diff --git a/examples/config/sft/lora.yaml b/examples/config/sft/lora.yaml index 9601cea349e..4fb09077f56 100644 --- a/examples/config/sft/lora.yaml +++ b/examples/config/sft/lora.yaml @@ -19,7 +19,7 @@ lora_rank: 8 ### finetuning # base -stage: SFT +stage: SFT-V2 fine_tuning: lora seed: 23 do_train: true diff --git a/examples/config/sft/lora_tp_pp.yaml b/examples/config/sft/lora_tp_pp.yaml index 8495b7935d2..e7b1199c31d 100644 --- a/examples/config/sft/lora_tp_pp.yaml +++ b/examples/config/sft/lora_tp_pp.yaml @@ -19,7 +19,7 @@ lora_rank: 8 ### finetuning # base -stage: SFT +stage: SFT-V2 fine_tuning: lora seed: 23 do_train: true diff --git a/examples/config/sft/lora_tp_pp_ep.yaml b/examples/config/sft/lora_tp_pp_ep.yaml index cea763c69fe..781e79acf93 100644 --- a/examples/config/sft/lora_tp_pp_ep.yaml +++ b/examples/config/sft/lora_tp_pp_ep.yaml @@ -19,7 +19,7 @@ lora_rank: 8 ### finetuning # base -stage: SFT +stage: SFT-V2 fine_tuning: lora seed: 23 do_train: true diff --git a/paddleformers/cli/train/dpo/__init__.py b/paddleformers/cli/train/dpo/__init__.py index 9bee88e4281..0c2f61ec291 100644 --- a/paddleformers/cli/train/dpo/__init__.py +++ b/paddleformers/cli/train/dpo/__init__.py @@ -14,4 +14,12 @@ from .workflow import run_dpo -__all__ = ["run_dpo"] +__all__ = ["run_dpo", "run_dpo_v2"] + + +def __getattr__(name): + if name == "run_dpo_v2": + from .workflow2 import run_dpo_v2 + + return run_dpo_v2 + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/paddleformers/cli/train/dpo/workflow2.py b/paddleformers/cli/train/dpo/workflow2.py new file mode 100644 index 00000000000..14646337f74 --- /dev/null +++ b/paddleformers/cli/train/dpo/workflow2.py @@ -0,0 +1,396 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""DPO workflow using datasets_v2 pipeline. + +Uses the new datasets_v2 module for data loading, encoding, and collation, +while reusing the existing DPOTrainer and DPO loss functions. +""" + +import logging + +logging.getLogger("paddleformers.transformers.modeling_rope_utils").setLevel(logging.ERROR) + +import math +import os +from functools import partial + +import paddle + +from paddleformers.cli.hparams import ( + DataArguments, + FinetuningArguments, + GeneratingArguments, + ModelArguments, +) +from paddleformers.cli.utils.process import add_new_special_tokens +from paddleformers.datasets_v2 import LazyEncodeDataset, get_template +from paddleformers.datasets_v2 import load_dataset as v2_load_dataset +from paddleformers.datasets_v2.datapipe.collate import collate_dpo +from paddleformers.datasets_v2.datapipe.encode import DPOEncodeConfig, encode_dpo +from paddleformers.nn.attention import AttentionInterface +from paddleformers.peft import LoRAConfig, LoRAModel +from paddleformers.trainer import ( + IntervalStrategy, + get_last_checkpoint, + set_random_seed, + set_seed, +) +from paddleformers.transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer +from paddleformers.transformers.configuration_utils import ( + LlmMetaConfig, + QuantizationConfig, +) +from paddleformers.utils.log import logger + +from .dpo_argument import DPOConfig +from .dpo_trainer import DPOTrainer + +os.environ["USE_CASUAL_MASK"] = "False" + + +def _detect_template(tokenizer, data_args) -> str: + """Detect which template to use based on config and tokenizer.""" + if hasattr(data_args, "use_template") and not data_args.use_template: + return "empty" + + if data_args.template: + return data_args.template + + model_path = getattr(data_args, "_model_name_or_path", "") + model_lower = model_path.lower() + if "qwen" in model_lower: + return "chatml" + elif "llama" in model_lower: + return "llama3" + elif "deepseek" in model_lower: + return "deepseek3" + elif "glm" in model_lower: + return "chatml" + + if tokenizer.chat_template is not None: + return "__jinja__" + return "chatml" + + +def run_dpo_v2( + model_args: "ModelArguments", + data_args: "DataArguments", + generating_args: "GeneratingArguments", + finetuning_args: "FinetuningArguments", +): + """Run DPO training using datasets_v2 pipeline.""" + + training_args = finetuning_args + training_args.max_seq_len = data_args.max_seq_len + training_args.model_name_or_path = model_args.model_name_or_path + training_args.download_hub = model_args.download_hub + training_args.copy_custom_file_list = model_args.copy_custom_file_list + + # DPO-specific loss type adjustments + if training_args.loss_type == "orpo": + training_args.reference_free = True + training_args.sft_loss_ratio = 1.0 + training_args.loss_type = "or" + logger.info("orpo loss_type is equal to sft_loss + pref_loss_ratio * or_loss.") + if training_args.loss_type in ["or", "simpo"] and not training_args.reference_free: + training_args.reference_free = True + logger.warning( + f"{training_args.loss_type} loss_type only supports reference_free. Set reference_free to True." + ) + + training_args.print_config(model_args, "Model") + training_args.print_config(data_args, "Data") + training_args.print_config(training_args, "Train") + + # Setup GPU & distributed training + paddle.set_device(training_args.device) + set_random_seed(seed_=training_args.seed) + set_seed(seed=training_args.seed) + logger.warning( + f"Process rank: {training_args.local_rank}, device: {training_args.device}, " + f"world_size: {training_args.world_size}, " + f"distributed training: {bool(training_args.local_rank != -1)}, " + f"16-bits training: {training_args.fp16 or training_args.bf16}" + ) + + # Detecting last checkpoint + last_checkpoint = None + if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir: + last_checkpoint = get_last_checkpoint(training_args.output_dir) + if last_checkpoint is not None and training_args.resume_from_checkpoint is None: + logger.info( + f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change " + "the `--output_dir` or add `--overwrite_output_dir` to train from scratch." + ) + + # ====== DPO Config ====== + dpo_config = DPOConfig( + beta=training_args.beta, + offset_alpha=training_args.offset_alpha, + simpo_gamma=training_args.simpo_gamma, + normalize_logps=training_args.normalize_logps, + ignore_eos_token=training_args.ignore_eos_token, + label_smoothing=training_args.label_smoothing, + loss_type=training_args.loss_type, + pref_loss_ratio=training_args.pref_loss_ratio, + sft_loss_ratio=training_args.sft_loss_ratio, + dpop_lambda=training_args.dpop_lambda, + ref_model_update_steps=training_args.ref_model_update_steps, + reference_free=training_args.reference_free, + lora=model_args.lora, + ) + + # ====== Model Setup ====== + if training_args.fp16_opt_level == "O2": + if training_args.fp16: + dtype = "float16" + elif training_args.bf16: + dtype = "bfloat16" + else: + raise ValueError("Please specific dtype: --fp16 or --bf16") + else: + dtype = "float32" + + if finetuning_args.weight_quantize_algo is not None: + quantization_config = dict( + weight_quantize_algo=finetuning_args.weight_quantize_algo, + ignore_modules=[".*out_linear.*"], + ) + else: + quantization_config = dict(weight_quantize_algo=finetuning_args.weight_quantize_algo) + quantization_config = QuantizationConfig.from_dict(quantization_config) + + model_config = AutoConfig.from_pretrained( + model_args.model_name_or_path, + dtype=dtype, + quantization_config=quantization_config, + ) + + LlmMetaConfig.set_llm_config(model_config, training_args) + + avaible_attn_impl = AttentionInterface._global_mapping.keys() + if model_args._attn_implementation not in avaible_attn_impl: + raise ValueError( + f"Invalid _attn_implementation: {model_args._attn_implementation}, available: {avaible_attn_impl}" + ) + + model_config.seq_length = data_args.max_seq_len + model_config.max_sequence_length = data_args.max_seq_len + model_config._attn_implementation = model_args._attn_implementation + model_config.is_lora = model_args.lora + model_config.dpo_config = dpo_config + + logger.info(f"Final model config: {model_config}") + logger.info("Loading model...") + + model_class = AutoModelForCausalLM + + # Reference model setup + ref_model = None + if not training_args.reference_free and not model_args.lora: + ref_model_config = AutoConfig.from_pretrained(model_args.model_name_or_path, dtype=dtype) + ref_model_config.max_sequence_length = data_args.max_seq_len + ref_model_config.seq_length = data_args.max_seq_len + ref_model_config._attn_implementation = model_args._attn_implementation + ref_model_config.dpo_config = dpo_config + LlmMetaConfig.set_llm_config(ref_model_config, training_args) + + if model_args.continue_training and not training_args.autotuner_benchmark: + model = model_class.from_pretrained( + model_args.model_name_or_path, + config=model_config, + convert_from_hf=training_args.convert_from_hf, + load_via_cpu=training_args.load_via_cpu, + load_checkpoint_format=training_args.load_checkpoint_format, + ) + if not training_args.reference_free and not model_args.lora: + ref_model = model_class.from_config(ref_model_config) + ref_model.set_state_dict(model.state_dict()) + else: + model = model_class.from_config(model_config, dtype=dtype) + if not training_args.reference_free and not model_args.lora: + ref_model = model_class.from_config(ref_model_config) + ref_model.set_state_dict(model.state_dict()) + + # ====== Tokenizer ====== + tokenizer = AutoTokenizer.from_pretrained(model_args.model_name_or_path) + add_new_special_tokens(tokenizer, data_args.new_special_tokens_path) + if tokenizer.pad_token_id is None: + tokenizer.pad_token_id = tokenizer.eos_token_id + + # ====== Dataset (datasets_v2) ====== + logger.info("[datasets_v2] Loading DPO dataset...") + + data_args._model_name_or_path = model_args.model_name_or_path + + # Determine template + template_name = _detect_template(tokenizer, data_args) + use_jinja = template_name == "__jinja__" + + if use_jinja: + template = None + logger.info("[datasets_v2] Using tokenizer's built-in chat_template (jinja)") + else: + template = get_template(template_name) + logger.info(f"[datasets_v2] Using template: {template_name}") + + # DPO Encode config + use_filtered_label_loss = getattr(model_config, "use_filtered_label_loss", True) + encode_config = DPOEncodeConfig( + max_seq_len=data_args.max_seq_len, + truncation="right", + label_shift=False, # DPO does NOT shift labels (response_labels are direct token IDs) + use_filtered_label_loss=use_filtered_label_loss, + ) + + encode_fn = partial( + encode_dpo, + tokenizer=tokenizer, + template=template, + config=encode_config, + ) + + # Load train dataset + train_dataset = None + eval_dataset = None + num_proc = getattr(data_args, "dataset_num_proc", 1) + train_format = getattr(data_args, "train_dataset_type", None) + eval_format = getattr(data_args, "eval_dataset_type", None) + + # DPO always uses map mode (no streaming packing needed, needs random access) + if training_args.do_train and training_args.should_load_dataset: + hf_ds = v2_load_dataset( + data_args.train_dataset_path, streaming=False, num_proc=num_proc, dataset_format=train_format + ) + train_dataset = LazyEncodeDataset(hf_ds, encode_fn, seed=training_args.seed) + logger.info(f"[datasets_v2] DPO train dataset loaded: {len(train_dataset)} samples") + + if training_args.do_eval and training_args.should_load_dataset: + hf_eval_ds = v2_load_dataset( + data_args.eval_dataset_path, streaming=False, num_proc=num_proc, dataset_format=eval_format + ) + eval_dataset = LazyEncodeDataset(hf_eval_ds, encode_fn, seed=training_args.seed) + logger.info(f"[datasets_v2] DPO eval dataset loaded: {len(eval_dataset)} samples") + + # ====== Collate Function ====== + max_seq_len = ( + data_args.max_seq_len if training_args.sequence_parallel or training_args.context_parallel_size > 1 else None + ) + logger.info(f"[datasets_v2] max_seq_len for collate: {max_seq_len}") + + use_startend = getattr(model_args, "use_attn_mask_startend_row_indices", True) + + data_collator = partial( + collate_dpo, + pad_token_id=tokenizer.pad_token_id, + max_seq_len=data_args.max_seq_len, + use_attn_mask_startend_row_indices=use_startend, + use_filtered_label_loss=use_filtered_label_loss, + ) + + # ====== LoRA (optional) ====== + if model_args.lora: + from paddleformers.cli.utils import get_lora_target_modules + + if model_args.lora_path is None: + target_modules = get_lora_target_modules(model) + lora_config = LoRAConfig( + target_modules=target_modules, + r=model_args.lora_rank, + lora_alpha=2 * model_args.lora_rank if not model_args.rslora else 4, + rslora=model_args.rslora, + lora_plus_scale=model_args.lora_plus_scale, + merge_weights=False, + tensor_model_parallel_size=training_args.tensor_model_parallel_size, + dtype=dtype, + base_model_name_or_path=model_args.model_name_or_path, + ) + model = LoRAModel(model, lora_config) + else: + model = LoRAModel.from_pretrained(model=model, lora_path=model_args.lora_path) + + # ====== max_steps calculation ====== + if training_args.max_steps == -1: + if training_args.should_load_dataset and paddle.distributed.get_rank() == 0: + training_args.max_steps = math.ceil(len(train_dataset) / training_args.global_batch_size) + training_args.max_steps *= training_args.num_train_epochs + logger.info( + f"len(train_dataset): {len(train_dataset)}, " + f"global_batch_size: {training_args.global_batch_size}, " + f"num_train_epochs: {training_args.num_train_epochs}, " + f"max_steps: {training_args.max_steps}" + ) + + if paddle.distributed.get_world_size() > 1: + paddle.distributed.barrier() + max_steps = paddle.to_tensor([training_args.max_steps]) + paddle.distributed.broadcast(max_steps, src=0) + training_args.max_steps = int(max_steps.item()) + + if training_args.max_steps <= 0: + raise ValueError(f"Invalid max_steps: {training_args.max_steps}. Please check your dataset") + logger.info(f"Re-setting training_args.max_steps to {training_args.max_steps}.") + + if training_args.decay_steps is None: + training_args.decay_steps = training_args.max_steps + + if training_args.save_strategy == IntervalStrategy.EPOCH: + training_args.save_strategy = IntervalStrategy.STEPS + training_args.save_steps = int(training_args.max_steps / training_args.num_train_epochs) + if training_args.evaluation_strategy == IntervalStrategy.EPOCH: + training_args.evaluation_strategy = IntervalStrategy.STEPS + training_args.eval_steps = int(training_args.max_steps / training_args.num_train_epochs) + if training_args.logging_strategy == IntervalStrategy.EPOCH: + training_args.logging_strategy = IntervalStrategy.STEPS + training_args.logging_steps = int(training_args.max_steps / training_args.num_train_epochs) + + # ====== Create DPO Trainer ====== + trainer = DPOTrainer( + model=model, + ref_model=ref_model, + dpo_config=dpo_config, + args=training_args, + train_dataset=(train_dataset if training_args.do_train and training_args.should_load_dataset else None), + eval_dataset=(eval_dataset if training_args.do_eval and training_args.should_load_dataset else None), + tokenizer=tokenizer, + data_collator=data_collator, + model_with_dpo_criterion=getattr(model_args, "model_with_dpo_criterion", False), + ) + + trainable_parameters = [ + p for p in model.parameters() if not p.stop_gradient or ("quantization_linear" in p.name and "w_1" in p.name) + ] + trainer.set_optimizer_grouped_parameters(trainable_parameters) + + # ====== Train ====== + if training_args.do_train: + checkpoint = None + if training_args.resume_from_checkpoint is not None: + checkpoint = training_args.resume_from_checkpoint + elif last_checkpoint is not None: + checkpoint = last_checkpoint + train_result = trainer.train(resume_from_checkpoint=checkpoint) + + if not training_args.autotuner_benchmark and not getattr(training_args, "benchmark", False): + trainer.save_model(merge_tensor_parallel=training_args.tensor_model_parallel_size > 1, last_fc_to_hf=True) + trainer.log_metrics("train", train_result.metrics) + trainer.save_metrics("train", train_result.metrics) + trainer.save_state() + + # ====== Eval ====== + if training_args.do_eval: + eval_result = trainer.evaluate() + trainer.log_metrics("eval", eval_result) + trainer.save_metrics("eval", eval_result) diff --git a/paddleformers/cli/train/sft/__init__.py b/paddleformers/cli/train/sft/__init__.py index 5fca2fb42fa..63568d2fc16 100644 --- a/paddleformers/cli/train/sft/__init__.py +++ b/paddleformers/cli/train/sft/__init__.py @@ -14,4 +14,16 @@ from .workflow import run_sft -__all__ = ["run_sft"] +__all__ = ["run_sft", "run_sft_v2", "run_vl_sft_v2"] + + +def __getattr__(name): + if name == "run_sft_v2": + from .workflow2 import run_sft_v2 + + return run_sft_v2 + if name == "run_vl_sft_v2": + from .workflow_vl_v2 import run_vl_sft_v2 + + return run_vl_sft_v2 + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/paddleformers/cli/train/sft/workflow2.py b/paddleformers/cli/train/sft/workflow2.py new file mode 100644 index 00000000000..f60868ea8c1 --- /dev/null +++ b/paddleformers/cli/train/sft/workflow2.py @@ -0,0 +1,473 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""SFT workflow using datasets_v2 pipeline. + +Simplified version of workflow.py that uses the new datasets_v2 module +for data loading, encoding, and collation. Supports packing + flashmask. +""" + +import logging + +import numpy as np + +logging.getLogger("paddleformers.transformers.modeling_rope_utils").setLevel(logging.ERROR) + +import math +import os +from functools import partial + +import paddle + +from paddleformers.cli.hparams import ( + DataArguments, + FinetuningArguments, + GeneratingArguments, + ModelArguments, +) +from paddleformers.cli.utils.process import add_new_special_tokens +from paddleformers.datasets_v2 import ( + EncodeConfig, + LazyEncodeDataset, + StreamingDataset, + encode_pt, + encode_sft, + get_template, +) +from paddleformers.datasets_v2 import load_dataset as v2_load_dataset +from paddleformers.datasets_v2.datapipe.collate import collate_sft +from paddleformers.nn.attention import AttentionInterface +from paddleformers.peft import LoRAConfig, LoRAModel +from paddleformers.trainer import ( + IntervalStrategy, + get_last_checkpoint, + set_random_seed, + set_seed, +) +from paddleformers.transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer +from paddleformers.transformers.configuration_utils import ( + LlmMetaConfig, + QuantizationConfig, +) +from paddleformers.utils.log import logger + +from .sft_trainer import SFTTrainer + +os.environ["USE_CASUAL_MASK"] = "False" + + +def _detect_template(tokenizer, data_args) -> str: + """Detect which template to use based on config and tokenizer. + + Priority: + 1. use_template=False → "empty" (no template formatting) + 2. Explicit template name in config → use that + 3. Tokenizer has chat_template → "__jinja__" (parse from tokenizer, aligned with V1 parse_template) + 4. Fallback → "chatml" + """ + if hasattr(data_args, "use_template") and not data_args.use_template: + return "empty" + + if data_args.template: + return data_args.template + + # Use tokenizer's chat_template if available (aligned with V1's parse_template behavior) + if tokenizer.chat_template is not None: + return "__jinja__" + return "chatml" + + +def run_sft_v2( + model_args: "ModelArguments", + data_args: "DataArguments", + generating_args: "GeneratingArguments", + finetuning_args: "FinetuningArguments", +): + """Run SFT training using datasets_v2 pipeline.""" + + training_args = finetuning_args + training_args.max_seq_len = data_args.max_seq_len + training_args.model_name_or_path = model_args.model_name_or_path + training_args.download_hub = model_args.download_hub + training_args.copy_custom_file_list = model_args.copy_custom_file_list + + training_args.print_config(model_args, "Model") + training_args.print_config(data_args, "Data") + training_args.print_config(training_args, "Train") + + # Setup GPU & distributed training + paddle.set_device(training_args.device) + set_random_seed(seed_=training_args.seed) + set_seed(seed=training_args.seed) + logger.warning( + f"Process rank: {training_args.local_rank}, device: {training_args.device}, " + f"world_size: {training_args.world_size}, " + f"distributed training: {bool(training_args.local_rank != -1)}, " + f"16-bits training: {training_args.fp16 or training_args.bf16}" + ) + + # Detecting last checkpoint + last_checkpoint = None + if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir: + last_checkpoint = get_last_checkpoint(training_args.output_dir) + if last_checkpoint is not None and training_args.resume_from_checkpoint is None: + logger.info( + f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change " + "the `--output_dir` or add `--overwrite_output_dir` to train from scratch." + ) + + # ====== Model Setup ====== + if training_args.fp16_opt_level == "O2": + if training_args.fp16: + dtype = "float16" + elif training_args.bf16: + dtype = "bfloat16" + else: + raise ValueError("Please specific dtype: --fp16 or --bf16") + else: + dtype = "float32" + + if finetuning_args.weight_quantize_algo is not None: + quantization_config = dict( + weight_quantize_algo=finetuning_args.weight_quantize_algo, + ignore_modules=[".*out_linear.*"], + ) + else: + quantization_config = dict(weight_quantize_algo=finetuning_args.weight_quantize_algo) + quantization_config = QuantizationConfig.from_dict(quantization_config) + + model_config = AutoConfig.from_pretrained( + model_args.model_name_or_path, + dtype=dtype, + quantization_config=quantization_config, + ) + + LlmMetaConfig.set_llm_config(model_config, training_args) + + if hasattr(model_config, "hidden_dropout_prob"): + model_config.hidden_dropout_prob = finetuning_args.hidden_dropout_prob + if hasattr(model_config, "attention_probs_dropout_prob"): + model_config.attention_probs_dropout_prob = finetuning_args.attention_probs_dropout_prob + if hasattr(model_config, "ignore_index"): + model_config.ignore_index = -100 + + avaible_attn_impl = AttentionInterface._global_mapping.keys() + if model_args._attn_implementation not in avaible_attn_impl: + raise ValueError( + f"Invalid _attn_implementation: {model_args._attn_implementation}, " f"available: {avaible_attn_impl}" + ) + + model_config.seq_length = data_args.max_seq_len + model_config.max_sequence_length = data_args.max_seq_len + model_config._attn_implementation = model_args._attn_implementation + model_config.is_lora = model_args.lora + + # MTP: extend model seq_length to accommodate extra prediction tokens + num_nextn_predict_layers = getattr(training_args, "num_nextn_predict_layers", 0) + if num_nextn_predict_layers > 0: + model_config.seq_length = data_args.max_seq_len + num_nextn_predict_layers + model_config.max_sequence_length = data_args.max_seq_len + num_nextn_predict_layers + + # Autoregressive MTP training: swap mtp_num_layers and num_nextn_predict_layers + if getattr(model_config, "mtp_num_layers", 0) > 1: + tmp = model_config.mtp_num_layers + model_config.mtp_num_layers = model_config.num_nextn_predict_layers + model_config.num_nextn_predict_layers = tmp + tmp = training_args.mtp_num_layers + training_args.mtp_num_layers = training_args.num_nextn_predict_layers + training_args.num_nextn_predict_layers = tmp + num_nextn_predict_layers = training_args.num_nextn_predict_layers + logger.info( + f"MTP args swapped for autoregressive mtp training: " + f"mtp_num_layers={model_config.mtp_num_layers}, " + f"num_nextn_predict_layers={model_config.num_nextn_predict_layers}" + ) + + logger.info(f"Final model config: {model_config}") + logger.info("Loading model...") + + model_class = AutoModelForCausalLM + if model_args.continue_training and not training_args.autotuner_benchmark: + model = model_class.from_pretrained( + model_args.model_name_or_path, + config=model_config, + convert_from_hf=training_args.convert_from_hf, + load_via_cpu=training_args.load_via_cpu, + load_checkpoint_format=training_args.load_checkpoint_format, + ) + else: + model = model_class.from_config(model_config, dtype=dtype) + + # ====== Tokenizer ====== + tokenizer = AutoTokenizer.from_pretrained(model_args.model_name_or_path) + add_new_special_tokens(tokenizer, data_args.new_special_tokens_path) + if tokenizer.pad_token_id is None: + tokenizer.pad_token_id = tokenizer.eos_token_id + + # ====== Dataset (datasets_v2) ====== + logger.info("[datasets_v2] Loading dataset...") + + data_args._model_name_or_path = model_args.model_name_or_path + + template_name = _detect_template(tokenizer, data_args) + use_jinja = template_name == "__jinja__" + + if use_jinja: + template = None + logger.info("[datasets_v2] Using tokenizer's built-in chat_template (jinja)") + else: + template = get_template(template_name) + logger.info(f"[datasets_v2] Using template: {template_name}") + + # Encode config + encode_config = EncodeConfig( + max_seq_len=data_args.max_seq_len, + truncation="right", + label_shift=True, + ) + + is_pretraining = "pt" in model_args.stage.lower() + + if is_pretraining: + encode_fn = partial( + encode_pt, + tokenizer=tokenizer, + config=encode_config, + ) + logger.info("[datasets_v2] Pretrain mode: using encode_pt") + else: + encode_fn = partial( + encode_sft, + tokenizer=tokenizer, + template=template, + config=encode_config, + ) + + # Load train dataset (and optionally split for eval) + train_dataset = None + eval_dataset = None + num_proc = getattr(data_args, "dataset_num_proc", 1) + dataset_type = getattr(data_args, "dataset_type", "iterable") + use_packing = getattr(data_args, "packing", False) + + # Auto-switch to map mode when packing is enabled (packing requires random access) + auto_switched_to_map = False + if dataset_type == "iterable" and use_packing: + dataset_type = "map" + auto_switched_to_map = True + logger.info("V2 pipeline: auto-switching dataset_type to 'map' (packing requires map mode).") + + is_iterable = dataset_type == "iterable" + + # Get dataset format hint (backward compatible with train_dataset_type config) + train_format = getattr(data_args, "train_dataset_type", None) + eval_format = getattr(data_args, "eval_dataset_type", None) + + if training_args.do_train and training_args.should_load_dataset: + if is_iterable: + if training_args.max_steps <= 0: + raise ValueError("dataset_type=iterable requires --max_steps to be explicitly set (no len available).") + + # Iterable mode: HF streaming=True, returns IterableDataset + hf_ds = v2_load_dataset( + data_args.train_dataset_path, streaming=True, num_proc=1, dataset_format=train_format + ) + + def _streaming_encode(row): + result = encode_fn(row) + if result is None: + return {"input_ids": [], "labels": [], "seq_len": 0} + return {"input_ids": result.input_ids, "labels": result.labels, "seq_len": result.seq_len} + + hf_ds = hf_ds.map(_streaming_encode, remove_columns=hf_ds.column_names or []) + hf_ds = hf_ds.filter(lambda x: x["seq_len"] > 0) + + random_shuffle = getattr(data_args, "random_shuffle", True) + train_dataset = StreamingDataset( + hf_ds, shuffle=random_shuffle, seed=training_args.seed, lazy=training_args.lazy_data_processing + ) + training_args.dataloader_num_workers = 0 + logger.info("[datasets_v2] Train dataset ready (iterable mode, streaming=True)") + else: + # Map mode: HF streaming=False, full download then random access + hf_ds = v2_load_dataset( + data_args.train_dataset_path, streaming=False, num_proc=num_proc, dataset_format=train_format + ) + + # Auto-split: if eval path is same as train or not set + need_auto_split = training_args.do_eval and ( + not data_args.eval_dataset_path or data_args.eval_dataset_path == data_args.train_dataset_path + ) + + if need_auto_split: + from paddleformers.datasets_v2 import split_dataset + + hf_ds, hf_eval_ds = split_dataset(hf_ds, test_ratio=0.1, shuffle=True, seed=training_args.seed) + eval_dataset = LazyEncodeDataset(hf_eval_ds, encode_fn, seed=training_args.seed) + logger.info(f"[datasets_v2] Auto-split: train={len(hf_ds)}, eval={len(hf_eval_ds)}") + + # Pre-shuffle to align with V1 ConcatDataset behavior. + # Explicit map: double shuffle (pre-shuffle + Trainer sampler). + # Auto-switched map (packing): single shuffle only (disable Trainer sampler). + random_shuffle = getattr(data_args, "random_shuffle", True) + if random_shuffle: + preshuffle_rng = np.random.RandomState(0 + training_args.seed) + indices = list(range(len(hf_ds))) + preshuffle_rng.shuffle(indices) + hf_ds = hf_ds.select(indices) + + if auto_switched_to_map: + training_args.dataloader_shuffle = False + logger.info( + "V2 pipeline: disabled Trainer shuffle (auto-switched map, aligning with V1 iterator)." + ) + + train_dataset = LazyEncodeDataset(hf_ds, encode_fn, seed=training_args.seed) + logger.info(f"[datasets_v2] Train dataset loaded: {len(train_dataset)} samples") + + # Eval always uses map mode (StreamingDataset's infinite iteration hangs eval loop) + if training_args.do_eval and training_args.should_load_dataset and eval_dataset is None: + hf_eval_ds = v2_load_dataset(data_args.eval_dataset_path, num_proc=num_proc, dataset_format=eval_format) + eval_dataset = LazyEncodeDataset(hf_eval_ds, encode_fn, seed=training_args.seed) + logger.info(f"[datasets_v2] Eval dataset loaded: {len(eval_dataset)} samples") + + # ====== Collate Function ====== + max_seq_len = ( + data_args.max_seq_len + if (data_args.packing or training_args.sequence_parallel or training_args.context_parallel_size > 1) + else None + ) + + use_startend = getattr(model_args, "use_attn_mask_startend_row_indices", True) + use_global_causal_attn = getattr(model_args, "use_global_causal_attn", False) + + data_collator = partial( + collate_sft, + pad_token_id=tokenizer.pad_token_id, + max_seq_len=data_args.max_seq_len, + packing=data_args.packing, + packing_method="binpacking" if getattr(data_args, "binpacking", False) else "greedy", + packing_interval=getattr(data_args, "packing_interval", 1000), + use_attn_mask_startend_row_indices=use_startend, + num_nextn_predict_layers=num_nextn_predict_layers, + eos_token_id=tokenizer.eos_token_id if num_nextn_predict_layers > 0 else None, + use_global_causal_attn=use_global_causal_attn, + ) + + # ====== LoRA (optional) ====== + if model_args.lora: + from paddleformers.cli.utils import get_lora_target_modules + + if model_args.lora_path is None: + target_modules = get_lora_target_modules(model) + lora_config = LoRAConfig( + target_modules=target_modules, + r=model_args.lora_rank, + lora_alpha=2 * model_args.lora_rank if not model_args.rslora else 4, + rslora=model_args.rslora, + lora_plus_scale=model_args.lora_plus_scale, + merge_weights=False, + tensor_model_parallel_size=training_args.tensor_model_parallel_size, + dtype=dtype, + base_model_name_or_path=model_args.model_name_or_path, + ) + model = LoRAModel(model, lora_config) + else: + model = LoRAModel.from_pretrained(model=model, lora_path=model_args.lora_path) + + # ====== max_steps calculation ====== + if training_args.max_steps == -1: + if is_iterable: + raise ValueError("Streaming (iterable) mode requires --max_steps to be explicitly set.") + if training_args.should_load_dataset and paddle.distributed.get_rank() == 0: + training_args.max_steps = math.ceil(len(train_dataset) / training_args.global_batch_size) + training_args.max_steps *= training_args.num_train_epochs + logger.info( + f"len(train_dataset): {len(train_dataset)}, " + f"global_batch_size: {training_args.global_batch_size}, " + f"num_train_epochs: {training_args.num_train_epochs}, " + f"max_steps: {training_args.max_steps}" + ) + + if paddle.distributed.get_world_size() > 1: + paddle.distributed.barrier() + max_steps = paddle.to_tensor([training_args.max_steps]) + paddle.distributed.broadcast(max_steps, src=0) + training_args.max_steps = int(max_steps.item()) + + if training_args.max_steps <= 0: + raise ValueError(f"Invalid max_steps: {training_args.max_steps}. Please check your dataset") + logger.info(f"Re-setting training_args.max_steps to {training_args.max_steps}.") + + if training_args.decay_steps is None: + training_args.decay_steps = training_args.max_steps + + if training_args.save_strategy == IntervalStrategy.EPOCH: + training_args.save_strategy = IntervalStrategy.STEPS + training_args.save_steps = int(training_args.max_steps / training_args.num_train_epochs) + if training_args.evaluation_strategy == IntervalStrategy.EPOCH: + training_args.evaluation_strategy = IntervalStrategy.STEPS + training_args.eval_steps = int(training_args.max_steps / training_args.num_train_epochs) + if training_args.logging_strategy == IntervalStrategy.EPOCH: + training_args.logging_strategy = IntervalStrategy.STEPS + training_args.logging_steps = int(training_args.max_steps / training_args.num_train_epochs) + + # ====== Create Trainer ====== + trainer = SFTTrainer( + model=model, + args=training_args, + train_dataset=(train_dataset if training_args.do_train and training_args.should_load_dataset else None), + eval_dataset=(eval_dataset if training_args.do_eval and training_args.should_load_dataset else None), + tokenizer=tokenizer, + data_collator=data_collator, + do_generation=False, + data_args=data_args, + ) + + # MTP-only training: freeze everything except MTP layers + if getattr(training_args, "train_mtp_only", False): + from paddleformers.cli.train.sft.workflow import freeze_param_except_mtp + + freeze_param_except_mtp(model, model_config) + + trainable_parameters = [ + p for p in model.parameters() if not p.stop_gradient or ("quantization_linear" in p.name and "w_1" in p.name) + ] + trainer.set_optimizer_grouped_parameters(trainable_parameters) + + # ====== Train ====== + if training_args.do_train: + checkpoint = None + if training_args.resume_from_checkpoint is not None: + checkpoint = training_args.resume_from_checkpoint + elif last_checkpoint is not None: + checkpoint = last_checkpoint + train_result = trainer.train(resume_from_checkpoint=checkpoint) + + total_tokens = ( + data_args.max_seq_len + * training_args.per_device_train_batch_size + * training_args.dataset_world_size + * training_args.gradient_accumulation_steps + * training_args.max_steps + ) + total_tokens_per_second_per_gpu = ( + total_tokens / train_result.metrics["train_runtime"] / training_args.world_size + ) + logger.info(f"Total_Tokens_per_second_per_gpu: {total_tokens_per_second_per_gpu}") + if not training_args.autotuner_benchmark: + trainer.save_model(merge_tensor_parallel=training_args.tensor_model_parallel_size > 1, last_fc_to_hf=True) + trainer.log_metrics("train", train_result.metrics) + trainer.save_metrics("train", train_result.metrics) + trainer.save_state() diff --git a/paddleformers/cli/train/sft/workflow_vl_v2.py b/paddleformers/cli/train/sft/workflow_vl_v2.py new file mode 100644 index 00000000000..480481c5109 --- /dev/null +++ b/paddleformers/cli/train/sft/workflow_vl_v2.py @@ -0,0 +1,432 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""VL-SFT workflow using datasets_v2 pipeline. + +Adds multimodal (image) support on top of the datasets_v2 architecture, +reusing the existing mm_plugin system for model-agnostic vision processing. +""" + +import logging + +logging.getLogger("paddleformers.transformers.modeling_rope_utils").setLevel(logging.ERROR) + +import math +import os +from functools import partial + +import paddle + +from paddleformers.cli.hparams import ( + DataArguments, + FinetuningArguments, + GeneratingArguments, + ModelArguments, +) +from paddleformers.cli.utils.process import add_new_special_tokens +from paddleformers.datasets_v2 import EncodeConfig, LazyEncodeDataset, get_template +from paddleformers.datasets_v2 import load_dataset as v2_load_dataset +from paddleformers.datasets_v2.datapipe.collate import collate_vl_sft +from paddleformers.datasets_v2.datapipe.encode import encode_vl_sft +from paddleformers.datasets_v2.mm_plugin import get_mm_plugin +from paddleformers.nn.attention import AttentionInterface +from paddleformers.peft import LoRAConfig, LoRAModel +from paddleformers.trainer import ( + IntervalStrategy, + get_last_checkpoint, + set_random_seed, + set_seed, +) +from paddleformers.transformers import ( + AutoConfig, + AutoModelForConditionalGeneration, + AutoProcessor, + AutoTokenizer, +) +from paddleformers.transformers.configuration_utils import ( + LlmMetaConfig, + QuantizationConfig, +) +from paddleformers.utils.log import logger + +from .sft_trainer import SFTTrainer + +os.environ["USE_CASUAL_MASK"] = "False" + + +_MM_PLUGIN_MAP = { + "qwen2_vl": "qwen2_vl", + "qwen2_5_vl": "qwen2_vl", + "qwen3_vl": "qwen3_vl", + "qwen3_vl_moe": "qwen3_vl", + "ernie4_5_moe_vl": "ernie_vl", + "paddleocr_vl": "paddleocr_vl", + "glm4v_moe": "glm4v", + "glm_ocr": "glm_ocr", +} + + +def _detect_mm_plugin_name(model_config, data_args) -> str: + """Auto-detect mm_plugin name from model architecture or config.""" + if getattr(data_args, "mm_plugin", None): + return data_args.mm_plugin + + model_type = getattr(model_config, "model_type", "").lower() + return _MM_PLUGIN_MAP.get(model_type, "base") + + +def _get_rope_func(model): + """Extract get_rope_index from model for 3D position_ids.""" + if hasattr(model, "get_rope_index"): + return model.get_rope_index + if hasattr(model, "model") and hasattr(model.model, "get_rope_index"): + return model.model.get_rope_index + return None + + +def _detect_template(tokenizer, data_args) -> str: + """Detect which template to use based on config and tokenizer.""" + if hasattr(data_args, "use_template") and not data_args.use_template: + return "empty" + + if data_args.template: + return data_args.template + + model_path = getattr(data_args, "_model_name_or_path", "") + model_lower = model_path.lower() + if "qwen" in model_lower: + return "chatml" + elif "llama" in model_lower: + return "llama3" + elif "deepseek" in model_lower: + return "deepseek3" + elif "glm" in model_lower: + return "chatml" + + if tokenizer.chat_template is not None: + return "__jinja__" + return "chatml" + + +def run_vl_sft_v2( + model_args: "ModelArguments", + data_args: "DataArguments", + generating_args: "GeneratingArguments", + finetuning_args: "FinetuningArguments", +): + """Run VL-SFT training using datasets_v2 pipeline with multimodal support.""" + + training_args = finetuning_args + training_args.max_seq_len = data_args.max_seq_len + training_args.model_name_or_path = model_args.model_name_or_path + training_args.download_hub = model_args.download_hub + training_args.copy_custom_file_list = model_args.copy_custom_file_list + + training_args.print_config(model_args, "Model") + training_args.print_config(data_args, "Data") + training_args.print_config(training_args, "Train") + + # Setup GPU & distributed training + paddle.set_device(training_args.device) + set_random_seed(seed_=training_args.seed) + set_seed(seed=training_args.seed) + logger.warning( + f"Process rank: {training_args.local_rank}, device: {training_args.device}, " + f"world_size: {training_args.world_size}, " + f"distributed training: {bool(training_args.local_rank != -1)}, " + f"16-bits training: {training_args.fp16 or training_args.bf16}" + ) + + # Detecting last checkpoint + last_checkpoint = None + if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir: + last_checkpoint = get_last_checkpoint(training_args.output_dir) + if last_checkpoint is not None and training_args.resume_from_checkpoint is None: + logger.info( + f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change " + "the `--output_dir` or add `--overwrite_output_dir` to train from scratch." + ) + + # ====== Model Setup ====== + if training_args.fp16_opt_level == "O2": + if training_args.fp16: + dtype = "float16" + elif training_args.bf16: + dtype = "bfloat16" + else: + raise ValueError("Please specific dtype: --fp16 or --bf16") + else: + dtype = "float32" + + if finetuning_args.weight_quantize_algo is not None: + quantization_config = dict( + weight_quantize_algo=finetuning_args.weight_quantize_algo, + ignore_modules=[".*out_linear.*"], + ) + else: + quantization_config = dict(weight_quantize_algo=finetuning_args.weight_quantize_algo) + quantization_config = QuantizationConfig.from_dict(quantization_config) + + model_config = AutoConfig.from_pretrained( + model_args.model_name_or_path, + dtype=dtype, + quantization_config=quantization_config, + ) + + LlmMetaConfig.set_llm_config(model_config, training_args) + + # Sync VL-specific config + if hasattr(model_config, "text_config"): + model_config.text_config.max_sequence_length = data_args.max_seq_len + if hasattr(model_config, "vision_config"): + model_config.vision_config._attn_implementation = model_args._attn_implementation + if hasattr(training_args, "recompute_granularity"): + model_config.vision_config.recompute_granularity = training_args.recompute_granularity + model_config.vision_config.recompute_method = training_args.recompute_method + model_config.vision_config.recompute_num_layers = training_args.recompute_num_layers + + if hasattr(model_config, "hidden_dropout_prob"): + model_config.hidden_dropout_prob = finetuning_args.hidden_dropout_prob + if hasattr(model_config, "attention_probs_dropout_prob"): + model_config.attention_probs_dropout_prob = finetuning_args.attention_probs_dropout_prob + if hasattr(model_config, "ignore_index"): + model_config.ignore_index = -100 + + avaible_attn_impl = AttentionInterface._global_mapping.keys() + if model_args._attn_implementation not in avaible_attn_impl: + raise ValueError( + f"Invalid _attn_implementation: {model_args._attn_implementation}, " f"available: {avaible_attn_impl}" + ) + + model_config.seq_length = data_args.max_seq_len + model_config.max_sequence_length = data_args.max_seq_len + model_config._attn_implementation = model_args._attn_implementation + model_config.is_lora = model_args.lora + + logger.info(f"Final model config: {model_config}") + logger.info("Loading VL model...") + + model_class = AutoModelForConditionalGeneration + if model_args.continue_training and not training_args.autotuner_benchmark: + model = model_class.from_pretrained( + model_args.model_name_or_path, + config=model_config, + convert_from_hf=training_args.convert_from_hf, + load_via_cpu=training_args.load_via_cpu, + load_checkpoint_format=training_args.load_checkpoint_format, + ) + else: + model = model_class.from_config(model_config, dtype=dtype) + + # ====== Tokenizer & Processor ====== + tokenizer = AutoTokenizer.from_pretrained(model_args.model_name_or_path) + add_new_special_tokens(tokenizer, data_args.new_special_tokens_path) + if tokenizer.pad_token_id is None: + tokenizer.pad_token_id = tokenizer.eos_token_id + + processor = AutoProcessor.from_pretrained(model_args.model_name_or_path) + + # ====== MM Plugin ====== + mm_plugin_name = _detect_mm_plugin_name(model_config, data_args) + logger.info(f"[VL-SFT-V2] Using mm_plugin: {mm_plugin_name}") + + image_token = getattr(processor, "image_token", "") + mm_plugin = get_mm_plugin( + name=mm_plugin_name, + image_token=image_token, + video_token=None, + audio_token=None, + ) + + # ====== Freeze vision/language (optional) ====== + freeze_config = getattr(model_args, "freeze_config", None) + if freeze_config: + from paddleformers.cli.utils.mllm_utils import freeze_model_parameters + + freeze_model_parameters(model, model_config, freeze_config) + logger.info(f"[VL-SFT-V2] Applied freeze config: {freeze_config}") + + # ====== Dataset (datasets_v2) ====== + logger.info("[VL-SFT-V2] Loading dataset...") + + data_args._model_name_or_path = model_args.model_name_or_path + + # Determine template + template_name = _detect_template(tokenizer, data_args) + use_jinja = template_name == "__jinja__" + + if use_jinja: + template = None + logger.info("[VL-SFT-V2] Using tokenizer's built-in chat_template (jinja)") + else: + template = get_template(template_name) + logger.info(f"[VL-SFT-V2] Using template: {template_name}") + + encode_config = EncodeConfig( + max_seq_len=data_args.max_seq_len, + truncation="right", + label_shift=True, + ) + + encode_fn = partial( + encode_vl_sft, + tokenizer=tokenizer, + template=template, + config=encode_config, + processor=processor, + mm_plugin=mm_plugin, + ) + + # Load train dataset + train_dataset = None + eval_dataset = None + num_proc = getattr(data_args, "dataset_num_proc", 1) + + if training_args.do_train and training_args.should_load_dataset: + hf_ds = v2_load_dataset(data_args.train_dataset_path, num_proc=num_proc) + + need_auto_split = training_args.do_eval and ( + not data_args.eval_dataset_path or data_args.eval_dataset_path == data_args.train_dataset_path + ) + + if need_auto_split: + from paddleformers.datasets_v2 import split_dataset + + hf_ds, hf_eval_ds = split_dataset(hf_ds, test_ratio=0.1, shuffle=True, seed=training_args.seed) + eval_dataset = LazyEncodeDataset(hf_eval_ds, encode_fn, seed=training_args.seed) + logger.info(f"[VL-SFT-V2] Auto-split: train={len(hf_ds)}, eval={len(hf_eval_ds)}") + + train_dataset = LazyEncodeDataset(hf_ds, encode_fn, seed=training_args.seed) + logger.info(f"[VL-SFT-V2] Train dataset loaded: {len(train_dataset)} samples") + + if training_args.do_eval and training_args.should_load_dataset and eval_dataset is None: + hf_eval_ds = v2_load_dataset(data_args.eval_dataset_path, num_proc=num_proc) + eval_dataset = LazyEncodeDataset(hf_eval_ds, encode_fn, seed=training_args.seed) + logger.info(f"[VL-SFT-V2] Eval dataset loaded: {len(eval_dataset)} samples") + + # ====== Collate Function ====== + use_startend = getattr(model_args, "use_attn_mask_startend_row_indices", True) + get_rope_func = _get_rope_func(model) + + data_collator = partial( + collate_vl_sft, + pad_token_id=tokenizer.pad_token_id, + max_seq_len=data_args.max_seq_len, + get_rope_func=get_rope_func, + use_attn_mask_startend_row_indices=use_startend, + ) + + # ====== LoRA (optional) ====== + if model_args.lora: + from paddleformers.cli.utils import get_lora_target_modules + + if model_args.lora_path is None: + target_modules = get_lora_target_modules(model) + lora_config = LoRAConfig( + target_modules=target_modules, + r=model_args.lora_rank, + lora_alpha=2 * model_args.lora_rank if not model_args.rslora else 4, + rslora=model_args.rslora, + lora_plus_scale=model_args.lora_plus_scale, + merge_weights=False, + tensor_model_parallel_size=training_args.tensor_model_parallel_size, + dtype=dtype, + base_model_name_or_path=model_args.model_name_or_path, + ) + model = LoRAModel(model, lora_config) + else: + model = LoRAModel.from_pretrained(model=model, lora_path=model_args.lora_path) + + # ====== max_steps calculation ====== + if training_args.max_steps == -1: + if training_args.should_load_dataset and paddle.distributed.get_rank() == 0: + training_args.max_steps = math.ceil(len(train_dataset) / training_args.global_batch_size) + training_args.max_steps *= training_args.num_train_epochs + logger.info( + f"len(train_dataset): {len(train_dataset)}, " + f"global_batch_size: {training_args.global_batch_size}, " + f"num_train_epochs: {training_args.num_train_epochs}, " + f"max_steps: {training_args.max_steps}" + ) + + if paddle.distributed.get_world_size() > 1: + paddle.distributed.barrier() + max_steps = paddle.to_tensor([training_args.max_steps]) + paddle.distributed.broadcast(max_steps, src=0) + training_args.max_steps = int(max_steps.item()) + + if training_args.max_steps <= 0: + raise ValueError(f"Invalid max_steps: {training_args.max_steps}. Please check your dataset") + logger.info(f"Re-setting training_args.max_steps to {training_args.max_steps}.") + + if training_args.decay_steps is None: + training_args.decay_steps = training_args.max_steps + + if training_args.save_strategy == IntervalStrategy.EPOCH: + training_args.save_strategy = IntervalStrategy.STEPS + training_args.save_steps = int(training_args.max_steps / training_args.num_train_epochs) + if training_args.evaluation_strategy == IntervalStrategy.EPOCH: + training_args.evaluation_strategy = IntervalStrategy.STEPS + training_args.eval_steps = int(training_args.max_steps / training_args.num_train_epochs) + if training_args.logging_strategy == IntervalStrategy.EPOCH: + training_args.logging_strategy = IntervalStrategy.STEPS + training_args.logging_steps = int(training_args.max_steps / training_args.num_train_epochs) + + # ====== Create Trainer ====== + trainer = SFTTrainer( + model=model, + args=training_args, + train_dataset=(train_dataset if training_args.do_train and training_args.should_load_dataset else None), + eval_dataset=(eval_dataset if training_args.do_eval and training_args.should_load_dataset else None), + tokenizer=tokenizer, + data_collator=data_collator, + do_generation=False, + data_args=data_args, + ) + + trainable_parameters = [ + p for p in model.parameters() if not p.stop_gradient or ("quantization_linear" in p.name and "w_1" in p.name) + ] + trainer.set_optimizer_grouped_parameters(trainable_parameters) + + # ====== Train ====== + if training_args.do_train: + checkpoint = None + if training_args.resume_from_checkpoint is not None: + checkpoint = training_args.resume_from_checkpoint + elif last_checkpoint is not None: + checkpoint = last_checkpoint + train_result = trainer.train(resume_from_checkpoint=checkpoint) + + if getattr(training_args, "benchmark", False): + total_tokens = ( + data_args.max_seq_len + * training_args.per_device_train_batch_size + * training_args.dataset_world_size + * training_args.gradient_accumulation_steps + * training_args.max_steps + ) + total_tokens_per_second_per_gpu = ( + total_tokens / train_result.metrics["train_runtime"] / training_args.world_size + ) + logger.info(f"Total_Tokens_per_second_per_gpu: {total_tokens_per_second_per_gpu}") + logger.info("Benchmark done.") + else: + if not training_args.autotuner_benchmark: + trainer.save_model( + merge_tensor_parallel=training_args.tensor_model_parallel_size > 1, last_fc_to_hf=True + ) + trainer.log_metrics("train", train_result.metrics) + trainer.save_metrics("train", train_result.metrics) + trainer.save_state() diff --git a/paddleformers/cli/train/tuner.py b/paddleformers/cli/train/tuner.py index 0852562261f..c47e33369e1 100644 --- a/paddleformers/cli/train/tuner.py +++ b/paddleformers/cli/train/tuner.py @@ -51,9 +51,24 @@ def _training_function(config: dict[str, Any]) -> None: if model_args.stage in ["SFT", "PT", "VL-SFT", "VL-PT"]: with paddle.amp.auto_cast(enable=False): run_sft(model_args, data_args, generating_args, finetuning_args) + elif model_args.stage in ["SFT-V2", "PT-V2"]: + from .sft import run_sft_v2 + + with paddle.amp.auto_cast(enable=False): + run_sft_v2(model_args, data_args, generating_args, finetuning_args) + elif model_args.stage in ["VL-SFT-V2"]: + from .sft import run_vl_sft_v2 + + with paddle.amp.auto_cast(enable=False): + run_vl_sft_v2(model_args, data_args, generating_args, finetuning_args) elif model_args.stage == "DPO" or model_args.stage == "VL-DPO": with paddle.amp.auto_cast(enable=False): run_dpo(model_args, data_args, generating_args, finetuning_args) + elif model_args.stage == "DPO-V2": + from .dpo import run_dpo_v2 + + with paddle.amp.auto_cast(enable=False): + run_dpo_v2(model_args, data_args, generating_args, finetuning_args) elif model_args.stage == "dsv3_pretrain": from .deepseek_v3_pretrain import run_dsv3_pretrain diff --git a/paddleformers/datasets_v2/__init__.py b/paddleformers/datasets_v2/__init__.py new file mode 100644 index 00000000000..64dab8426bf --- /dev/null +++ b/paddleformers/datasets_v2/__init__.py @@ -0,0 +1,91 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .loaders import load_dataset, load_datasets +from .ops import ( + concat_datasets, + interleave, + sample_dataset, + shuffle_dataset, + split_dataset, +) +from .preprocessors import ( + AlpacaPreprocessor, + AutoPreprocessor, + BasePreprocessor, + MessagesPreprocessor, + ResponsePreprocessor, + TextPreprocessor, +) + +# Load built-in dataset registry on import +from .registry import ( + DatasetMeta, + DatasetSpec, + _load_builtin_registry, + get_dataset_meta, + list_datasets, + parse_dataset_string, + register_dataset, + register_dataset_info, +) +from .schema import ( + PAIR_KEYS, + ROLES, + STANDARD_KEYS, + GroundingObjects, + ImageMedia, + Message, + StandardRow, + cast_images, + cast_media_list, + check_messages, + remove_non_standard_keys, +) + +_load_builtin_registry() + +from .datapipe import ( + DPOEncodeConfig, + DPOEncodedSample, + EncodeConfig, + EncodedSample, + FunctionCall, + TemplateMeta, + VLEncodedSample, + collate_dpo, + collate_sft, + collate_vl_sft, + encode_dpo, + encode_multiturn, + encode_multiturn_reasoning, + encode_pt, + encode_sft, + encode_vl_sft, + fix_special_tokens, + get_template, + get_template_and_fix_tokenizer, + get_tool_utils, + greedy_pack, + list_templates, + parse_template, + register_template, +) +from .dataset import LazyEncodeDataset, StreamingDataset +from .grounding_plugin import ( + BaseGroundingPlugin, + get_grounding_plugin, + register_grounding_plugin, +) +from .mm_plugin import get_mm_plugin, register_mm_plugin diff --git a/paddleformers/datasets_v2/augment_utils.py b/paddleformers/datasets_v2/augment_utils.py new file mode 100644 index 00000000000..69f2f778e99 --- /dev/null +++ b/paddleformers/datasets_v2/augment_utils.py @@ -0,0 +1,103 @@ +# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Image processor class for PaddleOCR-VL.""" + +import io +import random + +from paddle.vision import transforms +from PIL import Image, ImageOps + + +class RandomApply: + def __init__(self, transforms, p=0.5): + self.transforms = transforms + self.p = p + + def __call__(self, x): + if random.random() < self.p: + for t in self.transforms: + x = t(x) + return x + + +class RandomDiscreteRotation: + def __init__(self, degrees, interpolation="nearest", expand=True): + self.degrees = degrees + self.interpolation = interpolation + self.expand = expand + + def __call__(self, img): + angle = random.choice(self.degrees) + return img.rotate(angle, self.interpolation, self.expand) + + +class JpegCompression: + def __init__(self, quality_range=(20, 80)): + self.quality_range = quality_range + + def __call__(self, img): + quality = random.randint(self.quality_range[0], self.quality_range[1]) + output = io.BytesIO() + img.convert("RGB").save(output, "JPEG", quality=quality) + output.seek(0) + return Image.open(output) + + +class RandomScale: + def __init__(self, scale_range=(0.7, 1.3), interpolation="bicubic"): + self.scale_range = scale_range + self.interpolation = interpolation + + def __call__(self, img): + scale = random.uniform(self.scale_range[0], self.scale_range[1]) + + original_width, original_height = img.size + new_width = int(original_width * scale) + new_height = int(original_height * scale) + new_size = (new_height, new_width) # transforms.Resize需要 (h, w) + + return transforms.functional.resize(img, new_size, self.interpolation) + + +class RandomSingleSidePadding: + def __init__(self, padding_range=(0, 20), fill="white"): + assert ( + isinstance(padding_range, (tuple, list)) and len(padding_range) == 2 + ), "padding_range must be the tuple or list like (min, max)" + self.min_pad, self.max_pad = padding_range + self.fill = fill + + def __call__(self, img): + + pad_amount = random.randint(self.min_pad, self.max_pad) + if pad_amount == 0: + return img + + chosen_edge = random.choice(["left", "top", "right", "bottom"]) + + pad_left, pad_top, pad_right, pad_bottom = 0, 0, 0, 0 + + if chosen_edge == "left": + pad_left = pad_amount + elif chosen_edge == "top": + pad_top = pad_amount + elif chosen_edge == "right": + pad_right = pad_amount + else: # 'bottom' + pad_bottom = pad_amount + + padding = (pad_left, pad_top, pad_right, pad_bottom) + return ImageOps.expand(img, border=padding, fill=self.fill) diff --git a/paddleformers/datasets_v2/datapipe/__init__.py b/paddleformers/datasets_v2/datapipe/__init__.py new file mode 100644 index 00000000000..5bab5c71994 --- /dev/null +++ b/paddleformers/datasets_v2/datapipe/__init__.py @@ -0,0 +1,43 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""datapipe: template encoding + packing + collation for datasets_v2.""" + +from .collate import collate_dpo, collate_sft, collate_vl_sft +from .encode import ( + DPOEncodeConfig, + DPOEncodedSample, + EncodeConfig, + EncodedSample, + VLEncodedSample, + encode_dpo, + encode_pt, + encode_sft, + encode_vl_sft, +) +from .packing import binpack_ffd, greedy_pack +from .template import ( + Slot, + TemplateMeta, + encode_multiturn, + encode_multiturn_jinja, + encode_multiturn_reasoning, + fix_special_tokens, + get_template, + get_template_and_fix_tokenizer, + list_templates, + parse_template, + register_template, +) +from .tool_utils import FunctionCall, ToolUtils, get_tool_utils diff --git a/paddleformers/datasets_v2/datapipe/collate.py b/paddleformers/datasets_v2/datapipe/collate.py new file mode 100644 index 00000000000..dc8e6a84996 --- /dev/null +++ b/paddleformers/datasets_v2/datapipe/collate.py @@ -0,0 +1,779 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Collation: batch of EncodedSamples → padded numpy dict for training. + +Supports: + - Non-packed (simple padding) and packed (greedy/binpacking with + block-diagonal attention mask) modes. + - MTP (Multi-Token Prediction) outputs: nbatch_pack_offset, + mtp_attn_mask / mtp_attn_mask_startend_row_indices, mtp_layer_mask. + +Output format is compatible with PaddleFormers' existing Trainer. +""" + +from typing import Any, Callable, Dict, List, Optional + +import numpy as np + +from .encode import EncodedSample +from .packing import binpack_ffd, greedy_pack + + +def collate_sft( + batch: List[EncodedSample], + pad_token_id: int, + max_seq_len: int, + packing: bool = False, + packing_method: str = "greedy", + packing_interval: int = 1000, + use_attn_mask_startend_row_indices: bool = False, + num_nextn_predict_layers: int = 0, + eos_token_id: Optional[int] = None, + use_global_causal_attn: bool = False, +) -> Dict[str, np.ndarray]: + """Collate a batch of EncodedSamples into training-ready numpy arrays. + + Args: + batch: List of EncodedSample from DataLoader. + pad_token_id: Tokenizer's pad token ID for input_ids padding. + max_seq_len: Maximum sequence length. Sequences are padded to + min(max_seq_len, max_length_in_batch) when not packing, + or exactly max_seq_len when packing. + packing: If True, use packing to combine short sequences. + packing_method: "greedy" (default) or "binpacking" (FFD). + packing_interval: Buffer size for binpack_ffd algorithm. + use_attn_mask_startend_row_indices: If True, output compact flashmask + format (attn_mask_startend_row_indices) instead of 4D attention_mask. + num_nextn_predict_layers: MTP depth D. When > 0, max_seq_len is extended + by D and MTP outputs are generated. + eos_token_id: Tokenizer's EOS token ID for mtp_layer_mask. + use_global_causal_attn: If True, MTP masks use single global causal block + instead of per-document block-causal. + + Returns: + Dict with numpy arrays: + - input_ids: [B, S] int64 + - labels: [B, S] int64, -100 for masked positions + - position_ids: [B, S] int64 + - attention_mask: [B, 1, S, S] float32 (when not using startend) + - attn_mask_startend_row_indices: [B, 1, S, 1] int32 (when using startend) + When num_nextn_predict_layers > 0, additionally: + - nbatch_pack_offset: [B, S] int32 + - mtp_attn_mask: [B, D, 1, S, S] float32 (or startend variant) + - mtp_layer_mask: [B, D, S] int32 + """ + # Handle dict input from streaming .map() — convert to EncodedSample + if batch and isinstance(batch[0], dict): + batch = [ + EncodedSample( + input_ids=item["input_ids"], + labels=item["labels"], + seq_len=item["seq_len"], + ) + for item in batch + ] + + if not batch: + raise ValueError("collate_sft received an empty batch") + + # Effective max_seq_len (extended for MTP) + effective_max_seq_len = max_seq_len + if num_nextn_predict_layers > 0: + effective_max_seq_len = max_seq_len + num_nextn_predict_layers + + if packing: + # Select packing algorithm + if packing_method == "binpacking": + groups = binpack_ffd(batch, max_seq_len, packing_interval) + else: + groups = greedy_pack(batch, max_seq_len) + return _collate_packed( + groups, + pad_token_id, + effective_max_seq_len, + use_attn_mask_startend_row_indices, + num_nextn_predict_layers, + eos_token_id, + use_global_causal_attn, + ) + else: + return _collate_simple( + batch, + pad_token_id, + effective_max_seq_len, + use_attn_mask_startend_row_indices, + num_nextn_predict_layers, + eos_token_id, + use_global_causal_attn, + ) + + +def _collate_simple( + batch: List[EncodedSample], + pad_token_id: int, + max_seq_len: int, + use_startend: bool = False, + num_nextn_predict_layers: int = 0, + eos_token_id: Optional[int] = None, + use_global_causal_attn: bool = False, +) -> Dict[str, np.ndarray]: + """Simple collation: pad each sample independently, causal mask.""" + batch_size = len(batch) + pad_len = min(max_seq_len, max(s.seq_len for s in batch)) + + input_ids = np.full((batch_size, pad_len), pad_token_id, dtype=np.int64) + labels = np.full((batch_size, pad_len), -100, dtype=np.int64) + position_ids = np.zeros((batch_size, pad_len), dtype=np.int64) + + for i, sample in enumerate(batch): + seq_len = min(sample.seq_len, pad_len) + input_ids[i, :seq_len] = sample.input_ids[:seq_len] + labels[i, :seq_len] = sample.labels[:seq_len] + position_ids[i, :seq_len] = np.arange(seq_len) + + result = { + "input_ids": input_ids, + "labels": labels, + "position_ids": position_ids, + } + + if use_startend: + result["attn_mask_startend_row_indices"] = _build_startend_simple(batch, pad_len) + else: + attention_mask = np.zeros((batch_size, 1, pad_len, pad_len), dtype=np.float32) + for i, sample in enumerate(batch): + seq_len = min(sample.seq_len, pad_len) + attention_mask[i, 0, :seq_len, :seq_len] = np.tril(np.ones((seq_len, seq_len))) + result["attention_mask"] = attention_mask + + # MTP outputs (single-sequence case: each sample is one "document") + if num_nextn_predict_layers > 0: + all_pack_offset = [] + all_mtp_mask = [] + all_mtp_layer_mask = [] + for i, sample in enumerate(batch): + seq_len = min(sample.seq_len, pad_len) + seq_lens_i = [seq_len] # single document + all_pack_offset.append(_build_nbatch_pack_offset(seq_lens_i, pad_len)) + if use_startend: + all_mtp_mask.append( + gen_mtp_attn_mask_startend_row_indices( + seq_lens_i, pad_len, num_nextn_predict_layers, use_global_causal_attn + ) + ) + else: + all_mtp_mask.append( + gen_mtp_attn_mask(seq_lens_i, pad_len, num_nextn_predict_layers, use_global_causal_attn) + ) + all_mtp_layer_mask.append( + gen_mtp_layer_mask(input_ids[i], seq_len, pad_len, num_nextn_predict_layers, eos_token_id) + ) + + result["nbatch_pack_offset"] = np.concatenate(all_pack_offset, axis=0) + if use_startend: + result["mtp_attn_mask_startend_row_indices"] = np.stack(all_mtp_mask) + else: + result["mtp_attn_mask"] = np.stack(all_mtp_mask) + result["mtp_layer_mask"] = np.stack(all_mtp_layer_mask) + + return result + + +def _collate_packed( + groups: List[List[EncodedSample]], + pad_token_id: int, + max_seq_len: int, + use_startend: bool = False, + num_nextn_predict_layers: int = 0, + eos_token_id: Optional[int] = None, + use_global_causal_attn: bool = False, +) -> Dict[str, np.ndarray]: + """Packed collation: groups already formed, build block-diagonal attention mask.""" + batch_size = len(groups) + + input_ids = np.full((batch_size, max_seq_len), pad_token_id, dtype=np.int64) + labels = np.full((batch_size, max_seq_len), -100, dtype=np.int64) + position_ids = np.zeros((batch_size, max_seq_len), dtype=np.int64) + + # Track sub-sequence lengths for each group + group_seq_lens: List[List[int]] = [] + + for i, group in enumerate(groups): + offset = 0 + seq_lens = [] + for sample in group: + seq_len = min(sample.seq_len, max_seq_len - offset) + if seq_len <= 0: + break + + input_ids[i, offset : offset + seq_len] = sample.input_ids[:seq_len] + labels[i, offset : offset + seq_len] = sample.labels[:seq_len] + position_ids[i, offset : offset + seq_len] = np.arange(seq_len) + + seq_lens.append(seq_len) + offset += seq_len + group_seq_lens.append(seq_lens) + + result = { + "input_ids": input_ids, + "labels": labels, + "position_ids": position_ids, + } + + if use_startend: + result["attn_mask_startend_row_indices"] = _build_startend_packed(group_seq_lens, max_seq_len) + else: + attention_mask = np.zeros((batch_size, 1, max_seq_len, max_seq_len), dtype=np.float32) + for i, seq_lens in enumerate(group_seq_lens): + offset = 0 + for seq_len in seq_lens: + causal_block = np.tril(np.ones((seq_len, seq_len), dtype=np.float32)) + attention_mask[i, 0, offset : offset + seq_len, offset : offset + seq_len] = causal_block + offset += seq_len + result["attention_mask"] = attention_mask + + # MTP outputs + if num_nextn_predict_layers > 0: + all_pack_offset = [] + all_mtp_mask = [] + all_mtp_layer_mask = [] + + for i, seq_lens in enumerate(group_seq_lens): + all_pack_offset.append(_build_nbatch_pack_offset(seq_lens, max_seq_len)) + + if use_startend: + all_mtp_mask.append( + gen_mtp_attn_mask_startend_row_indices( + seq_lens, max_seq_len, num_nextn_predict_layers, use_global_causal_attn + ) + ) + else: + all_mtp_mask.append( + gen_mtp_attn_mask(seq_lens, max_seq_len, num_nextn_predict_layers, use_global_causal_attn) + ) + + total_len = sum(seq_lens) + all_mtp_layer_mask.append( + gen_mtp_layer_mask(input_ids[i], total_len, max_seq_len, num_nextn_predict_layers, eos_token_id) + ) + + result["nbatch_pack_offset"] = np.concatenate(all_pack_offset, axis=0) + if use_startend: + result["mtp_attn_mask_startend_row_indices"] = np.stack(all_mtp_mask) + else: + result["mtp_attn_mask"] = np.stack(all_mtp_mask) + result["mtp_layer_mask"] = np.stack(all_mtp_layer_mask) + + return result + + +# --------------------------------------------------------------------------- +# Attention mask helpers +# --------------------------------------------------------------------------- + + +def _build_startend_simple( + batch: List[EncodedSample], + pad_len: int, +) -> np.ndarray: + """Build attn_mask_startend_row_indices for non-packed batch. + + Each sample is a single sequence: all tokens attend to [0, seq_len). + Padding tokens get identity (each attends only to itself). + + Returns: [B, 1, S, 1] int32 + """ + batch_size = len(batch) + indices = np.zeros((batch_size, 1, pad_len, 1), dtype=np.int32) + for i, sample in enumerate(batch): + seq_len = min(sample.seq_len, pad_len) + indices[i, 0, :seq_len, 0] = seq_len + indices[i, 0, seq_len:, 0] = np.arange(seq_len, pad_len) + return indices + + +def _build_startend_packed( + group_seq_lens: List[List[int]], + max_seq_len: int, +) -> np.ndarray: + """Build attn_mask_startend_row_indices for packed batch. + + For each group, each sub-sequence of length L at offset O: + all L positions get value O+L (meaning "can attend up to position O+L-1"). + Padding positions get identity (self-attend only). + + Returns: [B, 1, S, 1] int32 + """ + batch_size = len(group_seq_lens) + indices = np.zeros((batch_size, 1, max_seq_len, 1), dtype=np.int32) + for i, seq_lens in enumerate(group_seq_lens): + offset = 0 + for seq_len in seq_lens: + indices[i, 0, offset : offset + seq_len, 0] = offset + seq_len + offset += seq_len + # Padding region: each position attends only to itself + if offset < max_seq_len: + indices[i, 0, offset:, 0] = np.arange(offset, max_seq_len) + return indices + + +# --------------------------------------------------------------------------- +# MTP (Multi-Token Prediction) helpers +# --------------------------------------------------------------------------- + + +def _build_nbatch_pack_offset( + group_seq_lens: List[int], + max_seq_len: int, +) -> np.ndarray: + """Build document boundary markers for MTP. + + Places 1 at the last position of each sub-sequence (except the final one). + This tells MTP layers where document boundaries are within a packed group. + + Args: + group_seq_lens: List of sub-sequence lengths in one packed group. + max_seq_len: Padded sequence length. + + Returns: [1, S] int32 + """ + offset_arr = np.zeros(max_seq_len, dtype=np.int32) + pos = 0 + for seq_len in group_seq_lens[:-1]: + pos += seq_len + offset_arr[pos - 1] = 1 + return offset_arr[None, :] # [1, S] + + +def gen_mtp_attn_mask( + group_seq_lens: List[int], + max_seq_len: int, + mtp_depth: int, + use_global_causal_attn: bool = False, +) -> np.ndarray: + """Generate MTP per-layer attention mask (2D matrix form). + + For each MTP layer d (0-indexed), document boundaries are shifted left + by d+1 positions. This creates the correct causal mask for each + speculative prediction depth. + + Args: + group_seq_lens: Sub-sequence lengths in one packed group. + max_seq_len: Padded sequence length (already extended by mtp_depth). + mtp_depth: Number of MTP prediction layers D. + use_global_causal_attn: If True, single global causal block. + + Returns: [D, 1, S, S] float32 + """ + total_len = sum(group_seq_lens) + + if use_global_causal_attn: + single = np.zeros((max_seq_len, max_seq_len), dtype=np.float32) + single[:total_len, :total_len] = np.tril(np.ones((total_len, total_len))) + return np.stack([single] * mtp_depth, axis=0)[:, None, :, :] + + # Compute internal boundaries (cumulative sum of seq_lens, excluding last) + boundaries = [] + offset = 0 + for sl in group_seq_lens[:-1]: + offset += sl + boundaries.append(offset) + + result = [] + for mtp_idx in range(mtp_depth): + shift = mtp_idx + 1 + shifted = [b - shift for b in boundaries if b - shift > 0] + [total_len] + mask = np.zeros((max_seq_len, max_seq_len), dtype=np.float32) + prev = 0 + for boundary in shifted: + if boundary > prev: + mask[prev:boundary, prev:boundary] = np.tril( + np.ones((boundary - prev, boundary - prev), dtype=np.float32) + ) + prev = boundary + result.append(mask) + + return np.stack(result, axis=0)[:, None, :, :] + + +def gen_mtp_attn_mask_startend_row_indices( + group_seq_lens: List[int], + max_seq_len: int, + mtp_depth: int, + use_global_causal_attn: bool = False, +) -> np.ndarray: + """Generate MTP per-layer attention mask (compressed startend form). + + Args: + group_seq_lens: Sub-sequence lengths in one packed group. + max_seq_len: Padded sequence length (already extended by mtp_depth). + mtp_depth: Number of MTP prediction layers D. + use_global_causal_attn: If True, single global block. + + Returns: [D, 1, S, 1] int32 + """ + total_len = sum(group_seq_lens) + pad_indices = list(range(total_len, max_seq_len)) + + if use_global_causal_attn: + row = [total_len] * total_len + pad_indices + return np.array([row] * mtp_depth, dtype=np.int32)[:, None, :, None] + + boundaries = [] + offset = 0 + for sl in group_seq_lens[:-1]: + offset += sl + boundaries.append(offset) + + result = [] + for mtp_idx in range(mtp_depth): + shift = mtp_idx + 1 + shifted = [b - shift for b in boundaries if b - shift > 0] + [total_len] + indices = [] + prev = 0 + for boundary in shifted: + indices.extend([boundary] * (boundary - prev)) + prev = boundary + result.append(indices + pad_indices) + + return np.array(result, dtype=np.int32)[:, None, :, None] + + +def gen_mtp_layer_mask( + input_ids_row: np.ndarray, + total_len: int, + max_seq_len: int, + mtp_depth: int, + eos_token_id: Optional[int] = None, +) -> np.ndarray: + """Generate MTP per-layer hidden inputs mask. + + Zeroes out positions where EOS token appears at the shifted offset, + preventing MTP layers from predicting across document boundaries. + + Args: + input_ids_row: The input_ids for one batch row, shape [S]. + total_len: Actual total token length (sum of sub-sequences). + max_seq_len: Padded sequence length. + mtp_depth: Number of MTP prediction layers D. + eos_token_id: If provided, zero out EOS positions in shifted input. + + Returns: [D, S] int32 + """ + if eos_token_id is None: + return np.ones((mtp_depth, max_seq_len), dtype=np.int32) + + result = [] + for mtp_idx in range(mtp_depth): + mask = np.ones(max_seq_len, dtype=np.int32) + shifted = input_ids_row[mtp_idx + 1 : total_len] + eos_positions = np.where(shifted == eos_token_id)[0] + mask[eos_positions] = 0 + result.append(mask) + + return np.stack(result, axis=0) + + +# --------------------------------------------------------------------------- +# VL (Vision-Language) collation +# --------------------------------------------------------------------------- + + +def collate_vl_sft( + batch: List[Any], + pad_token_id: int, + max_seq_len: int, + get_rope_func: Optional[Callable] = None, + use_attn_mask_startend_row_indices: bool = False, +) -> Dict[str, Any]: + """Collate a batch of VLEncodedSamples into training-ready arrays. + + Handles pixel_values concatenation and 3D position_ids via get_rope_func. + + Args: + batch: List of VLEncodedSample from DataLoader. + pad_token_id: Tokenizer's pad token ID. + max_seq_len: Maximum sequence length for padding. + get_rope_func: Model's get_rope_index method for 3D position_ids. + If None, falls back to simple 1D position_ids. + use_attn_mask_startend_row_indices: If True, use compact flashmask format. + + Returns: + Dict with numpy/tensor arrays: + - input_ids, labels, position_ids, attention_mask (same as collate_sft) + - pixel_values: concatenated vision features + - image_grid_thw: stacked grid dimensions + """ + import paddle + + batch_size = len(batch) + pad_len = min(max_seq_len, max(s.seq_len for s in batch)) + + input_ids = np.full((batch_size, pad_len), pad_token_id, dtype=np.int64) + labels = np.full((batch_size, pad_len), -100, dtype=np.int64) + + for i, sample in enumerate(batch): + seq_len = min(sample.seq_len, pad_len) + input_ids[i, :seq_len] = sample.input_ids[:seq_len] + labels[i, :seq_len] = sample.labels[:seq_len] + + # Collect vision inputs + all_pixel_values = [] + all_image_grid_thw = [] + + for sample in batch: + mm = sample.mm_inputs + if "pixel_values" in mm and mm["pixel_values"] is not None: + pv = mm["pixel_values"] + if not isinstance(pv, paddle.Tensor): + pv = paddle.to_tensor(pv) + all_pixel_values.append(pv) + if "image_grid_thw" in mm and mm["image_grid_thw"] is not None: + grid = mm["image_grid_thw"] + if not isinstance(grid, paddle.Tensor): + grid = paddle.to_tensor(grid, dtype="int64") + all_image_grid_thw.append(grid) + + pixel_values = paddle.concat(all_pixel_values, axis=0) if all_pixel_values else None + image_grid_thw = paddle.concat(all_image_grid_thw, axis=0) if all_image_grid_thw else None + + # Compute position_ids + if get_rope_func is not None and image_grid_thw is not None: + input_ids_tensor = paddle.to_tensor(input_ids, dtype="int64") + attention_mask_tensor = paddle.ones_like(input_ids_tensor) + for i, sample in enumerate(batch): + seq_len = min(sample.seq_len, pad_len) + attention_mask_tensor[i, seq_len:] = 0 + + position_ids, rope_deltas = get_rope_func( + input_ids=input_ids_tensor, + image_grid_thw=image_grid_thw, + video_grid_thw=None, + attention_mask=attention_mask_tensor, + ) + position_ids = position_ids.numpy() + else: + position_ids = np.zeros((batch_size, pad_len), dtype=np.int64) + for i, sample in enumerate(batch): + seq_len = min(sample.seq_len, pad_len) + position_ids[i, :seq_len] = np.arange(seq_len) + + result: Dict[str, Any] = { + "input_ids": input_ids, + "labels": labels, + "position_ids": position_ids, + } + + if pixel_values is not None: + result["pixel_values"] = pixel_values + if image_grid_thw is not None: + result["image_grid_thw"] = image_grid_thw + + if use_attn_mask_startend_row_indices: + result["attn_mask_startend_row_indices"] = _build_startend_simple(batch, pad_len) + else: + attention_mask = np.zeros((batch_size, 1, pad_len, pad_len), dtype=np.float32) + for i, sample in enumerate(batch): + seq_len = min(sample.seq_len, pad_len) + attention_mask[i, 0, :seq_len, :seq_len] = np.tril(np.ones((seq_len, seq_len))) + result["attention_mask"] = attention_mask + + return result + + +# --------------------------------------------------------------------------- +# DPO collation +# --------------------------------------------------------------------------- + + +def collate_dpo( + batch: List[Any], + pad_token_id: int, + max_seq_len: int, + use_attn_mask_startend_row_indices: bool = False, + use_filtered_label_loss: bool = True, +) -> Dict[str, np.ndarray]: + """Collate a batch of DPOEncodedSamples into training-ready numpy arrays. + + Produces the format expected by DPOTrainer.compute_loss(): + - Block-causal attention mask (rejected cannot see chosen tokens) + - response_indexs for slicing chosen/rejected log-probs + + Args: + batch: List of DPOEncodedSample (or dicts from streaming). + pad_token_id: Tokenizer's pad token ID. + max_seq_len: Maximum sequence length for padding. + use_attn_mask_startend_row_indices: If True, use compact flashmask format. + use_filtered_label_loss: Controls response_indexs computation mode. + + Returns: + Dict with numpy arrays: + - input_ids: [B, S] int64 + - position_ids: [B, S] int64 + - response_labels: [B, S] int64 + - response_indexs: [N, 4] int32 where N=batch_size + - attention_mask: [B, 1, S, S] float32 (or attn_mask_startend_row_indices) + - score_deltas: [N, 1] float32 + """ + from .encode import DPOEncodedSample + + # Handle dict input from streaming .map() + if batch and isinstance(batch[0], dict): + batch = [ + DPOEncodedSample( + input_ids=item["input_ids"], + position_ids=item["position_ids"], + response_labels=item["response_labels"], + response_index=item["response_index"], + seq_len=item["seq_len"], + score_delta=item.get("score_delta", 1.0), + ) + for item in batch + ] + + if not batch: + raise ValueError("collate_dpo received an empty batch") + + batch_size = len(batch) + pad_len = min(max_seq_len, max(s.seq_len for s in batch)) + + input_ids = np.full((batch_size, pad_len), pad_token_id, dtype=np.int64) + position_ids = np.zeros((batch_size, pad_len), dtype=np.int64) + response_labels = np.full((batch_size, pad_len), -100, dtype=np.int64) + + for i, sample in enumerate(batch): + seq_len = min(sample.seq_len, pad_len) + input_ids[i, :seq_len] = sample.input_ids[:seq_len] + position_ids[i, :seq_len] = sample.position_ids[:seq_len] + response_labels[i, :seq_len] = sample.response_labels[:seq_len] + + # Build response_indexs: [batch_idx, chosen_start, rejected_start, rejected_end] + response_indexs_list = [] + if use_filtered_label_loss: + # Absolute indices into flattened (filtered) logps + sequence_sum_flatten = 0 + for i, sample in enumerate(batch): + ri = sample.response_index # [chosen_start, rejected_start, rejected_end] + response_indexs_list.append( + [ + i, + sequence_sum_flatten + ri[0], + sequence_sum_flatten + ri[1], + sequence_sum_flatten + ri[2], + ] + ) + sequence_sum_flatten += ri[2] # rejected_end = total response length + else: + # Per-sequence relative indices + for i, sample in enumerate(batch): + ri = sample.response_index + response_indexs_list.append([i, ri[0], ri[1], ri[2]]) + + response_indexs = np.array(response_indexs_list, dtype=np.int32) + + # Score deltas + score_deltas = np.array([[s.score_delta] for s in batch], dtype=np.float32) + + result = { + "input_ids": input_ids, + "position_ids": position_ids, + "response_labels": response_labels, + "response_indexs": response_indexs, + "score_deltas": score_deltas, + } + + # Build attention mask (block-causal for DPO) + if use_attn_mask_startend_row_indices: + result["attn_mask_startend_row_indices"] = _build_dpo_startend(batch, pad_len) + else: + result["attention_mask"] = _build_dpo_attention_mask(batch, pad_len) + + return result + + +def _build_dpo_startend( + batch: List[Any], + pad_len: int, +) -> np.ndarray: + """Build attn_mask_startend_row_indices for DPO batch. + + DPO attention pattern: + - Prompt tokens: can attend to full sequence (value = seq_len) + - Chosen tokens: can attend up to end of chosen (value = prompt_len + chosen_len - 1) + - Fork token: can attend to full sequence (value = seq_len) + - Rejected tokens: can attend to full sequence (value = seq_len) + - Padding: self-attend only (value = position) + + Returns: [B, 1, S, 1] int32 + """ + batch_size = len(batch) + indices = np.zeros((batch_size, 1, pad_len, 1), dtype=np.int32) + + for i, sample in enumerate(batch): + seq_len = min(sample.seq_len, pad_len) + ri = sample.response_index + + # Derive structure lengths + chosen_len = ri[1] - ri[0] + rejected_len = ri[2] - ri[1] + prompt_len = seq_len - chosen_len - rejected_len + 1 + + # DPO startend pattern + indices[i, 0, :prompt_len, 0] = seq_len + chosen_end_pos = prompt_len + chosen_len - 1 + indices[i, 0, prompt_len:chosen_end_pos, 0] = chosen_end_pos + indices[i, 0, chosen_end_pos, 0] = seq_len # fork token + if chosen_end_pos + 1 < seq_len: + indices[i, 0, chosen_end_pos + 1 : seq_len, 0] = seq_len # rejected + + # Padding: each position attends only to itself + if seq_len < pad_len: + indices[i, 0, seq_len:, 0] = np.arange(seq_len, pad_len) + + return indices + + +def _build_dpo_attention_mask( + batch: List[Any], + pad_len: int, +) -> np.ndarray: + """Build 4D attention mask for DPO batch. + + Block-causal: lower triangular, but rejected tokens cannot attend to chosen tokens. + + Returns: [B, 1, S, S] float32 + """ + batch_size = len(batch) + attention_mask = np.zeros((batch_size, 1, pad_len, pad_len), dtype=np.float32) + + for i, sample in enumerate(batch): + seq_len = min(sample.seq_len, pad_len) + ri = sample.response_index + + chosen_len = ri[1] - ri[0] + rejected_len = ri[2] - ri[1] + prompt_len = seq_len - chosen_len - rejected_len + 1 + + # Start with full causal mask + attention_mask[i, 0, :seq_len, :seq_len] = np.tril(np.ones((seq_len, seq_len), dtype=np.float32)) + + # Block rejected tokens from seeing chosen tokens + # Fork token is at: prompt_len + chosen_len - 1 + # Rejected tokens: [fork_pos, seq_len) + # Chosen tokens occupy: [prompt_len - 1, fork_pos) + fork_pos = prompt_len + chosen_len - 1 + attention_mask[i, 0, fork_pos:seq_len, prompt_len - 1 : fork_pos] = 0.0 + + return attention_mask diff --git a/paddleformers/datasets_v2/datapipe/encode.py b/paddleformers/datasets_v2/datapipe/encode.py new file mode 100644 index 00000000000..2c99d54c56d --- /dev/null +++ b/paddleformers/datasets_v2/datapipe/encode.py @@ -0,0 +1,731 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Single-sample encoding: messages → input_ids + labels. + +Supports SFT, PT, VL-SFT, and DPO encoding modes. +""" + +import logging +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional, Set, Tuple + +from .template import ( + TemplateMeta, + encode_multiturn, + encode_multiturn_jinja, + encode_multiturn_reasoning, +) + +logger = logging.getLogger(__name__) + + +@dataclass +class EncodeConfig: + """Configuration for SFT encoding.""" + + max_seq_len: int = 4096 + truncation: str = "right" # "right" | "left" | "oral" | "delete" + label_shift: bool = True + auto_add_bos: bool = False + placeholder_tokens: List[int] = field(default_factory=list) + + +@dataclass +class EncodedSample: + """Output of encode_sft / encode_pt.""" + + input_ids: List[int] + labels: List[int] + seq_len: int + position_ids: List[int] = field(default_factory=list) + + +@dataclass +class VLEncodedSample: + """Output of encode_vl_sft.""" + + input_ids: List[int] + labels: List[int] + seq_len: int + mm_inputs: Dict[str, Any] + position_ids: List[int] = field(default_factory=list) + + +# --------------------------------------------------------------------------- +# Shared helpers +# --------------------------------------------------------------------------- + +_ERNIE_THINK_MODE_SYSTEM = "\nthink_mode=True\n" + + +def _inject_ernie_think_system(messages: List[Dict[str, Any]], template: TemplateMeta) -> List[Dict[str, Any]]: + if not template.name.startswith("ernie"): + return messages + if template.enable_thinking is None: + return messages + if messages and messages[0].get("role") == "system": + return messages + return [{"role": "system", "content": _ERNIE_THINK_MODE_SYSTEM}] + messages + + +def _extract_loss_mask(messages: List[Dict[str, Any]]) -> List[bool]: + mask: List[bool] = [] + for msg in messages: + if msg.get("role") == "assistant": + mask.append(msg.get("loss", True)) + return mask + + +def _dispatch_encode( + template: Optional[TemplateMeta], + tokenizer: Any, + messages: List[Dict[str, Any]], + tools: Optional[Any] = None, +) -> List[Tuple[List[int], List[int]]]: + if template is not None: + if template.enable_thinking is not None: + return encode_multiturn_reasoning(template, tokenizer, messages, tools=tools) + else: + return encode_multiturn(template, tokenizer, messages, tools=tools) + else: + return encode_multiturn_jinja(tokenizer, messages) + + +def _get_sep_token_len(template: Optional[TemplateMeta], tokenizer: Any) -> int: + if template is not None and template.chat_sep: + return len(tokenizer.encode(template.chat_sep, add_special_tokens=False)) + return 0 + + +def _flatten_turns( + pairs: List[Tuple[List[int], List[int]]], + loss_mask: List[bool], + sep_token_len: int, +) -> Tuple[List[int], List[int]]: + """Flatten (prompt_ids, response_ids) pairs into token_ids + labels.""" + token_ids: List[int] = [] + labels: List[int] = [] + num_pairs = len(pairs) + + for turn_idx, (prompt_ids, response_ids) in enumerate(pairs): + token_ids += prompt_ids + labels += [-100] * len(prompt_ids) + + token_ids += response_ids + if turn_idx < len(loss_mask) and not loss_mask[turn_idx]: + labels += [-100] * len(response_ids) + else: + if sep_token_len > 0 and turn_idx != (num_pairs - 1) and len(response_ids) > sep_token_len: + labels += response_ids[: len(response_ids) - sep_token_len] + [-100] * sep_token_len + else: + labels += response_ids + + return token_ids, labels + + +def _add_dynamic_eos(input_ids: List[int], labels: List[int], suffix_tokens_id: List[int]) -> None: + suffix_len = len(suffix_tokens_id) + if suffix_len == 0: + return + + start = 0 + for i in range(1, len(labels) + 1): + if labels[i - 1] >= 0 and i < len(labels) and labels[i] == -100: + start = i + elif start > 0 and labels[i - 1] == -100 and (i == len(labels) or labels[i] >= 0): + length = i - start + if length >= suffix_len and input_ids[start : start + suffix_len] == suffix_tokens_id: + labels[start : start + suffix_len] = suffix_tokens_id + + +def _get_suffix_ids(template: Optional[TemplateMeta], tokenizer: Any) -> List[int]: + """Get suffix token IDs using tokenize+convert_tokens_to_ids (aligned with V1 SFTDataset).""" + if template is None or not template.suffix: + return [] + return tokenizer.convert_tokens_to_ids(tokenizer.tokenize(template.suffix[-1])) + + +def _apply_dynamic_eos(token_ids: List[int], labels: List[int], template: Optional[TemplateMeta], tokenizer: Any): + if template is not None and template.suffix: + suffix_ids = _get_suffix_ids(template, tokenizer) + _add_dynamic_eos(token_ids, labels, suffix_ids) + + +def _apply_efficient_eos(token_ids: List[int], labels: List[int], template: Optional[TemplateMeta], tokenizer: Any): + if template and template.efficient_eos: + if template.suffix: + suffix_ids = _get_suffix_ids(template, tokenizer) + token_ids.extend(suffix_ids) + labels.extend(suffix_ids) + elif tokenizer.eos_token_id is not None: + token_ids.append(tokenizer.eos_token_id) + labels.append(tokenizer.eos_token_id) + elif template is None: + # Jinja mode: append eos_token to align with V1's efficient_eos behavior + # V1's parse_template always sets efficient_eos=True with suffix=[tokenizer.eos_token] + if tokenizer.eos_token_id is not None: + token_ids.append(tokenizer.eos_token_id) + labels.append(tokenizer.eos_token_id) + + +def _apply_label_shift(labels: List[int], config: EncodeConfig) -> List[int]: + if config.label_shift: + return labels[1:] + [-100] + return labels + + +def _truncate_placeholder_aware( + input_ids: List[int], + labels: List[int], + max_seq_len: int, + placeholder_set: Set[int], + strategy: str = "right", +) -> Tuple[List[int], List[int]]: + is_placeholder = [tok in placeholder_set for tok in input_ids] + placeholder_idx = [i for i, v in enumerate(is_placeholder) if v] + + if len(placeholder_idx) >= max_seq_len: + keep_idx = placeholder_idx[:max_seq_len] + else: + remain = max_seq_len - len(placeholder_idx) + non_placeholder_idx = [i for i, v in enumerate(is_placeholder) if not v] + + if strategy == "left": + extra_idx = non_placeholder_idx[-remain:] + else: + extra_idx = non_placeholder_idx[:remain] + + keep_idx = sorted(placeholder_idx + extra_idx) + + input_ids = [input_ids[i] for i in keep_idx] + labels = [labels[i] for i in keep_idx] + return input_ids, labels + + +def _apply_truncation( + token_ids: List[int], + labels: List[int], + config: EncodeConfig, +) -> Optional[Tuple[List[int], List[int]]]: + """Apply truncation. Returns None if strategy is 'delete'.""" + if len(token_ids) <= config.max_seq_len: + return token_ids, labels + + if config.truncation == "delete": + return None + elif config.placeholder_tokens: + return _truncate_placeholder_aware( + token_ids, labels, config.max_seq_len, set(config.placeholder_tokens), config.truncation + ) + elif config.truncation == "left": + return token_ids[-config.max_seq_len :], labels[-config.max_seq_len :] + else: + return token_ids[: config.max_seq_len], labels[: config.max_seq_len] + + +def _apply_auto_bos( + token_ids: List[int], + labels: List[int], + config: EncodeConfig, + tokenizer: Any, +) -> Tuple[List[int], List[int]]: + if config.auto_add_bos and tokenizer.bos_token_id is not None: + if not token_ids or token_ids[0] != tokenizer.bos_token_id: + token_ids = [tokenizer.bos_token_id] + token_ids + labels = [-100] + labels + if len(token_ids) > config.max_seq_len: + token_ids = token_ids[: config.max_seq_len] + labels = labels[: config.max_seq_len] + return token_ids, labels + + +def _validate_and_build( + token_ids: List[int], + labels: List[int], +) -> Optional[Tuple[List[int], List[int], List[int]]]: + """Validate and return (token_ids, labels, position_ids) or None.""" + if not token_ids: + return None + if all(x == -100 for x in labels): + logger.warning("[SKIP] all labels set to -100") + return None + position_ids = list(range(len(token_ids))) + return token_ids, labels, position_ids + + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- + + +def encode_sft( + example: Dict[str, Any], + tokenizer: Any, + template: Optional[TemplateMeta], + config: EncodeConfig, +) -> Optional[EncodedSample]: + """Encode a single SFT sample: messages → token_ids + labels.""" + messages = example.get("messages") + if not messages or len(messages) < 2: + return None + + # Extract tools definition (for function calling) + tools = example.get("tools") + + if template is not None: + messages = _inject_ernie_think_system(messages, template) + + loss_mask = _extract_loss_mask(messages) + pairs = _dispatch_encode(template, tokenizer, messages, tools=tools) + if not pairs: + return None + + if config.truncation == "oral": + return _encode_oral_truncation(pairs, loss_mask, tokenizer, template, config) + + sep_token_len = _get_sep_token_len(template, tokenizer) + token_ids, labels = _flatten_turns(pairs, loss_mask, sep_token_len) + if not token_ids: + return None + + _apply_dynamic_eos(token_ids, labels, template, tokenizer) + _apply_efficient_eos(token_ids, labels, template, tokenizer) + + result = _apply_truncation(token_ids, labels, config) + if result is None: + return None + token_ids, labels = result + + labels = _apply_label_shift(labels, config) + + token_ids, labels = _apply_auto_bos(token_ids, labels, config, tokenizer) + + validated = _validate_and_build(token_ids, labels) + if validated is None: + return None + token_ids, labels, position_ids = validated + return EncodedSample(input_ids=token_ids, labels=labels, seq_len=len(token_ids), position_ids=position_ids) + + +def _encode_oral_truncation( + pairs: List[Tuple[List[int], List[int]]], + loss_mask: List[bool], + tokenizer: Any, + template: Optional[TemplateMeta], + config: EncodeConfig, +) -> Optional[EncodedSample]: + """Oral truncation: process turns in reverse order, newest first.""" + max_seq_len = config.max_seq_len + sep_token_len = _get_sep_token_len(template, tokenizer) + + num_pairs = len(pairs) + tokens_chunks: List[List[int]] = [] + labels_chunks: List[List[int]] = [] + cur_len = 0 + + for turn_index in range(num_pairs - 1, -1, -1): + prompt_ids, response_ids = pairs[turn_index] + + if len(response_ids) == 0: + logger.warning("[SKIP] The length of encoded assistant tokens is 0") + return None + + remaining_len = max_seq_len - cur_len + if len(prompt_ids) + len(response_ids) > remaining_len: + if len(prompt_ids) > remaining_len: + break + else: + response_ids = response_ids[: remaining_len - len(prompt_ids)] + + labels_src = [-100] * len(prompt_ids) + + if turn_index < len(loss_mask) and not loss_mask[turn_index]: + labels_target = [-100] * len(response_ids) + else: + if sep_token_len > 0 and turn_index != (num_pairs - 1) and len(response_ids) > sep_token_len: + labels_target = list(response_ids[: len(response_ids) - sep_token_len]) + [-100] * sep_token_len + else: + labels_target = list(response_ids) + + tokens_chunks.append(list(prompt_ids) + list(response_ids)) + labels_chunks.append(labels_src + labels_target) + cur_len += len(prompt_ids) + len(response_ids) + + tokens_chunks.reverse() + labels_chunks.reverse() + + token_ids: List[int] = [] + labels: List[int] = [] + for tc in tokens_chunks: + token_ids.extend(tc) + for lc in labels_chunks: + labels.extend(lc) + + if not token_ids: + return None + + _apply_dynamic_eos(token_ids, labels, template, tokenizer) + token_ids, labels = _apply_auto_bos(token_ids, labels, config, tokenizer) + _apply_efficient_eos(token_ids, labels, template, tokenizer) + + if len(token_ids) > max_seq_len: + token_ids = token_ids[:max_seq_len] + labels = labels[:max_seq_len] + + labels = _apply_label_shift(labels, config) + + validated = _validate_and_build(token_ids, labels) + if validated is None: + return None + token_ids, labels, position_ids = validated + return EncodedSample(input_ids=token_ids, labels=labels, seq_len=len(token_ids), position_ids=position_ids) + + +def encode_vl_sft( + example: Dict[str, Any], + tokenizer: Any, + template: Optional[TemplateMeta], + config: EncodeConfig, + processor: Any, + mm_plugin: Any, +) -> Optional[VLEncodedSample]: + """Encode a VL-SFT sample: process images + messages → token_ids + labels + mm_inputs.""" + messages = example.get("messages") + if not messages or len(messages) < 2: + return None + + images = example.get("images", []) + if not images: + return None + + images = [img["path"] if isinstance(img, dict) else img for img in images] + + mm_inputs = mm_plugin.get_mm_inputs(images=images, videos=[], audios=[], processor=processor) + messages = mm_plugin.process_messages( + messages=messages, images=images, videos=[], audios=[], mm_inputs=mm_inputs, processor=processor + ) + + if template is not None: + messages = _inject_ernie_think_system(messages, template) + + loss_mask = _extract_loss_mask(messages) + tools = example.get("tools") + pairs = _dispatch_encode(template, tokenizer, messages, tools=tools) + if not pairs: + return None + + sep_token_len = _get_sep_token_len(template, tokenizer) + token_ids, labels = _flatten_turns(pairs, loss_mask, sep_token_len) + if not token_ids: + return None + + _apply_dynamic_eos(token_ids, labels, template, tokenizer) + _apply_efficient_eos(token_ids, labels, template, tokenizer) + + # VL-specific: mm_plugin overrides labels for image/video placeholder regions + labels = mm_plugin.process_tokens(token_ids, processor) + + labels = _apply_label_shift(labels, config) + + # VL: oral not supported + if len(token_ids) > config.max_seq_len: + if config.truncation == "delete" or config.truncation == "oral": + logger.warning("[SKIP] VL data too long, discarding") + return None + elif config.placeholder_tokens: + token_ids, labels = _truncate_placeholder_aware( + token_ids, labels, config.max_seq_len, set(config.placeholder_tokens), config.truncation + ) + elif config.truncation == "left": + token_ids = token_ids[-config.max_seq_len :] + labels = labels[-config.max_seq_len :] + else: + token_ids = token_ids[: config.max_seq_len] + labels = labels[: config.max_seq_len] + + token_ids, labels = _apply_auto_bos(token_ids, labels, config, tokenizer) + + position_ids = list(range(len(token_ids))) + seq_len = len(token_ids) + return VLEncodedSample( + input_ids=token_ids, labels=labels, seq_len=seq_len, mm_inputs=mm_inputs, position_ids=position_ids + ) + + +def encode_pt( + example: Dict[str, Any], + tokenizer: Any, + config: EncodeConfig, +) -> Optional[EncodedSample]: + """Encode a single pretrain sample: plain text → token_ids with full loss.""" + messages = example.get("messages") + if not messages: + return None + + content = messages[0].get("content", "") + if not content: + return None + + tokens = tokenizer.convert_tokens_to_ids(tokenizer.tokenize(content)) + if tokenizer.eos_token_id is not None: + tokens = tokens + [tokenizer.eos_token_id] + + if len(tokens) < 2: + return None + + if len(tokens) > config.max_seq_len + 1: + tokens = tokens[: config.max_seq_len + 1] + + input_ids = tokens[:-1] + labels = tokens[1:] + position_ids = list(range(len(input_ids))) + + return EncodedSample(input_ids=input_ids, labels=labels, seq_len=len(input_ids), position_ids=position_ids) + + +# --------------------------------------------------------------------------- +# DPO Encoding +# --------------------------------------------------------------------------- + + +@dataclass +class DPOEncodeConfig(EncodeConfig): + """Configuration for DPO encoding.""" + + use_filtered_label_loss: bool = True + + +@dataclass +class DPOEncodedSample: + """Output of encode_dpo. + + The DPO sequence layout: + [prompt_tokens, chosen_response[:-1], prompt_last_token, rejected_response[:-1]] + + Position IDs fork at the prompt boundary: + [0..P-1, P..P+C-2, P-1, P..P+R-2] + + Response labels: + [-100]*(P-1) + chosen_with_eos + rejected_with_eos + """ + + input_ids: List[int] + position_ids: List[int] + response_labels: List[int] + response_index: List[int] # [chosen_start, rejected_start, rejected_end] + seq_len: int + score_delta: float = 1.0 + + +def _get_suffix_ids(template: Optional[TemplateMeta], tokenizer: Any) -> List[int]: + """Get EOS/suffix token IDs for DPO response endings.""" + if template is not None and template.suffix: + return tokenizer.encode(template.suffix[-1], add_special_tokens=False) + if tokenizer.eos_token_id is not None: + return [tokenizer.eos_token_id] + return [] + + +def _find_divergence_index(messages: List[Dict], rejected_messages: List[Dict]) -> int: + """Find the index where chosen and rejected message lists diverge. + + Returns the index of the first differing message. Typically this is + len(messages)-1 for standard DPO (only the last assistant turn differs). + """ + min_len = min(len(messages), len(rejected_messages)) + for i in range(min_len): + if messages[i].get("role") != rejected_messages[i].get("role"): + return i + if messages[i].get("content") != rejected_messages[i].get("content"): + return i + return min_len + + +def encode_dpo( + example: Dict[str, Any], + tokenizer: Any, + template: Optional[TemplateMeta], + config: DPOEncodeConfig, +) -> Optional[DPOEncodedSample]: + """Encode a single DPO sample: chosen/rejected messages → forked sequence. + + Expected input format: + - messages: Full chosen conversation [system?, user, assistant, ...] + - rejected_messages: Full rejected conversation [system?, user, assistant, ...] + + The two lists share a common prefix (prompt). This function: + 1. Finds the divergence point + 2. Encodes prompt, chosen response, and rejected response + 3. Concatenates into the forked DPO format + + Args: + example: Dict with 'messages' and 'rejected_messages'. + tokenizer: Tokenizer instance. + template: TemplateMeta for encoding. + config: DPOEncodeConfig. + + Returns: + DPOEncodedSample or None if encoding fails. + """ + messages = example.get("messages") + rejected_messages = example.get("rejected_messages") + + if not messages or not rejected_messages: + logger.warning("[SKIP] DPO sample missing messages or rejected_messages") + return None + + if len(messages) < 2 or len(rejected_messages) < 2: + logger.warning("[SKIP] DPO messages too short") + return None + + # Inject system prompt for ERNIE models if needed + if template is not None: + messages = _inject_ernie_think_system(messages, template) + rejected_messages = _inject_ernie_think_system(rejected_messages, template) + + # Find shared prefix (prompt) vs diverging response + diverge_idx = _find_divergence_index(messages, rejected_messages) + if diverge_idx == 0: + logger.warning("[SKIP] DPO messages diverge at index 0 (no shared prompt)") + return None + + # Encode the full chosen and rejected conversations + tools = example.get("tools") + chosen_pairs = _dispatch_encode(template, tokenizer, messages, tools=tools) + rejected_pairs = _dispatch_encode(template, tokenizer, rejected_messages, tools=tools) + + if not chosen_pairs or not rejected_pairs: + logger.warning("[SKIP] DPO encoding produced empty pairs") + return None + + # Determine the split point in encoded pairs. + # Each pair corresponds to one (user, assistant) turn. + # diverge_idx is in message-space. Convert to pair-space: + # Messages: [sys?, user, asst, user, asst, ...] + # Pairs: [pair0(user+asst), pair1(user+asst), ...] + # If system message exists, it's folded into pair0's prompt. + + # Count how many complete (user, assistant) turns are in the shared prefix + prompt_messages = messages[:diverge_idx] + # Count assistant messages in prompt (= number of complete turns in prompt) + prompt_turns = sum(1 for m in prompt_messages if m.get("role") == "assistant") + + # Split pairs into prompt_pairs and response_pairs + split_pair_idx = prompt_turns + if split_pair_idx >= len(chosen_pairs) or split_pair_idx >= len(rejected_pairs): + # Edge case: all turns are shared or encoding mismatch + logger.warning("[SKIP] DPO split point beyond encoded pairs") + return None + + # Build prompt token sequence from shared pairs + prompt_token_ids: List[int] = [] + for i in range(split_pair_idx): + q, a = chosen_pairs[i] + prompt_token_ids += q + a + + # Add the prompt part (query) of the split pair + chosen_split_q, chosen_split_a = chosen_pairs[split_pair_idx] + rejected_split_q, rejected_split_a = rejected_pairs[split_pair_idx] + prompt_token_ids += chosen_split_q + + # Build chosen response tokens (remaining turns after split) + chosen_response_ids: List[int] = list(chosen_split_a) + for i in range(split_pair_idx + 1, len(chosen_pairs)): + q, a = chosen_pairs[i] + chosen_response_ids += q + a + + # Build rejected response tokens (remaining turns after split) + rejected_response_ids: List[int] = list(rejected_split_a) + for i in range(split_pair_idx + 1, len(rejected_pairs)): + q, a = rejected_pairs[i] + rejected_response_ids += q + a + + # Add EOS/suffix to both responses + suffix_ids = _get_suffix_ids(template, tokenizer) + efficient_eos = template.efficient_eos if template else False + if efficient_eos and suffix_ids: + chosen_response_ids += suffix_ids + rejected_response_ids += suffix_ids + + # Check minimum lengths + if not chosen_response_ids or not rejected_response_ids: + logger.warning("[SKIP] DPO chosen or rejected response is empty after encoding") + return None + + prompt_len = len(prompt_token_ids) + chosen_len = len(chosen_response_ids) + rejected_len = len(rejected_response_ids) + + # Total length: prompt + (chosen-1) + 1(fork token) + (rejected-1) = prompt + chosen + rejected - 1 + total_len = prompt_len + chosen_len + rejected_len - 1 + + # Truncate prompt from the front if too long + max_seq_len = config.max_seq_len + if total_len > max_seq_len: + excess = total_len - max_seq_len + if excess >= prompt_len: + logger.warning("[SKIP] DPO sequence too long even without prompt") + return None + prompt_token_ids = prompt_token_ids[excess:] + prompt_len = len(prompt_token_ids) + total_len = max_seq_len + + if prompt_len == 0: + logger.warning("[SKIP] DPO prompt became empty after truncation") + return None + + # Construct the forked DPO sequence: + # input_ids = prompt + chosen[:-1] + [prompt_last] + rejected[:-1] + input_ids = prompt_token_ids + chosen_response_ids[:-1] + [prompt_token_ids[-1]] + rejected_response_ids[:-1] + + # Position IDs: fork at prompt end + # prompt: [0, 1, ..., P-1] + # chosen: [P, P+1, ..., P+C-2] + # fork: [P-1] + # rejected: [P, P+1, ..., P+R-2] + position_ids = ( + list(range(prompt_len)) + + list(range(prompt_len, prompt_len + chosen_len - 1)) + + [prompt_len - 1] + + list(range(prompt_len, prompt_len + rejected_len - 1)) + ) + + # Response labels: + # [-100]*(P-1) + chosen_response_with_eos + rejected_response_with_eos + response_labels = [-100] * (prompt_len - 1) + chosen_response_ids + rejected_response_ids + + # Response index: [chosen_start, rejected_start, rejected_end] + if config.use_filtered_label_loss: + response_index = [0, chosen_len, chosen_len + rejected_len] + else: + response_index = [ + prompt_len - 1, + prompt_len - 1 + chosen_len, + prompt_len - 1 + chosen_len + rejected_len, + ] + + # Sanity checks + assert len(input_ids) == total_len, f"input_ids len {len(input_ids)} != total_len {total_len}" + assert len(position_ids) == total_len, f"position_ids len {len(position_ids)} != total_len {total_len}" + assert len(response_labels) == total_len, f"response_labels len {len(response_labels)} != total_len {total_len}" + + return DPOEncodedSample( + input_ids=input_ids, + position_ids=position_ids, + response_labels=response_labels, + response_index=response_index, + seq_len=total_len, + score_delta=example.get("score_delta", 1.0), + ) diff --git a/paddleformers/datasets_v2/datapipe/packing.py b/paddleformers/datasets_v2/datapipe/packing.py new file mode 100644 index 00000000000..c733ccd04ee --- /dev/null +++ b/paddleformers/datasets_v2/datapipe/packing.py @@ -0,0 +1,127 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Packing algorithms: bin multiple short sequences into max_seq_len slots. + +Two strategies: + - greedy_pack: largest-remaining-capacity first-fit (fast, no external deps) + - binpack_ffd: First-Fit Decreasing via `binpacking` library (better packing efficiency) +""" + +from typing import List + +import numpy as np + +from .encode import EncodedSample + + +def greedy_pack( + samples: List[EncodedSample], + max_seq_len: int, +) -> List[List[EncodedSample]]: + """Pack samples into bins using greedy first-fit (largest gap first). + + For each sample, finds the open bin with the most remaining capacity. + If the sample fits, it's placed there. Otherwise, a new bin is opened. + + Args: + samples: List of encoded samples to pack. + max_seq_len: Maximum total token count per bin. + + Returns: + List of packed groups. Each group is a list of samples whose + combined seq_len <= max_seq_len. + """ + if not samples: + return [] + + n = len(samples) + remaining = np.full(n, -1, dtype=np.int64) + remaining[0] = max_seq_len + packs: List[List[EncodedSample]] = [[]] + next_bin = 1 + + for sample in samples: + seq_len = sample.seq_len + if seq_len > max_seq_len: + # Sample exceeds max_seq_len — put it alone in its own bin (will be truncated by collate) + packs.append([sample]) + continue + + best_bin = int(np.argmax(remaining)) + if seq_len <= remaining[best_bin]: + packs[best_bin].append(sample) + remaining[best_bin] -= seq_len + else: + # Open a new bin + if next_bin >= n: + # Extend array if needed + remaining = np.append(remaining, np.full(n, -1, dtype=np.int64)) + remaining[next_bin] = max_seq_len - seq_len + packs.append([sample]) + next_bin += 1 + + return [p for p in packs if p] + + +def binpack_ffd( + samples: List[EncodedSample], + max_seq_len: int, + packing_interval: int = 1000, +) -> List[List[EncodedSample]]: + """Bin-pack using First-Fit Decreasing (binpacking library). + + Processes samples in chunks of packing_interval for memory efficiency. + The last incomplete bin in each chunk is carried over to the next chunk, + allowing cross-chunk optimization. + + Args: + samples: List of encoded samples to pack. + max_seq_len: Maximum total token count per bin. + packing_interval: Buffer size — how many samples to accumulate + before running the FFD algorithm. Larger values give better + packing efficiency but use more memory. + + Returns: + List of packed groups. Each group is a list of samples whose + combined seq_len <= max_seq_len. + """ + import binpacking + + if not samples: + return [] + + # Build (sample, weight) tuples for the binpacking library + items = [(s, s.seq_len) for s in samples] + accumulated = [] + result: List[List[EncodedSample]] = [] + + for start in range(0, len(items), packing_interval): + end = min(start + packing_interval, len(items)) + is_finished = end == len(items) + accumulated += items[start:end] + + bins = binpacking.to_constant_volume(accumulated, max_seq_len, weight_pos=1) + + if bins and not is_finished: + # Emit all complete bins, keep last one for next iteration + for b in bins[:-1]: + result.append([item[0] for item in b]) + accumulated = bins[-1] + else: + for b in bins: + result.append([item[0] for item in b]) + accumulated = [] + + return [g for g in result if g] diff --git a/paddleformers/datasets_v2/datapipe/template.py b/paddleformers/datasets_v2/datapipe/template.py new file mode 100644 index 00000000000..52acfa4e637 --- /dev/null +++ b/paddleformers/datasets_v2/datapipe/template.py @@ -0,0 +1,1063 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Independent template system for datasets_v2. + +Provides a simple, data-driven template definition and encoding mechanism. +Templates are defined as lists of "slots" (strings with {{content}} placeholders, +dicts for special token lookup, or sets for built-in token IDs). + +Registration is one function call with plain strings — no Formatter classes needed. + +Supports: +- Multi-turn SFT encoding (user/assistant pairs) +- Reasoning/Thinking mode (auto-add tags) +- Tool calling (function role + observation role) +- Jinja-based encoding (using tokenizer's built-in chat_template) +- fix_special_tokens (auto fix eos/pad) +""" + +import json +import logging +import re +from copy import deepcopy +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional, Set, Tuple, Union + +from .tool_utils import FunctionCall, get_tool_utils + +logger = logging.getLogger(__name__) + +# A slot is one of: +# str: text fragment, may contain {{content}} placeholder +# dict: {"token": ""} → looked up by name in tokenizer vocab +# set: {"bos_token"} or {"eos_token"} → resolved to tokenizer's special token ID +Slot = Union[str, Dict[str, str], Set[str]] + + +@dataclass +class TemplateMeta: + """Pure-data template definition.""" + + name: str + user: List[Slot] + assistant: List[Slot] + system: List[Slot] = field(default_factory=lambda: ["{{content}}"]) + prefix: List[Slot] = field(default_factory=list) + # Observation role slot (for tool response messages) + observation: Optional[List[Slot]] = None + # Function role slot (for tool call messages from assistant) + function: Optional[List[Slot]] = None + chat_sep: str = "" + default_system: str = "" + suffix: List[str] = field(default_factory=list) + stop_tokens: List[str] = field(default_factory=list) + efficient_eos: bool = True + # Reasoning/Thinking support + thought_words: Tuple[str, str] = ("\n", "\n\n\n") + enable_thinking: Optional[bool] = None # None = not a reasoning template + # Tool format name (references tool_utils.py) + tool_format: Optional[str] = None + + +# ============================================================ +# Registry +# ============================================================ + +_TEMPLATE_REGISTRY: Dict[str, TemplateMeta] = {} + + +def register_template( + name: str, + *, + user: List[Slot], + assistant: List[Slot], + system: Optional[List[Slot]] = None, + prefix: Optional[List[Slot]] = None, + observation: Optional[List[Slot]] = None, + function: Optional[List[Slot]] = None, + chat_sep: str = "", + default_system: str = "", + suffix: Optional[List[str]] = None, + stop_tokens: Optional[List[str]] = None, + efficient_eos: bool = True, + thought_words: Optional[Tuple[str, str]] = None, + enable_thinking: Optional[bool] = None, + tool_format: Optional[str] = None, + exist_ok: bool = False, +) -> None: + """Register a template by name. + + Example: + register_template( + "chatml", + user=["<|im_start|>user\\n{{content}}<|im_end|>\\n<|im_start|>assistant\\n"], + assistant=["{{content}}"], + system=["<|im_start|>system\\n{{content}}<|im_end|>\\n"], + chat_sep="<|im_end|>\\n", + suffix=["<|im_end|>"], + ) + """ + if not exist_ok and name in _TEMPLATE_REGISTRY: + raise ValueError(f"Template '{name}' is already registered. Use exist_ok=True to overwrite.") + _TEMPLATE_REGISTRY[name] = TemplateMeta( + name=name, + user=user, + assistant=assistant, + system=system or ["{{content}}"], + prefix=prefix or [], + observation=observation, + function=function, + chat_sep=chat_sep, + default_system=default_system, + suffix=suffix or [], + stop_tokens=stop_tokens or [], + efficient_eos=efficient_eos, + thought_words=thought_words or ("\n", "\n\n\n"), + enable_thinking=enable_thinking, + tool_format=tool_format, + ) + + +def get_template(name: str) -> TemplateMeta: + """Look up a registered template by name. + + Raises: + KeyError: if name not found. + """ + if name not in _TEMPLATE_REGISTRY: + available = ", ".join(sorted(_TEMPLATE_REGISTRY.keys())) + raise KeyError(f"Template '{name}' not found. Available: {available}") + return _TEMPLATE_REGISTRY[name] + + +def list_templates() -> List[str]: + """Return sorted list of all registered template names.""" + return sorted(_TEMPLATE_REGISTRY.keys()) + + +# ============================================================ +# Slot substitution +# ============================================================ + + +def _substitute_slots(slots: List[Slot], **kwargs: str) -> List[Slot]: + """Replace {{key}} placeholders in string slots with values.""" + result: List[Slot] = [] + for slot in slots: + if isinstance(slot, str): + s = slot + for key, value in kwargs.items(): + s = s.replace("{{" + key + "}}", value, 1) + result.append(s) + else: + result.append(slot) + return result + + +# ============================================================ +# Slot → token IDs conversion +# ============================================================ + + +def _slots_to_ids(tokenizer: Any, slots: List[Slot]) -> List[int]: + """Convert a list of slots to token IDs using the given tokenizer.""" + token_ids: List[int] = [] + for slot in slots: + if isinstance(slot, str): + if len(slot) > 0: + token_ids += tokenizer.encode(slot, add_special_tokens=False) + elif isinstance(slot, dict): + token_name = slot.get("token", "") + if token_name: + tid = tokenizer.convert_tokens_to_ids(token_name) + if tid is not None: + token_ids.append(tid) + elif isinstance(slot, set): + if "bos_token" in slot and tokenizer.bos_token_id is not None: + token_ids.append(tokenizer.bos_token_id) + elif "eos_token" in slot and tokenizer.eos_token_id is not None: + token_ids.append(tokenizer.eos_token_id) + return token_ids + + +# ============================================================ +# Tool call formatting helpers +# ============================================================ + + +def _format_function_content(content, tool_format: str, thought_words: Optional[Tuple[str, str]] = None) -> str: + """Format function call content using tool_utils. + + Parses JSON tool calls and formats them according to the tool_format. + content can be a JSON string or an already-parsed list/dict. + """ + thought = None + # Extract thought only from string content + if isinstance(content, str) and thought_words and len(thought_words) == 2 and len(content) > 0: + regex = re.compile(rf"{re.escape(thought_words[0])}(.*?){re.escape(thought_words[1])}", re.DOTALL) + thought = re.search(regex, content) + if thought: + content = content.replace(thought.group(0), "") + + tool_utils = get_tool_utils(tool_format) + functions: list[FunctionCall] = [] + try: + # Handle both string (JSON) and already-parsed list/dict + if isinstance(content, str): + tool_calls = json.loads(content) + else: + tool_calls = content + + if not isinstance(tool_calls, list): + tool_calls = [tool_calls] + + for tool_call in tool_calls: + if "type" in tool_call and tool_call["type"] == "function": + tool_call = tool_call["function"] + arguments = tool_call["arguments"] + if not isinstance(arguments, str): + arguments = json.dumps(arguments, ensure_ascii=False) + functions.append(FunctionCall(tool_call["name"], arguments)) + + except (json.JSONDecodeError, TypeError): + raise RuntimeError(f"Invalid format in function message: {str([content])}.") + + function_str = tool_utils.function_formatter(functions) + if thought: + function_str = thought.group(0) + function_str + + return function_str + + +def _format_tools_content(content, tool_format: str) -> str: + """Format tools description using tool_utils. + + content can be a JSON string or an already-parsed list. + """ + tool_utils = get_tool_utils(tool_format) + try: + if isinstance(content, str): + tools = json.loads(content) + else: + tools = content + return tool_utils.tool_formatter(tools) if len(tools) != 0 else "" + except (json.JSONDecodeError, TypeError): + raise RuntimeError(f"Invalid format in tool description: {str([content])}.") + + +# ============================================================ +# Reasoning/Thinking helpers +# ============================================================ + + +def _remove_thought(content: str, thought_words: Tuple[str, str]) -> str: + """Remove thought tags from content.""" + pattern = re.compile(f"{re.escape(thought_words[0])}(.*?){re.escape(thought_words[1])}", re.DOTALL) + return re.sub(pattern, "", content).lstrip("\n") + + +_GLM5_TEMPLATES = {"glm_moe_dsa"} + + +def _get_thought_word_ids(tokenizer: Any, thought_words: Tuple[str, str], template_name: str = "") -> List[int]: + """Get token IDs for empty thought. GLM5 uses only closing tag.""" + if template_name in _GLM5_TEMPLATES: + return tokenizer.encode(thought_words[1], add_special_tokens=False) + return tokenizer.encode(f"{thought_words[0]}{thought_words[1]}", add_special_tokens=False) + + +# ============================================================ +# Core encoding +# ============================================================ + + +def encode_multiturn( + template: TemplateMeta, + tokenizer: Any, + messages: List[Dict[str, str]], + system: Optional[str] = None, + tools: Optional[Any] = None, +) -> List[Tuple[List[int], List[int]]]: + """Encode a multi-turn conversation into (prompt_ids, response_ids) pairs. + + Args: + template: Template definition. + tokenizer: Tokenizer with encode() and convert_tokens_to_ids(). + messages: List of {"role": "user"|"assistant"|"system"|"function"|"observation", "content": str}. + system: Override system message. If None, uses template.default_system. + tools: JSON string describing available tools. Appended to system message. + + Returns: + List of (prompt_ids, response_ids) tuples, one per turn. + prompt_ids includes user message tokens (and prefix/system for first turn). + response_ids includes assistant response tokens. + """ + # Extract system message if present as first message + actual_messages = messages + if messages and messages[0].get("role") == "system": + system = system or messages[0]["content"] + actual_messages = messages[1:] + + system = system or template.default_system + + encoded_messages: List[List[int]] = [] + for i, message in enumerate(actual_messages): + elements: List[Slot] = [] + + if i == 0: + # Prefix (BOS, etc.) + if template.prefix: + elements += template.prefix + # System message and tools description + if system or tools: + tool_text = "" + if tools and template.tool_format: + tool_text = _format_tools_content(tools, template.tool_format) + sys_content = (system or "") + tool_text + elements += _substitute_slots(template.system, content=sys_content) + + role = message.get("role", "") + content = message.get("content", "") + + if role == "user": + elements += _substitute_slots(template.user, content=content) + elif role == "assistant": + # Check if this message has tool_calls + if "tool_calls" in message and template.tool_format: + func_content = _format_function_content( + message["tool_calls"], template.tool_format, template.thought_words + ) + # Preserve thinking from content field if present + if content and template.thought_words: + think_start, think_end = template.thought_words + thought_regex = re.compile(rf"{re.escape(think_start)}.*?{re.escape(think_end)}", re.DOTALL) + thought_match = thought_regex.search(content if isinstance(content, str) else "") + if thought_match: + func_content = thought_match.group(0) + func_content + func_slots = template.function or template.assistant + elements += _substitute_slots(func_slots, content=func_content) + else: + elements += _substitute_slots(template.assistant, content=content) + if i < len(actual_messages) - 1 and template.chat_sep: + elements.append(template.chat_sep) + elif role in ("function", "tool_call"): + # Function call from assistant + if template.tool_format: + func_content = _format_function_content(content, template.tool_format, template.thought_words) + func_slots = template.function or template.assistant + elements += _substitute_slots(func_slots, content=func_content) + else: + elements += _substitute_slots(template.assistant, content=content) + if i < len(actual_messages) - 1 and template.chat_sep: + elements.append(template.chat_sep) + elif role in ("observation", "tool_response", "tool"): + # Tool response + obs_slots = template.observation or template.user + obs_content = content if isinstance(content, str) else json.dumps(content, ensure_ascii=False) + elements += _substitute_slots(obs_slots, content=obs_content) + else: + # Fallback: treat unknown roles as user + elements += _substitute_slots(template.user, content=content) + + encoded_messages.append(_slots_to_ids(tokenizer, elements)) + + # Pair up: (prompt, response) for each turn + pairs: List[Tuple[List[int], List[int]]] = [] + for i in range(0, len(encoded_messages) - 1, 2): + prompt_ids = encoded_messages[i] + response_ids = encoded_messages[i + 1] if i + 1 < len(encoded_messages) else [] + pairs.append((prompt_ids, response_ids)) + + return pairs + + +def encode_multiturn_reasoning( + template: TemplateMeta, + tokenizer: Any, + messages: List[Dict[str, str]], + system: Optional[str] = None, + tools: Optional[Any] = None, +) -> List[Tuple[List[int], List[int]]]: + """Encode with reasoning/thinking support. + + Aligned with old ReasoningTemplate.encode_multiturn: + - If enable_thinking is False: removes CoT from ALL assistant messages + - For each turn without thought tags: inserts empty thought tokens + - If enable_thinking is truthy: thought IDs go into response (trained on) + - If enable_thinking is falsy: thought IDs go into prompt (not trained on) + """ + messages = deepcopy(messages) + thought_words = template.thought_words + + # Extract system + actual_messages = messages + if messages and messages[0].get("role") == "system": + system = system or messages[0]["content"] + actual_messages = messages[1:] + + # If enable_thinking is False, remove all CoT from all assistant messages + if template.enable_thinking is False: + for i in range(1, len(actual_messages), 2): + if actual_messages[i].get("role") == "assistant": + actual_messages[i]["content"] = _remove_thought(actual_messages[i]["content"], thought_words) + + # Rebuild messages with system + if system: + rebuild = [{"role": "system", "content": system}] + actual_messages + else: + rebuild = actual_messages + + # Encode using normal path + pairs = encode_multiturn(template, tokenizer, rebuild, system=system, tools=tools) + + # Add empty thought to ALL turns that don't have thought tags + for i in range(0, len(actual_messages), 2): + if i + 1 >= len(actual_messages): + break + assistant_content = actual_messages[i + 1].get("content", "") + has_thought = thought_words[0].strip() in assistant_content and thought_words[1].strip() in assistant_content + pair_idx = i // 2 + if pair_idx >= len(pairs): + break + + if not has_thought: + thought_ids = _get_thought_word_ids(tokenizer, thought_words, template.name) + prompt_ids, response_ids = pairs[pair_idx] + if not template.enable_thinking: + prompt_ids = prompt_ids + thought_ids + else: + response_ids = thought_ids + response_ids + pairs[pair_idx] = (prompt_ids, response_ids) + + return pairs + + +# ============================================================ +# fix_special_tokens +# ============================================================ + + +def fix_special_tokens(tokenizer: Any, template: TemplateMeta) -> None: + """Fix eos and pad tokens in the tokenizer based on template config. + + - If tokenizer has no eos_token, adds '<|endoftext|>' + - If tokenizer has no pad_token, uses eos_token + - Adds stop_tokens as additional special tokens + """ + if tokenizer.eos_token_id is None: + num_added = tokenizer.add_special_tokens({"eos_token": "<|endoftext|>"}) + logger.warning(f"Add eos token: {tokenizer.eos_token}.") + if num_added > 0: + logger.warning("New tokens have been added, make sure `resize_vocab` is True.") + + if tokenizer.pad_token_id is None: + tokenizer.pad_token = tokenizer.eos_token + logger.info(f"Add pad token: {tokenizer.pad_token}") + + if template.stop_tokens: + num_added = tokenizer.add_special_tokens( + dict(additional_special_tokens=template.stop_tokens), replace_extra_special_tokens=False + ) + logger.info("Add {} to stop words.".format(",".join(template.stop_tokens))) + if num_added > 0: + logger.warning("New tokens have been added, make sure `resize_vocab` is True.") + + +# ============================================================ +# parse_template: auto-detect from tokenizer's chat_template +# ============================================================ + + +def parse_template(tokenizer: Any) -> TemplateMeta: + """Extract a template from the tokenizer's built-in chat_template. + + Useful when no explicit template name is specified — auto-detect from Jinja. + """ + + def find_diff(short_str: str, long_str: str) -> str: + i, j = 0, 0 + diff = "" + while i < len(short_str) and j < len(long_str): + if short_str[i] == long_str[j]: + i += 1 + j += 1 + else: + diff += long_str[j] + j += 1 + return diff + + prefix = tokenizer.decode(tokenizer.encode("")) + + messages = [{"role": "system", "content": "{{content}}"}] + system_slot = tokenizer.apply_chat_template(messages, add_generation_prompt=False, tokenize=False)[len(prefix) :] + + messages = [{"role": "system", "content": ""}, {"role": "user", "content": "{{content}}"}] + user_slot_empty_system = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False) + user_slot_empty_system = user_slot_empty_system[len(prefix) :] + + messages = [{"role": "user", "content": "{{content}}"}] + user_slot = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False) + user_slot = user_slot[len(prefix) :] + + messages = [{"role": "user", "content": "{{content}}"}, {"role": "assistant", "content": "{{content}}"}] + assistant_slot = tokenizer.apply_chat_template(messages, add_generation_prompt=False, tokenize=False) + assistant_slot = assistant_slot[len(prefix) + len(user_slot) :] + + # Detect reasoning template + is_reasoning = "" in assistant_slot + assistant_slot = assistant_slot.replace("", "").replace("", "").lstrip("\n") + + if len(user_slot) > len(user_slot_empty_system): + default_system = find_diff(user_slot_empty_system, user_slot) + sole_system = system_slot.replace("{{content}}", default_system, 1) + user_slot = user_slot[len(sole_system) :] + else: + default_system = "" + + eos_token = getattr(tokenizer, "eos_token", None) or "" + efficient_eos = eos_token not in assistant_slot + + return TemplateMeta( + name="_auto_parsed", + user=[user_slot], + assistant=[assistant_slot], + system=[system_slot], + prefix=[prefix] if prefix else [], + default_system=default_system, + suffix=[], + stop_tokens=[], + thought_words=("\n", "\n\n\n"), + enable_thinking=True if is_reasoning else None, + efficient_eos=efficient_eos, + ) + + +# ============================================================ +# get_template_and_fix_tokenizer: high-level entry point +# ============================================================ + + +def get_template_and_fix_tokenizer( + tokenizer: Any, + template_name: Optional[str] = None, + tool_format: Optional[str] = None, + default_system: Optional[str] = None, +) -> TemplateMeta: + """Get template by name (or auto-parse) and fix tokenizer special tokens. + + Args: + tokenizer: The tokenizer instance. + template_name: Registered template name, or None to auto-parse. + tool_format: Override the tool_format in the template. + default_system: Override the default system message. + + Returns: + The TemplateMeta instance (possibly modified). + """ + if template_name is None: + if isinstance(getattr(tokenizer, "chat_template", None), str): + logger.warning("`template` was not specified, try parsing the chat template from the tokenizer.") + template = parse_template(tokenizer) + else: + logger.warning("`template` was not specified, use `empty` template.") + template = get_template("empty") + else: + template = get_template(template_name) + + # Override tool_format + if tool_format is not None: + template = TemplateMeta( + name=template.name, + user=template.user, + assistant=template.assistant, + system=template.system, + prefix=template.prefix, + observation=template.observation, + function=template.function, + chat_sep=template.chat_sep, + default_system=template.default_system, + suffix=template.suffix, + stop_tokens=template.stop_tokens, + efficient_eos=template.efficient_eos, + thought_words=template.thought_words, + enable_thinking=template.enable_thinking, + tool_format=tool_format, + ) + + # Override default_system + if default_system is not None: + template.default_system = default_system + + # Ensure suffix is set + if not template.suffix: + template.suffix = [tokenizer.eos_token] + logger.warning("suffix is not specified, using eos token as suffix.") + + # Fix special tokens + fix_special_tokens(tokenizer, template) + + return template + + +# ============================================================ +# Jinja-based encoding +# ============================================================ + + +def encode_multiturn_jinja( + tokenizer: Any, + messages: List[Dict[str, str]], +) -> List[Tuple[List[int], List[int]]]: + """Encode using the tokenizer's built-in chat_template (Jinja2). + + For models that ship with a chat_template in their tokenizer config. + Encodes each turn incrementally to separate prompt and response IDs. + + Args: + tokenizer: Tokenizer with apply_chat_template() method. + messages: List of {"role": str, "content": str}. + + Returns: + List of (prompt_ids, response_ids) tuples. + """ + pairs: List[Tuple[List[int], List[int]]] = [] + + # Build up conversation turn by turn + accumulated: List[Dict[str, str]] = [] + prev_len = 0 + i = 0 + while i < len(messages): + # Skip system at start — it gets included with first user turn + if i == 0 and messages[i].get("role") == "system": + accumulated.append(messages[i]) + i += 1 + continue + + # Expect user then assistant + if i < len(messages) and messages[i].get("role") == "user": + accumulated.append(messages[i]) + # Encode up to this user message with generation prompt + prompt_ids = tokenizer.apply_chat_template( + accumulated, add_generation_prompt=True, tokenize=True, return_dict=False + ) + + if i + 1 < len(messages) and messages[i + 1].get("role") == "assistant": + accumulated.append(messages[i + 1]) + # Encode with assistant response included + full_ids = tokenizer.apply_chat_template( + accumulated, add_generation_prompt=False, tokenize=True, return_dict=False + ) + response_ids = full_ids[len(prompt_ids) :] + pairs.append((prompt_ids if not pairs else prompt_ids[prev_len:], response_ids)) + prev_len = len(full_ids) + i += 2 + else: + # No response — just prompt + pairs.append((prompt_ids if not pairs else prompt_ids[prev_len:], [])) + prev_len = len(prompt_ids) + i += 1 + else: + i += 1 + + # Simpler approach: encode full conversation, then use incremental to split + if not pairs: + pairs = _encode_jinja_simple(tokenizer, messages) + + return pairs + + +def _encode_jinja_simple( + tokenizer: Any, + messages: List[Dict[str, str]], +) -> List[Tuple[List[int], List[int]]]: + """Simple jinja encoding: encode incrementally turn by turn.""" + pairs: List[Tuple[List[int], List[int]]] = [] + prev_len = 0 + accumulated: List[Dict[str, str]] = [] + + i = 0 + # Handle leading system message + if messages and messages[0].get("role") == "system": + accumulated.append(messages[0]) + i = 1 + + while i < len(messages): + msg = messages[i] + if msg.get("role") == "user": + accumulated.append(msg) + prompt_ids_full = tokenizer.apply_chat_template( + accumulated, add_generation_prompt=True, tokenize=True, return_dict=False + ) + prompt_ids = prompt_ids_full[prev_len:] + + if i + 1 < len(messages) and messages[i + 1].get("role") == "assistant": + accumulated.append(messages[i + 1]) + full_ids = tokenizer.apply_chat_template( + accumulated, add_generation_prompt=False, tokenize=True, return_dict=False + ) + response_ids = full_ids[len(prompt_ids_full) :] + pairs.append((prompt_ids, response_ids)) + prev_len = len(full_ids) + i += 2 + else: + pairs.append((prompt_ids, [])) + prev_len = len(prompt_ids_full) + i += 1 + else: + accumulated.append(msg) + i += 1 + + return pairs + + +# ============================================================ +# Built-in template definitions +# ============================================================ + +# ChatML format (Qwen, Qwen2, Qwen3, etc.) +register_template( + "chatml", + user=["<|im_start|>user\n{{content}}<|im_end|>\n<|im_start|>assistant\n"], + assistant=["{{content}}"], + system=["<|im_start|>system\n{{content}}<|im_end|>\n"], + chat_sep="<|im_end|>\n", + suffix=["<|im_end|>"], + stop_tokens=["<|im_end|>", "<|endoftext|>"], + default_system="You are Qwen, created by Alibaba Cloud. You are a helpful assistant.", +) + +# Qwen3 (reasoning model, same format as chatml but with thinking) +register_template( + "qwen3", + user=["<|im_start|>user\n{{content}}<|im_end|>\n<|im_start|>assistant\n"], + assistant=["{{content}}"], + system=["<|im_start|>system\n{{content}}<|im_end|>\n"], + observation=[ + "<|im_start|>user\n\n{{content}}\n<|im_end|>\n<|im_start|>assistant\n" + ], + function=["{{content}}"], + chat_sep="<|im_end|>\n", + suffix=["<|im_end|>\n"], + enable_thinking=True, + tool_format="qwen", +) + +# Qwen3 without thinking +register_template( + "qwen3_nothink", + user=["<|im_start|>user\n{{content}}<|im_end|>\n<|im_start|>assistant\n"], + assistant=["{{content}}"], + system=["<|im_start|>system\n{{content}}<|im_end|>\n"], + observation=[ + "<|im_start|>user\n\n{{content}}\n<|im_end|>\n<|im_start|>assistant\n" + ], + function=["{{content}}"], + chat_sep="<|im_end|>\n", + suffix=["<|im_end|>\n"], + tool_format="qwen", +) + +# Qwen3.5 variant (assistant includes im_end in slot) +register_template( + "qwen3_5", + user=["<|im_start|>user\n{{content}}<|im_end|>\n<|im_start|>assistant\n"], + assistant=["{{content}}<|im_end|>\n"], + system=["<|im_start|>system\n{{content}}<|im_end|>\n"], + observation=[ + "<|im_start|>user\n\n{{content}}\n<|im_end|>\n<|im_start|>assistant\n" + ], + function=["{{content}}<|im_end|>\n"], + stop_tokens=["<|im_end|>"], + enable_thinking=True, + tool_format="qwen3_5", + efficient_eos=False, +) + +# Qwen3.5 without thinking +register_template( + "qwen3_5_nothink", + user=["<|im_start|>user\n{{content}}<|im_end|>\n<|im_start|>assistant\n"], + assistant=["{{content}}<|im_end|>\n"], + system=["<|im_start|>system\n{{content}}<|im_end|>\n"], + observation=[ + "<|im_start|>user\n\n{{content}}\n<|im_end|>\n<|im_start|>assistant\n" + ], + function=["{{content}}<|im_end|>\n"], + stop_tokens=["<|im_end|>"], + tool_format="qwen3_5", + efficient_eos=False, +) + +# Qwen2-VL / Qwen3-VL +register_template( + "qwen2_vl", + user=["<|im_start|>user\n{{content}}<|im_end|>\n<|im_start|>assistant\n"], + assistant=["{{content}}"], + system=["<|im_start|>system\n{{content}}<|im_end|>\n"], + observation=[ + "<|im_start|>user\n\n{{content}}\n<|im_end|>\n<|im_start|>assistant\n" + ], + function=["{{content}}"], + default_system="You are a helpful assistant.", + chat_sep="<|im_end|>\n", + suffix=["<|im_end|>\n"], + tool_format="qwen", +) + +register_template( + "qwen3_vl", + user=["<|im_start|>user\n{{content}}<|im_end|>\n<|im_start|>assistant\n"], + assistant=["{{content}}"], + system=["<|im_start|>system\n{{content}}<|im_end|>\n"], + observation=[ + "<|im_start|>user\n\n{{content}}\n<|im_end|>\n<|im_start|>assistant\n" + ], + function=["{{content}}"], + chat_sep="<|im_end|>\n", + suffix=["<|im_end|>\n"], + enable_thinking=True, + tool_format="qwen", +) + +register_template( + "qwen3_vl_nothink", + user=["<|im_start|>user\n{{content}}<|im_end|>\n<|im_start|>assistant\n"], + assistant=["{{content}}"], + system=["<|im_start|>system\n{{content}}<|im_end|>\n"], + observation=[ + "<|im_start|>user\n\n{{content}}\n<|im_end|>\n<|im_start|>assistant\n" + ], + function=["{{content}}"], + chat_sep="<|im_end|>\n", + suffix=["<|im_end|>\n"], + tool_format="qwen", +) + +# Chatml EOS variant +register_template( + "chatml_eos", + user=["<|im_start|>user\n{{content}}<|im_end|>\n<|im_start|>assistant\n"], + assistant=["{{content}}<|im_end|>\n"], + system=["<|im_start|>system\n{{content}}<|im_end|>\n"], + stop_tokens=["<|im_end|>"], + efficient_eos=False, +) + +# Llama3 format +register_template( + "llama3", + prefix=[{"token": "<|begin_of_text|>"}], + user=[ + "<|start_header_id|>user<|end_header_id|>\n\n{{content}}<|eot_id|>" + "<|start_header_id|>assistant<|end_header_id|>\n\n" + ], + assistant=["{{content}}<|eot_id|>"], + system=["<|start_header_id|>system<|end_header_id|>\n\n{{content}}<|eot_id|>"], + observation=[ + "<|start_header_id|>ipython<|end_header_id|>\n\n{{content}}<|eot_id|>" + "<|start_header_id|>assistant<|end_header_id|>\n\n" + ], + function=["{{content}}"], + stop_tokens=["<|eot_id|>"], + chat_sep="<|eot_id|>", + efficient_eos=False, + tool_format="llama3", +) + +# DeepSeek v3 format (uses fullwidth vertical bars) +register_template( + "deepseek3", + user=["<｜User｜>{{content}}\n\n<｜Assistant｜>"], + assistant=["{{content}}"], + system=["{{content}}\n\n"], + prefix=[{"token": "<｜begin▁of▁sentence｜>"}], + stop_tokens=["<｜end▁of▁sentence｜>"], + chat_sep="<｜end▁of▁sentence｜>", +) + +# GLM4 format +register_template( + "glm4", + user=["<|user|>\n{{content}}<|assistant|>\n"], + assistant=["\n{{content}}"], + system=["<|system|>\n{{content}}"], + observation=["<|observation|>\n{{content}}<|assistant|>"], + function=["{{content}}"], + stop_tokens=["<|endoftext|>", "<|user|>"], + tool_format="glm4", +) + +# GLM4 MOE +register_template( + "glm4_moe", + user=["<|user|>\n{{content}}<|assistant|>\n"], + assistant=["\n{{content}}"], + system=["<|system|>\n{{content}}"], + observation=["<|observation|>\n{{content}}<|assistant|>"], + function=["{{content}}"], + prefix=["[gMASK][gMASK]"], + suffix=["<|user|>"], + thought_words=("", ""), + enable_thinking=True, + tool_format="glm4_moe", +) + +# GLM5 / GLM-MOE-DSA +register_template( + "glm_moe_dsa", + user=["<|user|>{{content}}<|assistant|>"], + assistant=["{{content}}"], + system=["[gMASK]<|system|>{{content}}"], + observation=["<|observation|>{{content}}<|assistant|>"], + function=["{{content}}"], + prefix=["[gMASK]"], + suffix=["<|user|>"], + thought_words=("", ""), + enable_thinking=True, + tool_format="glm_moe_dsa", +) + +# GLM4V (vision) +register_template( + "glm4v", + user=["<|user|>\n{{content}}<|assistant|>"], + assistant=["\n{{content}}"], + system=["<|system|>\n{{content}}"], + observation=["<|observation|>\n{{content}}<|assistant|>"], + function=["{{content}}"], + prefix=["[gMASK]"], + suffix=["<|user|>"], + enable_thinking=True, + tool_format="glm4", +) + +# GLM4V MOE (vision) +register_template( + "glm4v_moe", + user=["<|user|>\n{{content}}<|assistant|>\n"], + assistant=["\n{{content}}"], + system=["<|system|>\n{{content}}"], + observation=["<|observation|>\n{{content}}<|assistant|>"], + function=["{{content}}"], + prefix=["[gMASK]"], + suffix=["<|user|>"], + stop_tokens=["<|user|>", "<|observation|>", ""], + thought_words=("", ""), + enable_thinking=True, + tool_format="glm4_moe", +) + +# GLM-OCR +register_template( + "glm_ocr", + user=["<|user|>\n{{content}}\n"], + assistant=["{{content}}"], + prefix=["[gMASK]"], + chat_sep="<|assistant|>\n", +) + +# ERNIE (Baidu) format +register_template( + "ernie", + user=["<|im_start|>user\n{{content}}<|im_end|>\n\n<|im_start|>assistant\n"], + assistant=["{{content}}"], + system=["<|im_start|>system\n{{content}}<|im_end|>\n"], + observation=[ + "<|im_start|>user\n\n{{content}}\n<|im_end|>\n\n<|im_start|>assistant\n" + ], + function=["\n{{content}}"], + chat_sep="<|im_end|>\n\n", + suffix=["<|im_end|>"], + stop_tokens=["<|im_end|>"], + thought_words=("", ""), + enable_thinking=True, + tool_format="ernie", +) + +# ERNIE VL +register_template( + "ernie_vl", + user=["User: {{content}}\nAssistant: "], + assistant=["{{content}}"], + system=["{{content}}\n"], + observation=["User: \n{{content}}\n\n\nAssistant: "], + function=["\n{{content}}"], + prefix=["<|begin_of_sentence|>"], + chat_sep="<|end_of_sentence|>", + suffix=[""], + thought_words=("\n\n", "\n\n\n"), + enable_thinking=True, + tool_format="ernie_vl", +) + +# PaddleOCR VL +register_template( + "paddleocr_vl", + user=["User: {{content}}\nAssistant: "], + assistant=["{{content}}"], + system=["{{content}}\n"], + prefix=["<|begin_of_sentence|>"], + chat_sep="<|end_of_sentence|>", +) + +# Gemma format +register_template( + "gemma", + user=["user\n{{content}}\nmodel\n"], + assistant=["{{content}}"], + system=["{{content}}\n\n"], + observation=["tool\n{{content}}\nmodel\n"], + prefix=[{"bos_token"}], + chat_sep="\n", + suffix=[""], + stop_tokens=[""], +) + +# Phi-4 format +register_template( + "phi4", + user=["<|im_start|>user<|im_sep|>{{content}}<|im_end|><|im_start|>assistant<|im_sep|>"], + assistant=["{{content}}"], + system=["<|im_start|>system<|im_sep|>{{content}}<|im_end|>"], + suffix=["<|im_end|>"], + chat_sep="<|im_end|>", +) + +# Kimi K2 +register_template( + "kimi_k2", + user=["<｜User｜>{{content}}\n\n<｜Assistant｜>"], + assistant=["{{content}}"], + system=["{{content}}\n\n"], + prefix=[{"bos_token"}], + chat_sep="<｜end▁of▁sentence｜>", + stop_tokens=["<｜end▁of▁sentence｜>"], +) + +# Simple default format +register_template( + "default", + user=["Human: {{content}}", {"eos_token"}, "\nAssistant:"], + assistant=["{{content}}", {"eos_token"}, "\n"], + system=["System: {{content}}", {"eos_token"}, "\n"], +) + +# Empty template — passthrough +register_template( + "empty", + user=["{{content}}"], + assistant=["{{content}}"], + efficient_eos=False, +) diff --git a/paddleformers/datasets_v2/datapipe/tool_utils.py b/paddleformers/datasets_v2/datapipe/tool_utils.py new file mode 100644 index 00000000000..8eedae45a97 --- /dev/null +++ b/paddleformers/datasets_v2/datapipe/tool_utils.py @@ -0,0 +1,357 @@ +# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# The file has been adapted from hiyouga LLaMA-Factory project +# Copyright (c) 2025 LLaMA-Factory +# Licensed under the Apache License - https://github.com/hiyouga/LLaMA-Factory/blob/main/LICENSE + + +import json +from abc import ABC, abstractmethod +from dataclasses import dataclass +from datetime import datetime +from typing import Any, NamedTuple + +from typing_extensions import override + + +class FunctionCall(NamedTuple): + name: str + arguments: str + + +DEFAULT_TOOL_PROMPT = ( + "You have access to the following tools:\n{tool_text}" + "Use the following format if using a tool:\n" + "```\n" + "Action: tool name (one of [{tool_names}])\n" + "Action Input: the input to the tool, in a JSON format representing the kwargs " + """(e.g. ```{{"input": "hello world", "num_beams": 5}}```)\n""" + "```\n" +) + +QWEN_TOOL_PROMPT = ( + "\n\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\n" + "You are provided with function signatures within XML tags:\n{tool_text}" + "\n\n\nFor each function call, return a json object with function name and arguments within " + """ XML tags:\n\n{{"name": , """ + """"arguments": }}\n""" +) + +ERNIE_TOOL_PROMPT = "\n\n\n[{tool_text}]\n" + +ERNIE_VL_TOOL_PROMPT = "\n\n[{tool_text}]\n\n" + + +GLM4_TOOL_PROMPT = ( + "你是一个名为 ChatGLM 的人工智能助手。你是基于智谱 AI 公司训练的语言模型 GLM-4 模型开发的，" "你的任务是针对用户的问题和要求提供适当的答复和支持。\n\n# 可用工具{tool_text}" +) + +GLM4_MOE_TOOL_PROMPT = ( + "\n\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\n" + "You are provided with function signatures within XML tags:\n{tool_text}" + "\n\n\nFor each function call, output the function name and arguments within the following XML format:" + "\n{{function-name}}" + "\n{{arg-key-1}}" + "\n{{arg-value-1}}" + "\n{{arg-key-2}}" + "\n{{arg-value-2}}" + "\n...\n\n" +) + +# GLM_5 template (no \n separators in tool_call format) +GLM_5_TOOL_PROMPT = ( + "\n\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\n" + "You are provided with function signatures within XML tags:\n{tool_text}" + "\n\n\nFor each function call, output the function name and arguments within the following XML format:" + "\n{{function-name}}" + "{{arg-key-1}}" + "{{arg-value-1}}" + "{{arg-key-2}}" + "{{arg-value-2}}" + "..." +) + +LLAMA3_TOOL_PROMPT = ( + "Cutting Knowledge Date: December 2023\nToday Date: {date}\n\n" + "You have access to the following functions. To call a function, please respond with JSON for a function call. " + """Respond in the format {{"name": function name, "parameters": dictionary of argument name and its value}}. """ + "Do not use variables.\n\n{tool_text}" +) + + +@dataclass +class ToolUtils(ABC): + """Base class for tool utilities.""" + + @staticmethod + @abstractmethod + def tool_formatter(tools: list[dict[str, Any]]) -> str: + r"""Generate the system message describing all the available tools.""" + ... + + @staticmethod + @abstractmethod + def function_formatter(functions: list["FunctionCall"]) -> str: + r"""Generate the assistant message including all the tool calls.""" + ... + + +class DefaultToolUtils(ToolUtils): + r"""Default tool using template.""" + + @override + @staticmethod + def tool_formatter(tools: list[dict[str, Any]]) -> str: + tool_text = "" + tool_names = [] + for tool in tools: + tool = tool.get("function", "") if tool.get("type") == "function" else tool + param_text = "" + for name, param in tool["parameters"]["properties"].items(): + required, enum, items = "", "", "" + if name in tool["parameters"].get("required", []): + required = ", required" + + if param.get("enum", None): + enum = ", should be one of [{}]".format(", ".join(param["enum"])) + + if param.get("items", None): + items = ", where each item should be {}".format(param["items"].get("type", "")) + + param_text += " - {name} ({type}{required}): {desc}{enum}{items}\n".format( + name=name, + type=param.get("type", ""), + required=required, + desc=param.get("description", ""), + enum=enum, + items=items, + ) + + tool_text += "> Tool Name: {name}\nTool Description: {desc}\nTool Args:\n{args}\n".format( + name=tool["name"], desc=tool.get("description", ""), args=param_text + ) + tool_names.append(tool["name"]) + + return DEFAULT_TOOL_PROMPT.format(tool_text=tool_text, tool_names=", ".join(tool_names)) + + @override + @staticmethod + def function_formatter(functions: list["FunctionCall"]) -> str: + return "\n".join([f"Action: {name}\nAction Input: {arguments}" for name, arguments in functions]) + + +class QwenToolUtils(ToolUtils): + r"""Qwen 2.5 tool using template.""" + + @override + @staticmethod + def tool_formatter(tools: list[dict[str, Any]]) -> str: + tool_text = "" + for tool in tools: + wrapped_tool = tool if tool.get("type") == "function" else {"type": "function", "function": tool} + tool_text += "\n" + json.dumps(wrapped_tool, ensure_ascii=False) + + return QWEN_TOOL_PROMPT.format(tool_text=tool_text) + + @override + @staticmethod + def function_formatter(functions: list["FunctionCall"]) -> str: + function_texts = [ + json.dumps({"name": name, "arguments": json.loads(arguments)}, ensure_ascii=False) + for name, arguments in functions + ] + return "\n".join([f"\n{text}\n" for text in function_texts]) + + +class GLM4ToolUtils(ToolUtils): + r"""GLM-4 tool using template.""" + + @override + @staticmethod + def tool_formatter(tools: list[dict[str, Any]]) -> str: + tool_text = "" + for tool in tools: + tool = tool.get("function", "") if tool.get("type") == "function" else tool + tool_text += "\n\n## {name}\n\n{body}\n在调用上述函数时，请使用 Json 格式表示调用的参数。".format( + name=tool["name"], body=json.dumps(tool, indent=4, ensure_ascii=False) + ) + + return GLM4_TOOL_PROMPT.format(tool_text=tool_text) + + @override + @staticmethod + def function_formatter(functions: list["FunctionCall"]) -> str: + if len(functions) > 1: + raise ValueError("GLM-4 does not support parallel functions.") + + return f"{functions[0].name}\n{functions[0].arguments}" + + +class GLM4MOEToolUtils(QwenToolUtils): + r"""GLM-4-MOE tool using template.""" + + @override + @staticmethod + def tool_formatter(tools: list[dict[str, Any]]) -> str: + tool_text = "" + for tool in tools: + wrapped_tool = tool if tool.get("type") == "function" else {"type": "function", "function": tool} + tool_text += "\n" + json.dumps(wrapped_tool, ensure_ascii=False) + + return GLM4_MOE_TOOL_PROMPT.format(tool_text=tool_text) + + @override + @staticmethod + def function_formatter(functions: list["FunctionCall"]) -> str: + function_json = [ + {"func_name": name, "func_key_values": json.loads(arguments)} for name, arguments in functions + ] + function_texts = [] + for func in function_json: + prompt = "\n" + func["func_name"] + for key, value in func["func_key_values"].items(): + prompt += "\n" + key + "" + if not isinstance(value, str): + value = json.dumps(value, ensure_ascii=False) + prompt += "\n" + value + "" + function_texts.append(prompt) + + return "\n".join(function_texts) + + +class GLM_5ToolUtils(QwenToolUtils): + r"""GLM-4.7/GLM-5 tool using template (GLM_5 style without \\n separators).""" + + @override + @staticmethod + def tool_formatter(tools: list[dict[str, Any]]) -> str: + tool_text = "" + for tool in tools: + wrapped_tool = tool if tool.get("type") == "function" else {"type": "function", "function": tool} + tool_text += "\n" + json.dumps(wrapped_tool, ensure_ascii=False) + + return GLM_5_TOOL_PROMPT.format(tool_text=tool_text) + + @override + @staticmethod + def function_formatter(functions: list["FunctionCall"]) -> str: + function_json = [ + {"func_name": name, "func_key_values": json.loads(arguments)} for name, arguments in functions + ] + function_texts = [] + for func in function_json: + prompt = "" + func["func_name"] + for key, value in func["func_key_values"].items(): + prompt += "" + key + "" + if not isinstance(value, str): + value = json.dumps(value, ensure_ascii=False) + prompt += "" + value + "" + prompt += "" + function_texts.append(prompt) + + return "".join(function_texts) + + +class Llama3ToolUtils(ToolUtils): + r"""Llama 3.x tool using template with `tools_in_user_message=False`. + + Reference: https://www.llama.com/docs/model-cards-and-prompt-formats/llama3_1/#json-based-tool-calling + """ + + @override + @staticmethod + def tool_formatter(tools: list[dict[str, Any]]) -> str: + date = datetime.now().strftime("%d %b %Y") + tool_text = "" + for tool in tools: + wrapped_tool = tool if tool.get("type") == "function" else {"type": "function", "function": tool} + tool_text += json.dumps(wrapped_tool, indent=4, ensure_ascii=False) + "\n\n" + + return LLAMA3_TOOL_PROMPT.format(date=date, tool_text=tool_text) + + @override + @staticmethod + def function_formatter(functions: list["FunctionCall"]) -> str: + function_objects = [{"name": name, "parameters": json.loads(arguments)} for name, arguments in functions] + return json.dumps(function_objects[0] if len(function_objects) == 1 else function_objects, ensure_ascii=False) + + +class ERNIEToolUtils(ToolUtils): + r"""ERNIE 4.5 tool using template.""" + + @override + @staticmethod + def tool_formatter(tools: list[dict[str, Any]]) -> str: + tool_text_list = [] + for tool in tools: + wrapped_tool = tool if tool.get("type") == "function" else {"type": "function", "function": tool} + tool_text_list.append(json.dumps(wrapped_tool, ensure_ascii=False)) + tool_text = ", ".join(tool_text_list) + + return ERNIE_TOOL_PROMPT.format(tool_text=tool_text) + + @override + @staticmethod + def function_formatter(functions: list["FunctionCall"]) -> str: + function_texts = [ + json.dumps({"name": name, "arguments": json.loads(arguments)}, ensure_ascii=False) + for name, arguments in functions + ] + return "\n".join([f"\n{text}\n\n" for text in function_texts]) + + +class ERNIEVLToolUtils(ToolUtils): + r"""ERNIE VL 4.5 tool using template.""" + + @override + @staticmethod + def tool_formatter(tools: list[dict[str, Any]]) -> str: + tool_text_list = [] + for tool in tools: + wrapped_tool = tool if tool.get("type") == "function" else {"type": "function", "function": tool} + tool_text_list.append(json.dumps(wrapped_tool, ensure_ascii=False)) + tool_text = ", ".join(tool_text_list) + + return ERNIE_VL_TOOL_PROMPT.format(tool_text=tool_text) + + @override + @staticmethod + def function_formatter(functions: list["FunctionCall"]) -> str: + function_texts = [ + json.dumps({"name": name, "arguments": json.loads(arguments)}, ensure_ascii=False) + for name, arguments in functions + ] + return "\n".join([f"\n{text}\n" for text in function_texts]) + + +TOOLS = { + "default": DefaultToolUtils(), + "ernie": ERNIEToolUtils(), + "ernie_vl": ERNIEVLToolUtils(), + "qwen": QwenToolUtils(), + "qwen3_5": QwenToolUtils(), + "glm4": GLM4ToolUtils(), + "glm4_moe": GLM4MOEToolUtils(), + "glm_moe_dsa": GLM_5ToolUtils(), + "llama3": Llama3ToolUtils(), +} + + +def get_tool_utils(name: str) -> "ToolUtils": + tool_utils = TOOLS.get(name, None) + if tool_utils is None: + raise ValueError(f"Tool utils `{name}` not found.") + + return tool_utils diff --git a/paddleformers/datasets_v2/dataset/__init__.py b/paddleformers/datasets_v2/dataset/__init__.py new file mode 100644 index 00000000000..24531f4de9d --- /dev/null +++ b/paddleformers/datasets_v2/dataset/__init__.py @@ -0,0 +1,18 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""dataset: training-ready Dataset classes for datasets_v2.""" + +from .lazy_dataset import LazyEncodeDataset +from .streaming_dataset import StreamingDataset diff --git a/paddleformers/datasets_v2/dataset/lazy_dataset.py b/paddleformers/datasets_v2/dataset/lazy_dataset.py new file mode 100644 index 00000000000..d8feda879a5 --- /dev/null +++ b/paddleformers/datasets_v2/dataset/lazy_dataset.py @@ -0,0 +1,90 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Lazy-encoding dataset wrapper for training. + +Wraps an HF Dataset + encode function into a map-style Dataset suitable for +DataLoader. Encoding happens lazily in __getitem__, with automatic retry +on failure (bad samples get skipped via random fallback). +""" + +import logging +import traceback +from typing import Any, Callable, Dict, Optional + +import numpy as np +import paddle + +from ..datapipe.encode import EncodedSample + +logger = logging.getLogger(__name__) + + +class LazyEncodeDataset(paddle.io.Dataset): + """Map-style dataset that lazily encodes HF Dataset rows. + + Each __getitem__ call fetches a row from the underlying HF Dataset, + applies encode_func to produce an EncodedSample. If encoding fails, + retries with random fallback indices. + + Inherits from paddle.io.Dataset for compatibility with PaddleFormers Trainer. + + Args: + dataset: HuggingFace Dataset with 'messages' column. + encode_func: Callable that takes a row dict → EncodedSample or None. + n_try_fetch: Max retry count on encoding failure. + seed: Random seed for reproducible fallback selection. + """ + + def __init__( + self, + dataset: Any, + encode_func: Callable[[Dict[str, Any]], Optional[EncodedSample]], + n_try_fetch: int = 10, + seed: Optional[int] = None, + ) -> None: + super().__init__() + self.dataset = dataset + self.encode_func = encode_func + self.n_try_fetch = min(n_try_fetch, max(1, len(dataset))) + self._rng = np.random.RandomState(seed) + self._fallback_indices = self._rng.permutation(len(dataset)).tolist() + self._fallback_ptr = 0 + self._traceback_counter = 0 + self._traceback_limit = 10 + + def __len__(self) -> int: + return len(self.dataset) + + def __getitem__(self, idx: int) -> EncodedSample: + for attempt in range(self.n_try_fetch): + i = idx if attempt == 0 else self._next_fallback() + try: + row = self.dataset[i] + result = self.encode_func(row) + if result is not None: + return result + except Exception: + if self._traceback_counter < self._traceback_limit: + logger.warning( + f"Encoding failed for index {i} (attempt {attempt + 1}):\n" f"{traceback.format_exc()}" + ) + self._traceback_counter += 1 + + raise RuntimeError(f"Failed to encode sample after {self.n_try_fetch} attempts " f"(started at index {idx}).") + + def _next_fallback(self) -> int: + idx = self._fallback_indices[self._fallback_ptr % len(self._fallback_indices)] + self._fallback_ptr += 1 + return idx diff --git a/paddleformers/datasets_v2/dataset/streaming_dataset.py b/paddleformers/datasets_v2/dataset/streaming_dataset.py new file mode 100644 index 00000000000..917b70fc3cf --- /dev/null +++ b/paddleformers/datasets_v2/dataset/streaming_dataset.py @@ -0,0 +1,85 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Streaming dataset wrapper: adapts HF IterableDataset to paddle.io.IterableDataset.""" + +import numpy as np +import paddle + + +class StreamingDataset(paddle.io.IterableDataset): + """Wraps a HuggingFace IterableDataset as a paddle.io.IterableDataset. + + This allows Trainer to recognize the dataset as iterable and use the + appropriate DataLoader path (no sampler, direct iteration). + + Supports two modes: + - lazy=True (default): True streaming. Yields directly from HF iterator + without downloading/materializing the full dataset. Suitable for large + remote datasets (e.g. fineweb-edu). Shuffle uses HF's buffer shuffle. + - lazy=False: Legacy V1-compatible mode. Materializes all data into memory + for epoch-based full-array shuffle. Only suitable for small datasets. + + Args: + hf_iterable: A HuggingFace IterableDataset instance (e.g. from + load_dataset(..., streaming=True) with .map()/.filter() applied). + shuffle: Whether to shuffle data each epoch. + seed: Random seed for shuffle. + lazy: If True, iterate directly from HF without materialization (true streaming). + If False, materialize all data for epoch-based shuffle (V1 compat). + """ + + def __init__(self, hf_iterable, shuffle: bool = False, seed: int = 0, lazy: bool = True): + super().__init__() + self._hf = hf_iterable + self._shuffle = shuffle + self._seed = seed + self._lazy = lazy + self._data = None # only used when lazy=False + + def _materialize(self): + """Load all data from HF iterable into memory for epoch-based shuffle.""" + if self._data is None: + self._data = list(self._hf) + + def __iter__(self): + if self._lazy: + # True streaming: yield directly from HF iterator. + # HF IterableDataset fetches parquet chunks on-demand, never downloads all. + # For shuffle, rely on HF's .shuffle(buffer_size=...) applied upstream. + yield from self._hf + else: + # Legacy mode: materialize all data, then epoch-based shuffle. + self._materialize() + data = self._data + n = len(data) + indices = list(range(n)) + rng = np.random.RandomState(self._seed) if self._shuffle else None + + while True: + if rng is not None: + rng.shuffle(indices) + else: + indices = list(range(n)) + + last_item = None + for idx in indices: + last_item = data[idx] + yield last_item + + # Align with V1 IteratorSFTDataset._generate_sequences behavior: + # V1 yields the last element twice per epoch due to a trailing + # `if len(batch_sequence) > 0: yield batch_sequence` after the loop. + if last_item is not None: + yield last_item diff --git a/paddleformers/datasets_v2/dataset_info.json b/paddleformers/datasets_v2/dataset_info.json new file mode 100644 index 00000000000..84f041bda81 --- /dev/null +++ b/paddleformers/datasets_v2/dataset_info.json @@ -0,0 +1,137 @@ +[ + { + "name": "alpaca_en", + "path": "llamafactory/alpaca_en", + "preprocessor": "alpaca", + "tags": ["sft", "en"] + }, + { + "name": "alpaca_zh", + "path": "llamafactory/alpaca_zh", + "preprocessor": "alpaca", + "tags": ["sft", "zh"] + }, + { + "name": "alpaca_gpt4_en", + "path": "llamafactory/alpaca_gpt4_en", + "preprocessor": "alpaca", + "tags": ["sft", "en"] + }, + { + "name": "alpaca_gpt4_zh", + "path": "llamafactory/alpaca_gpt4_zh", + "preprocessor": "alpaca", + "tags": ["sft", "zh"] + }, + { + "name": "belle_2m", + "path": "BelleGroup/train_2M_CN", + "preprocessor": "alpaca", + "tags": ["sft", "zh"] + }, + { + "name": "belle_1m", + "path": "BelleGroup/train_1M_CN", + "preprocessor": "alpaca", + "tags": ["sft", "zh"] + }, + { + "name": "belle_0.5m", + "path": "BelleGroup/train_0.5M_CN", + "preprocessor": "alpaca", + "tags": ["sft", "zh"] + }, + { + "name": "belle_chat", + "path": "BelleGroup/generated_chat_0.4M", + "preprocessor": "alpaca", + "tags": ["sft", "zh"] + }, + { + "name": "belle_math", + "path": "BelleGroup/school_math_0.25M", + "preprocessor": "alpaca", + "tags": ["sft", "zh", "math"] + }, + { + "name": "code_alpaca", + "path": "sahil2801/CodeAlpaca-20k", + "preprocessor": "alpaca", + "tags": ["sft", "en", "code"] + }, + { + "name": "math_instruct", + "path": "TIGER-Lab/MathInstruct", + "preprocessor": "alpaca", + "tags": ["sft", "en", "math"] + }, + { + "name": "firefly", + "path": "YeungNLP/firefly-train-1.1M", + "preprocessor": "response", + "columns": {"query": "input", "response": "target"}, + "tags": ["sft", "zh"] + }, + { + "name": "webqa", + "path": "suolyer/webqa", + "preprocessor": "response", + "columns": {"query": "input", "response": "output"}, + "tags": ["sft", "zh"] + }, + { + "name": "openo1_sft", + "path": "llamafactory/OpenO1-SFT", + "preprocessor": "response", + "tags": ["sft", "en", "reasoning"] + }, + { + "name": "deepseek_r1_distill", + "path": "Congliu/Chinese-DeepSeek-R1-Distill-data-110k-SFT", + "preprocessor": "alpaca", + "tags": ["sft", "zh", "reasoning"] + }, + { + "name": "sharegpt_gpt4", + "path": "shibing624/sharegpt_gpt4", + "preprocessor": "messages", + "tags": ["sft", "zh"] + }, + { + "name": "lima", + "path": "llamafactory/lima", + "preprocessor": "messages", + "tags": ["sft", "en"] + }, + { + "name": "slim_orca", + "path": "Open-Orca/SlimOrca", + "preprocessor": "messages", + "tags": ["sft", "en"] + }, + { + "name": "sharegpt_hyperfiltered", + "path": "totally-not-an-llm/sharegpt-hyperfiltered-3k", + "preprocessor": "messages", + "tags": ["sft", "en"] + }, + { + "name": "neo_sft", + "path": "m-a-p/neo_sft_phase2", + "preprocessor": "messages", + "tags": ["sft", "en"] + }, + { + "name": "orca_dpo", + "path": "Intel/orca_dpo_pairs", + "preprocessor": "response", + "columns": {"query": "question"}, + "tags": ["dpo", "en"] + }, + { + "name": "german_rag_dpo", + "path": "avemio/German-RAG-DPO-ShareGPT-HESSIAN-AI", + "preprocessor": "messages", + "tags": ["dpo", "de"] + } +] diff --git a/paddleformers/datasets_v2/grounding_plugin.py b/paddleformers/datasets_v2/grounding_plugin.py new file mode 100644 index 00000000000..4dfdacd7456 --- /dev/null +++ b/paddleformers/datasets_v2/grounding_plugin.py @@ -0,0 +1,76 @@ +# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass +from typing import List + + +@dataclass +class BaseGroundingPlugin: + def normalize_bbox(self, bbox: List[float]) -> List[int]: + return [int(coord) for coord in bbox] + + def format_ref_object(self, obj_name: str) -> str: + return f"<|object_ref_start|>{obj_name}<|object_ref_end|>" + + def format_bbox(self, bbox: List[float]) -> str: + normalized = self.normalize_bbox(bbox) + return f"<|box_start|>({normalized[0]},{normalized[1]}),({normalized[2]},{normalized[3]})<|box_end|>" + + def process_messages(self, messages, objects): + + ref_objects = objects.get("ref", []) + bboxes = objects.get("bbox", []) + + ref_idx = 0 + bbox_idx = 0 + + for message in messages: + content = message.get("content", "") + ref_count = content.count("") + bbox_count = content.count("") + current_refs = ref_objects[ref_idx : ref_idx + ref_count] + current_bboxes = bboxes[bbox_idx : bbox_idx + bbox_count] + + for ref in current_refs: + message["content"] = message["content"].replace("", self.format_ref_object(ref), 1) + for bbox in current_bboxes: + message["content"] = message["content"].replace("", self.format_bbox(bbox), 1) + + ref_idx += ref_count + bbox_idx += bbox_count + + return messages + + +PLUGINS = { + "base": BaseGroundingPlugin, +} + + +def register_grounding_plugin(name, plugin_class): + if name in PLUGINS: + raise ValueError(f"Grounding plugin {name} already exists.") + + PLUGINS[name] = plugin_class + + +def get_grounding_plugin( + name: str, + **kwargs, +): + if name not in PLUGINS: + raise ValueError(f"Grounding plugin `{name}` not found.") + + return PLUGINS[name](**kwargs) diff --git a/paddleformers/datasets_v2/loaders.py b/paddleformers/datasets_v2/loaders.py new file mode 100644 index 00000000000..e73cbe3815a --- /dev/null +++ b/paddleformers/datasets_v2/loaders.py @@ -0,0 +1,428 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Dataset loading: local files, directories, and HuggingFace Hub. + +Provides a unified load_dataset() entry point that handles: +- Local files (json/jsonl/csv/tsv/parquet/arrow/txt) +- Local directories (auto-detect data files) +- HuggingFace Hub datasets (by repo_id) +- Streaming mode (IterableDataset) +- Automatic preprocessing via registry metadata +""" + +import logging +import os +from typing import Dict, List, Optional, Union + +from datasets import Dataset as HfMapDataset +from datasets import load_dataset as hf_load_dataset + +from .registry import DatasetMeta, get_dataset_meta, parse_dataset_string +from .schema import DATASET_TYPE + +logger = logging.getLogger(__name__) + +# ============================================================ +# File extension mapping +# ============================================================ + +_EXT_TO_FORMAT = { + "json": "json", + "jsonl": "json", + "csv": "csv", + "tsv": "csv", + "parquet": "parquet", + "arrow": "arrow", + "txt": "text", +} + + +# ============================================================ +# Internal helpers +# ============================================================ + + +def _detect_file_format(path: str) -> str: + """Detect the HF-compatible format name from file extension.""" + ext = os.path.splitext(path)[1].lstrip(".").lower() + fmt = _EXT_TO_FORMAT.get(ext) + if fmt is None: + raise ValueError(f"Unsupported file format '.{ext}' for: {path}. " f"Supported: {list(_EXT_TO_FORMAT.keys())}") + return fmt + + +def _resolve_source(name_or_path: str) -> str: + """Determine the source type of a dataset identifier. + + Returns one of: "file", "directory", "hub" + """ + if os.path.isfile(name_or_path): + return "file" + if os.path.isdir(name_or_path): + return "directory" + return "hub" + + +def _load_local_file( + path: str, + *, + split: str = "train", + streaming: bool = False, + **kwargs, +) -> DATASET_TYPE: + """Load a dataset from a local file.""" + if not os.path.isfile(path): + raise FileNotFoundError(f"Data file not found: {path}") + + fmt = _detect_file_format(path) + load_kwargs = {"split": split, "streaming": streaming, **kwargs} + + if fmt == "csv": + load_kwargs.setdefault("na_filter", False) + if path.endswith(".tsv"): + load_kwargs.setdefault("delimiter", "\t") + + dataset = hf_load_dataset(fmt, data_files=path, **load_kwargs) + logger.info(f"Loaded local file: {path} (format={fmt}, streaming={streaming})") + return dataset + + +def _load_local_directory( + directory: str, + *, + split: str = "train", + streaming: bool = False, + **kwargs, +) -> DATASET_TYPE: + """Load a dataset from a local directory. + + If the directory contains a loading script or dataset_infos.json, + delegates directly to HF. Otherwise, auto-detects data files. + """ + if not os.path.isdir(directory): + raise FileNotFoundError(f"Directory not found: {directory}") + + # Check if HF can handle the directory directly (has script or metadata) + has_script = any(f.endswith(".py") for f in os.listdir(directory)) + has_info = os.path.isfile(os.path.join(directory, "dataset_infos.json")) + has_info = has_info or os.path.isfile(os.path.join(directory, "dataset_info.json")) + + if has_script or has_info: + dataset = hf_load_dataset(directory, split=split, streaming=streaming, **kwargs) + logger.info(f"Loaded directory (HF script/metadata): {directory}") + return dataset + + # Auto-detect data files by extension + data_files = [] + detected_fmt = None + for fname in sorted(os.listdir(directory)): + ext = os.path.splitext(fname)[1].lstrip(".").lower() + if ext in _EXT_TO_FORMAT: + if detected_fmt is None: + detected_fmt = _EXT_TO_FORMAT[ext] + elif _EXT_TO_FORMAT[ext] == detected_fmt: + pass + else: + # Mixed formats: use the first one found + continue + data_files.append(os.path.join(directory, fname)) + + if not data_files: + raise ValueError( + f"No supported data files found in directory: {directory}. " + f"Supported extensions: {list(_EXT_TO_FORMAT.keys())}" + ) + + load_kwargs = {"split": split, "streaming": streaming, **kwargs} + if detected_fmt == "csv": + load_kwargs.setdefault("na_filter", False) + + dataset = hf_load_dataset(detected_fmt, data_files=data_files, **load_kwargs) + logger.info(f"Loaded directory: {directory} ({len(data_files)} files, format={detected_fmt})") + return dataset + + +def _ensure_proxy(): + """Ensure proxy env vars are propagated for hub downloads. + + Reads from PADDLEFORMERS_PROXY env var. If set, applies it as + http/https proxy (respecting any existing proxy settings). + """ + proxy = os.environ.get("PADDLEFORMERS_PROXY", "") + if not proxy: + return + for key in ("https_proxy", "HTTPS_PROXY", "http_proxy", "HTTP_PROXY"): + os.environ.setdefault(key, proxy) + no_proxy = os.environ.get("PADDLEFORMERS_NO_PROXY", "") + if no_proxy: + os.environ.setdefault("no_proxy", no_proxy) + os.environ.setdefault("NO_PROXY", no_proxy) + + +def _load_hub_dataset( + repo_id: str, + *, + subset: Optional[str] = None, + split: str = "train", + streaming: bool = False, + token: Optional[str] = None, + **kwargs, +) -> DATASET_TYPE: + """Load a dataset from HuggingFace Hub.""" + _ensure_proxy() + + load_kwargs = { + "split": split, + "streaming": streaming, + **kwargs, + } + if subset is not None: + load_kwargs["name"] = subset + if token is not None: + load_kwargs["token"] = token + + dataset = hf_load_dataset(repo_id, **load_kwargs) + logger.info(f"Loaded from Hub: {repo_id} (subset={subset}, split={split}, streaming={streaming})") + return dataset + + +def _get_preprocessor(meta: DatasetMeta): + """Resolve preprocessor from DatasetMeta. + + Returns a callable preprocessor instance, or None if no preprocessing. + """ + preprocessor = meta.preprocessor + if preprocessor is None: + return None + if callable(preprocessor): + return preprocessor + + # Lazy import to avoid circular dependencies + from .preprocessors import ( + AlpacaPreprocessor, + AutoPreprocessor, + ErnieKitPreprocessor, + MessagesPreprocessor, + ResponsePreprocessor, + TextPreprocessor, + ) + + _PREPROCESSOR_MAP = { + "auto": AutoPreprocessor, + "messages": MessagesPreprocessor, + "response": ResponsePreprocessor, + "alpaca": AlpacaPreprocessor, + "text": TextPreprocessor, + "erniekit": ErnieKitPreprocessor, + } + + if preprocessor not in _PREPROCESSOR_MAP: + raise ValueError(f"Unknown preprocessor '{preprocessor}'. " f"Available: {list(_PREPROCESSOR_MAP.keys())}") + + cls = _PREPROCESSOR_MAP[preprocessor] + columns = meta.columns or {} + return cls(columns=columns) + + +_FORMAT_TO_PREPROCESSOR = { + "erniekit": "erniekit", + "messages": "messages", + "alpaca": "alpaca", + "text": "text", + "response": "response", +} + + +def _resolve_preprocessor(meta, dataset_format, columns): + """Resolve preprocessor: dataset_format hint > registry meta > AutoPreprocessor.""" + from .preprocessors import ( + AlpacaPreprocessor, + AutoPreprocessor, + ErnieKitPreprocessor, + MessagesPreprocessor, + ResponsePreprocessor, + TextPreprocessor, + ) + + _FORMAT_CLASS_MAP = { + "erniekit": ErnieKitPreprocessor, + "messages": MessagesPreprocessor, + "alpaca": AlpacaPreprocessor, + "text": TextPreprocessor, + "response": ResponsePreprocessor, + } + + # Priority 1: explicit dataset_format hint + if dataset_format and dataset_format in _FORMAT_CLASS_MAP: + cls = _FORMAT_CLASS_MAP[dataset_format] + return cls(columns=columns) + + # Priority 2: registry metadata + if meta and meta.preprocessor is None: + return None + if meta: + return _get_preprocessor(meta) + + # Priority 3: auto-detect + return AutoPreprocessor(columns=columns) + + +# ============================================================ +# Main entry points +# ============================================================ + + +def load_dataset( + dataset: str, + *, + split: Optional[str] = None, + subset: Optional[str] = None, + streaming: bool = False, + token: Optional[str] = None, + num_proc: int = 1, + preprocess: bool = True, + columns: Optional[Dict[str, str]] = None, + strict: bool = False, + dataset_format: Optional[str] = None, +) -> DATASET_TYPE: + """Load a dataset from any source: local file, directory, or HuggingFace Hub. + + Supports: + - Registered dataset names: "alpaca" + - Local file paths: "/data/train.jsonl" + - Local directories: "/data/my_dataset/" + - HuggingFace Hub IDs: "tatsu-lab/alpaca" + - Sample syntax: "dataset_name#500" (sample 500 rows after loading) + + Args: + dataset: Dataset specification string. + split: Split to load. Overrides registry default. + subset: HF subset/config. Overrides registry default. + streaming: If True, returns an IterableDataset. + token: HuggingFace Hub token for private datasets. + num_proc: Number of workers for preprocessing (map-style only). + preprocess: Whether to apply the registered preprocessor. + columns: Column rename mapping. Merged with registry columns. + strict: If True, raise on preprocessing errors. + dataset_format: Explicit dataset format hint (e.g. "erniekit", "messages"). + When provided and recognized, directly selects the corresponding preprocessor + instead of using auto-detection. Falls back to AutoPreprocessor if unrecognized. + + Returns: + The loaded (and optionally preprocessed) dataset. + """ + # 1. Parse string syntax + spec = parse_dataset_string(dataset) + + # 2. Registry lookup + meta = get_dataset_meta(spec.name) + + # 3. Resolve effective parameters + effective_split = split or (meta.split if meta else "train") + effective_subset = subset or spec.subset or (meta.subset if meta else None) + effective_path = (meta.path if meta else None) or spec.name + + # Merge column mappings: registry columns + explicit columns (explicit wins) + effective_columns = {} + if meta and meta.columns: + effective_columns.update(meta.columns) + if columns: + effective_columns.update(columns) + + # 4. Determine source and load + source_type = _resolve_source(effective_path) + + if source_type == "file": + ds = _load_local_file(effective_path, split=effective_split, streaming=streaming) + elif source_type == "directory": + ds = _load_local_directory(effective_path, split=effective_split, streaming=streaming) + else: + ds = _load_hub_dataset( + effective_path, + subset=effective_subset, + split=effective_split, + streaming=streaming, + token=token, + ) + + # 5. Apply preprocessing + if preprocess: + preprocessor = _resolve_preprocessor(meta, dataset_format, effective_columns) + + if preprocessor is not None: + proc_num_proc = num_proc if not streaming else 1 + ds = preprocessor(ds, num_proc=proc_num_proc, strict=strict) + + # 6. Apply sampling + if spec.sample is not None: + if isinstance(ds, HfMapDataset): + from .ops import sample_dataset + + ds = sample_dataset(ds, spec.sample) + else: + ds = ds.take(spec.sample) + + return ds + + +def load_datasets( + datasets: Union[str, List[str]], + *, + split: Optional[str] = None, + streaming: bool = False, + token: Optional[str] = None, + num_proc: int = 1, + preprocess: bool = True, + columns: Optional[Dict[str, str]] = None, + strict: bool = False, +) -> DATASET_TYPE: + """Load multiple datasets and concatenate them. + + Args: + datasets: A single dataset string or list of dataset strings. + split: Default split for all datasets. + streaming: If True, all datasets are loaded in streaming mode. + token: HuggingFace Hub token. + num_proc: Number of workers for preprocessing. + preprocess: Whether to apply preprocessors. + columns: Column rename mapping applied to all datasets. + strict: Strict preprocessing mode. + + Returns: + A single concatenated dataset. + """ + if isinstance(datasets, str): + datasets = [datasets] + + loaded = [] + for ds_str in datasets: + ds = load_dataset( + ds_str, + split=split, + streaming=streaming, + token=token, + num_proc=num_proc, + preprocess=preprocess, + columns=columns, + strict=strict, + ) + loaded.append(ds) + + if len(loaded) == 1: + return loaded[0] + + from .ops import concat_datasets + + return concat_datasets(loaded) diff --git a/paddleformers/datasets_v2/mm_plugin.py b/paddleformers/datasets_v2/mm_plugin.py new file mode 100644 index 00000000000..1bac642c9a0 --- /dev/null +++ b/paddleformers/datasets_v2/mm_plugin.py @@ -0,0 +1,1529 @@ +# Copyright 2025 HuggingFace Inc. and the LlamaFactory team. +# +# This code is inspired by the HuggingFace's Transformers library. +# https://github.com/huggingface/transformers/blob/v4.40.0/src/transformers/models/llava/processing_llava.py +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# The file has been adapted from hiyouga LLaMA-Factory project +# Copyright (c) 2025 LLaMA-Factory +# Licensed under the Apache License - https://github.com/hiyouga/LLaMA-Factory/blob/main/LICENSE + +import copy +import inspect +import io +import math +import os +import random +import sys +from copy import deepcopy +from dataclasses import dataclass +from typing import Any, BinaryIO, Optional + +import librosa +import numpy as np +import paddle +import requests +from PIL import Image +from PIL.Image import Image as ImageObject +from transformers.image_utils import is_valid_image +from typing_extensions import override + +from paddleformers.transformers.qwen2_vl.vision_process import fetch_image, fetch_video +from paddleformers.transformers.qwen3_omni_moe.processor import ( + Qwen3OmniMoeProcessorKwargs, +) +from paddleformers.utils.log import logger + +from .augment_utils import ( + JpegCompression, + RandomApply, + RandomDiscreteRotation, + RandomScale, + RandomSingleSidePadding, + transforms, +) + +IMAGE_PLACEHOLDER = os.getenv("IMAGE_PLACEHOLDER", "") +VIDEO_PLACEHOLDER = os.getenv("VIDEO_PLACEHOLDER", "