From 68d9e01e071d9002df9e4a95f33c921dafe3b8a8 Mon Sep 17 00:00:00 2001 From: weiyixuanxx <864644273@qq.com> Date: Tue, 12 May 2026 19:24:04 +0800 Subject: [PATCH 01/17] feat: add datasets_v2 module and SFT-V2 training pipeline Introduce a new data loading and encoding pipeline (datasets_v2) with: - Schema-based dataset registry with preprocessor auto-detection - Independent template system (chatml, llama3, deepseek3, etc.) - Lazy encoding dataset with packing and flashmask support - SFT-V2 workflow (workflow2.py) integrated via stage routing Co-Authored-By: Claude Opus 4.7 --- docs/zh/datasets_v2_design.md | 245 ++++++++ paddleformers/cli/train/sft/__init__.py | 3 +- paddleformers/cli/train/sft/workflow2.py | 377 +++++++++++ paddleformers/cli/train/tuner.py | 5 +- paddleformers/datasets_v2/__init__.py | 70 +++ .../datasets_v2/datapipe/__init__.py | 28 + paddleformers/datasets_v2/datapipe/collate.py | 199 ++++++ paddleformers/datasets_v2/datapipe/encode.py | 140 +++++ paddleformers/datasets_v2/datapipe/packing.py | 74 +++ .../datasets_v2/datapipe/template.py | 426 +++++++++++++ paddleformers/datasets_v2/dataset_info.json | 137 ++++ paddleformers/datasets_v2/loaders.py | 371 +++++++++++ paddleformers/datasets_v2/ops.py | 155 +++++ .../datasets_v2/preprocessors/__init__.py | 19 + .../datasets_v2/preprocessors/auto.py | 61 ++ .../datasets_v2/preprocessors/base.py | 193 ++++++ .../datasets_v2/preprocessors/extra.py | 55 ++ .../datasets_v2/preprocessors/messages.py | 150 +++++ .../datasets_v2/preprocessors/response.py | 98 +++ paddleformers/datasets_v2/registry.py | 246 ++++++++ paddleformers/datasets_v2/schema.py | 192 ++++++ tests/datasets_v2/__init__.py | 13 + tests/datasets_v2/debug_pipeline.py | 142 +++++ tests/datasets_v2/test_datapipe.py | 585 ++++++++++++++++++ tests/datasets_v2/test_loaders.py | 365 +++++++++++ tests/datasets_v2/test_preprocessors.py | 316 ++++++++++ tests/datasets_v2/test_schema.py | 13 + 27 files changed, 4676 insertions(+), 2 deletions(-) create mode 100644 docs/zh/datasets_v2_design.md create mode 100644 paddleformers/cli/train/sft/workflow2.py create mode 100644 paddleformers/datasets_v2/__init__.py create mode 100644 paddleformers/datasets_v2/datapipe/__init__.py create mode 100644 paddleformers/datasets_v2/datapipe/collate.py create mode 100644 paddleformers/datasets_v2/datapipe/encode.py create mode 100644 paddleformers/datasets_v2/datapipe/packing.py create mode 100644 paddleformers/datasets_v2/datapipe/template.py create mode 100644 paddleformers/datasets_v2/dataset_info.json create mode 100644 paddleformers/datasets_v2/loaders.py create mode 100644 paddleformers/datasets_v2/ops.py create mode 100644 paddleformers/datasets_v2/preprocessors/__init__.py create mode 100644 paddleformers/datasets_v2/preprocessors/auto.py create mode 100644 paddleformers/datasets_v2/preprocessors/base.py create mode 100644 paddleformers/datasets_v2/preprocessors/extra.py create mode 100644 paddleformers/datasets_v2/preprocessors/messages.py create mode 100644 paddleformers/datasets_v2/preprocessors/response.py create mode 100644 paddleformers/datasets_v2/registry.py create mode 100644 paddleformers/datasets_v2/schema.py create mode 100644 tests/datasets_v2/__init__.py create mode 100644 tests/datasets_v2/debug_pipeline.py create mode 100644 tests/datasets_v2/test_datapipe.py create mode 100644 tests/datasets_v2/test_loaders.py create mode 100644 tests/datasets_v2/test_preprocessors.py create mode 100644 tests/datasets_v2/test_schema.py diff --git a/docs/zh/datasets_v2_design.md b/docs/zh/datasets_v2_design.md new file mode 100644 index 00000000000..ba180af4b34 --- /dev/null +++ b/docs/zh/datasets_v2_design.md @@ -0,0 +1,245 @@ +# datasets_v2 架构设计文档 + +## 1. 概述 + +datasets_v2 是 PaddleFormers 的新一代数据处理模块，参考 ms-swift 的 dataset 架构，基于 HuggingFace datasets 库构建，目标是提供统一的、可扩展的数据加载与预处理管线。 + +**核心设计原则：** +- 显式 schema 作为管线契约（所有模块基于同一份字段定义交互） +- 基于 HF Dataset.map() 的批处理框架（兼容 Map 式和 Iterable 流式） +- 多格式归一化：无论输入是什么格式，输出统一为 messages 标准格式 +- 面向多模态 + 纯文本双重场景 + +**对应 ms-swift 源码：** `/ms-swift/swift/dataset/` + +--- + +## 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 — 数据集操作 + +对应 ms-swift 的 `dataset/utils.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 类型、列映射）的注册机制 +- 支持内置数据集和用户自定义注册 +- 对应 ms-swift 的 `dataset/register.py` + +### 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. 与 ms-swift 的对应关系 + +| paddleformers datasets_v2 | ms-swift dataset | 说明 | +|---|---|---| +| schema.py | preprocessor/core.py (常量部分) | 独立成文件，更清晰 | +| preprocessors/base.py | preprocessor/core.py:RowPreprocessor | 去除采样逻辑，纯化批处理 | +| preprocessors/response.py | preprocessor/core.py:ResponsePreprocessor | 基本一致 | +| preprocessors/messages.py | preprocessor/core.py:MessagesPreprocessor | 基本一致 | +| preprocessors/extra.py | preprocessor/core.py:AlpacaPreprocessor | 基本一致 | +| preprocessors/auto.py | preprocessor/core.py:AutoPreprocessor | 基本一致 | +| ops.py | dataset/utils.py | 待实现 | +| registry.py | dataset/register.py | 待实现 | +| loaders.py | dataset/loader/ | 待实现 | +| builder.py | dataset/load_dataset.py | 待实现 | +| adapters.py | (trainer 层处理) | 待实现 | + +--- + +## 7. 关键设计决策记录 + +1. **schema 独立成文件** — ms-swift 将字段定义散落在 preprocessor 中，我们集中到 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 访问），性能优于逐行收集 + +--- + +## 8. 测试 + +测试文件：`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/paddleformers/cli/train/sft/__init__.py b/paddleformers/cli/train/sft/__init__.py index 5fca2fb42fa..9efce437d02 100644 --- a/paddleformers/cli/train/sft/__init__.py +++ b/paddleformers/cli/train/sft/__init__.py @@ -13,5 +13,6 @@ # limitations under the License. from .workflow import run_sft +from .workflow2 import run_sft_v2 -__all__ = ["run_sft"] +__all__ = ["run_sft", "run_sft_v2"] diff --git a/paddleformers/cli/train/sft/workflow2.py b/paddleformers/cli/train/sft/workflow2.py new file mode 100644 index 00000000000..086c7deb487 --- /dev/null +++ b/paddleformers/cli/train/sft/workflow2.py @@ -0,0 +1,377 @@ +# 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 + +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, + 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.""" + if data_args.template: + return data_args.template + + # Auto-detect from model name + 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" + + # Default: use jinja if tokenizer has chat_template, else chatml + 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 + + 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...") + + # Stash model path for template detection + 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}") + + # Encode config + encode_config = EncodeConfig( + max_seq_len=data_args.max_seq_len, + truncation="right", + label_shift=True, + ) + + # Build encode function + 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) + + 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) + + # Auto-split: if eval path is same as train or not set, split from train + 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)}") + + train_dataset = LazyEncodeDataset(hf_ds, encode_fn, seed=training_args.seed) + logger.info(f"[datasets_v2] Train dataset loaded: {len(train_dataset)} samples") + + # Load eval dataset (only if not auto-split above) + 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"[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 + ) + logger.info(f"[datasets_v2] max_seq_len for collate: {max_seq_len}") + + # Determine whether to use flashmask compact format + use_startend = getattr(model_args, "use_attn_mask_startend_row_indices", True) + + data_collator = partial( + collate_sft, + pad_token_id=tokenizer.pad_token_id, + max_seq_len=data_args.max_seq_len, + packing=data_args.packing, + 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 training_args.benchmark: + 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..f3fb3273fde 100644 --- a/paddleformers/cli/train/tuner.py +++ b/paddleformers/cli/train/tuner.py @@ -19,7 +19,7 @@ from ..hparams import get_train_args, read_args from .auto_parallel import run_auto_parallel from .dpo import run_dpo -from .sft import run_sft +from .sft import run_sft, run_sft_v2 def check_path(path): @@ -51,6 +51,9 @@ 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"]: + with paddle.amp.auto_cast(enable=False): + run_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) diff --git a/paddleformers/datasets_v2/__init__.py b/paddleformers/datasets_v2/__init__.py new file mode 100644 index 00000000000..99342041045 --- /dev/null +++ b/paddleformers/datasets_v2/__init__.py @@ -0,0 +1,70 @@ +# 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, +) + +# 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 ( + EncodeConfig, + EncodedSample, + TemplateMeta, + collate_sft, + encode_multiturn, + encode_sft, + get_template, + greedy_pack, + list_templates, + register_template, +) +from .dataset import LazyEncodeDataset diff --git a/paddleformers/datasets_v2/datapipe/__init__.py b/paddleformers/datasets_v2/datapipe/__init__.py new file mode 100644 index 00000000000..03443e22817 --- /dev/null +++ b/paddleformers/datasets_v2/datapipe/__init__.py @@ -0,0 +1,28 @@ +# 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_sft +from .encode import EncodeConfig, EncodedSample, encode_sft +from .packing import greedy_pack +from .template import ( + Slot, + TemplateMeta, + encode_multiturn, + encode_multiturn_jinja, + get_template, + list_templates, + register_template, +) diff --git a/paddleformers/datasets_v2/datapipe/collate.py b/paddleformers/datasets_v2/datapipe/collate.py new file mode 100644 index 00000000000..820f8ba7dda --- /dev/null +++ b/paddleformers/datasets_v2/datapipe/collate.py @@ -0,0 +1,199 @@ +# 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 both non-packed (simple padding) and packed (greedy packing with +block-diagonal attention mask) modes. + +Output format is compatible with PaddleFormers' existing Trainer. +""" + +from typing import Dict, List + +import numpy as np + +from .encode import EncodedSample +from .packing import greedy_pack + + +def collate_sft( + batch: List[EncodedSample], + pad_token_id: int, + max_seq_len: int, + packing: bool = False, + use_attn_mask_startend_row_indices: 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 greedy packing to combine short sequences. + use_attn_mask_startend_row_indices: If True, output compact flashmask + format (attn_mask_startend_row_indices) instead of 4D attention_mask. + + 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 use_attn_mask_startend_row_indices=False) + - attn_mask_startend_row_indices: [B, 1, S, 1] int32 (when use_attn_mask_startend_row_indices=True) + """ + if packing: + return _collate_packed(batch, pad_token_id, max_seq_len, use_attn_mask_startend_row_indices) + else: + return _collate_simple(batch, pad_token_id, max_seq_len, use_attn_mask_startend_row_indices) + + +def _collate_simple( + batch: List[EncodedSample], + pad_token_id: int, + max_seq_len: int, + use_startend: 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 + + return result + + +def _collate_packed( + batch: List[EncodedSample], + pad_token_id: int, + max_seq_len: int, + use_startend: bool = False, +) -> Dict[str, np.ndarray]: + """Packed collation: bin samples together, block-diagonal attention mask.""" + # Pack samples into groups + groups = greedy_pack(batch, max_seq_len) + 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 (needed for startend format) + 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 + + return result + + +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 diff --git a/paddleformers/datasets_v2/datapipe/encode.py b/paddleformers/datasets_v2/datapipe/encode.py new file mode 100644 index 00000000000..55ff9f4b060 --- /dev/null +++ b/paddleformers/datasets_v2/datapipe/encode.py @@ -0,0 +1,140 @@ +# 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 SFT encoding: messages → input_ids + labels. + +Takes a datasets_v2 row (with 'messages' column in OpenAI format) and produces +token IDs with loss masking, ready for collation into a training batch. +""" + +import logging +from dataclasses import dataclass +from typing import Any, Dict, List, Optional + +from .template import TemplateMeta, encode_multiturn, encode_multiturn_jinja + +logger = logging.getLogger(__name__) + + +@dataclass +class EncodeConfig: + """Configuration for SFT encoding.""" + + max_seq_len: int = 4096 + truncation: str = "right" # "right" | "left" + label_shift: bool = True # causal LM: labels = labels[1:] + [-100] + + +@dataclass +class EncodedSample: + """Output of encode_sft.""" + + input_ids: List[int] + labels: List[int] # -100 for no-loss positions + seq_len: int # actual token count (before any padding) + + +def encode_sft( + example: Dict[str, Any], + tokenizer: Any, + template: Optional[TemplateMeta], + config: EncodeConfig, +) -> Optional[EncodedSample]: + """Encode a single datasets_v2 row into token_ids + labels for SFT. + + Args: + example: Dict with 'messages' key. Each message has: + - role: "user" | "assistant" | "system" + - content: str + - loss (optional): bool, whether this assistant turn contributes to loss. + Defaults to True for assistant turns. + tokenizer: Tokenizer with encode() and convert_tokens_to_ids(). + template: TemplateMeta for custom encoding. If None, uses jinja path. + config: EncodeConfig. + + Returns: + EncodedSample or None if the sample is invalid. + """ + messages = example.get("messages") + if not messages or len(messages) < 2: + return None + + # Extract per-turn loss mask (for assistant turns only) + loss_mask = _extract_loss_mask(messages) + + # Encode messages → (prompt_ids, response_ids) pairs + if template is not None: + pairs = encode_multiturn(template, tokenizer, messages) + else: + pairs = encode_multiturn_jinja(tokenizer, messages) + + if not pairs: + return None + + # Build token_ids and labels + token_ids: List[int] = [] + labels: List[int] = [] + + for turn_idx, (prompt_ids, response_ids) in enumerate(pairs): + # Prompt tokens: no loss + token_ids += prompt_ids + labels += [-100] * len(prompt_ids) + + # Response tokens: loss if enabled for this turn + token_ids += response_ids + if turn_idx < len(loss_mask) and not loss_mask[turn_idx]: + labels += [-100] * len(response_ids) + else: + labels += response_ids + + if not token_ids: + return None + + # Append EOS if template uses efficient_eos + if template and template.efficient_eos: + eos_id = tokenizer.eos_token_id + if eos_id is not None: + token_ids.append(eos_id) + labels.append(eos_id) + + # Label shift: labels[i] = labels[i+1] for causal LM + if config.label_shift: + labels = labels[1:] + [-100] + + # Truncation + if len(token_ids) > config.max_seq_len: + if config.truncation == "right": + token_ids = token_ids[: config.max_seq_len] + labels = labels[: config.max_seq_len] + elif config.truncation == "left": + token_ids = token_ids[-config.max_seq_len :] + labels = labels[-config.max_seq_len :] + else: + return None # skip samples that exceed max_seq_len + + seq_len = len(token_ids) + return EncodedSample(input_ids=token_ids, labels=labels, seq_len=seq_len) + + +def _extract_loss_mask(messages: List[Dict[str, Any]]) -> List[bool]: + """Extract per-assistant-turn loss mask from messages. + + Returns a list of booleans, one per assistant turn, indicating + whether that turn should contribute to the loss. + """ + mask: List[bool] = [] + for msg in messages: + if msg.get("role") == "assistant": + mask.append(msg.get("loss", True)) + return mask diff --git a/paddleformers/datasets_v2/datapipe/packing.py b/paddleformers/datasets_v2/datapipe/packing.py new file mode 100644 index 00000000000..3642744bed0 --- /dev/null +++ b/paddleformers/datasets_v2/datapipe/packing.py @@ -0,0 +1,74 @@ +# 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. + +"""Greedy packing: bin multiple short sequences into max_seq_len slots. + +Algorithm matches PaddleFormers' existing greedy_intokens approach: +largest-remaining-capacity first-fit. +""" + +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] diff --git a/paddleformers/datasets_v2/datapipe/template.py b/paddleformers/datasets_v2/datapipe/template.py new file mode 100644 index 00000000000..6a8ccc9817d --- /dev/null +++ b/paddleformers/datasets_v2/datapipe/template.py @@ -0,0 +1,426 @@ +# 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. +""" + +import logging +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional, Set, Tuple, Union + +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) + 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 + + +# ============================================================ +# 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, + chat_sep: str = "", + default_system: str = "", + suffix: Optional[List[str]] = None, + stop_tokens: Optional[List[str]] = None, + efficient_eos: bool = True, + 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 [], + chat_sep=chat_sep, + default_system=default_system, + suffix=suffix or [], + stop_tokens=stop_tokens or [], + efficient_eos=efficient_eos, + ) + + +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 + + +# ============================================================ +# Core encoding +# ============================================================ + + +def encode_multiturn( + template: TemplateMeta, + tokenizer: Any, + messages: List[Dict[str, str]], + system: Optional[str] = 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", "content": str}. + system: Override system message. If None, uses template.default_system. + + 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 + if system: + elements += _substitute_slots(template.system, content=system) + + role = message.get("role", "") + content = message.get("content", "") + + if role == "user": + elements += _substitute_slots(template.user, content=content) + elif role == "assistant": + elements += _substitute_slots(template.assistant, content=content) + if i < len(actual_messages) - 1 and template.chat_sep: + elements.append(template.chat_sep) + 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_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) + + 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) + 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 + # For v1, provide a straightforward implementation + 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) + 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) + 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.5 variant (assistant includes im_end in slot) +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|>"], + stop_tokens=["<|eot_id|>"], + efficient_eos=False, +) + +# 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｜>"], +) + +# GLM4 format +register_template( + "glm4", + user=["<|user|>\n{{content}}<|assistant|>\n"], + assistant=["\n{{content}}"], + system=["<|system|>\n{{content}}"], + stop_tokens=["<|endoftext|>", "<|user|>"], +) + +# Gemma format +register_template( + "gemma", + user=["user\n{{content}}\nmodel\n"], + assistant=["{{content}}"], + system=["{{content}}\n"], + chat_sep="\n", + suffix=[""], + stop_tokens=[""], +) + +# 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"], + chat_sep="<|im_end|>\n\n", + suffix=["<|im_end|>"], + stop_tokens=["<|im_end|>"], +) + +# 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}}"], +) 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/loaders.py b/paddleformers/datasets_v2/loaders.py new file mode 100644 index 00000000000..2644049fa13 --- /dev/null +++ b/paddleformers/datasets_v2/loaders.py @@ -0,0 +1,371 @@ +# 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 + +Corresponds to ms-swift's dataset/loader.py. +""" + +import logging +import os +from typing import Dict, List, Optional, Union + +from datasets import Dataset as HfMapDataset +from datasets import IterableDataset as HfIterableDataset # noqa: F401 +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 _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.""" + load_kwargs = { + "split": split, + "streaming": streaming, + "trust_remote_code": True, + **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, + MessagesPreprocessor, + ResponsePreprocessor, + ) + + _PREPROCESSOR_MAP = { + "auto": AutoPreprocessor, + "messages": MessagesPreprocessor, + "response": ResponsePreprocessor, + "alpaca": AlpacaPreprocessor, + } + + 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) + + +# ============================================================ +# 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_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. + + 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 (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: + if meta and meta.preprocessor is None: + # Explicitly disabled preprocessing via registry + preprocessor = None + elif meta: + preprocessor = _get_preprocessor(meta) + else: + # No registry meta: default to AutoPreprocessor + from .preprocessors import AutoPreprocessor + + preprocessor = AutoPreprocessor(columns=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/ops.py b/paddleformers/datasets_v2/ops.py new file mode 100644 index 00000000000..d96b0047105 --- /dev/null +++ b/paddleformers/datasets_v2/ops.py @@ -0,0 +1,155 @@ +# 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 operations: sample, split, concat, interleave, shuffle. + +Thin wrappers over HuggingFace datasets utilities with added support +for oversampling, streaming (IterableDataset), and unified API. + +Corresponds to ms-swift's dataset/utils.py + dataset/dataset_meta.py. +""" + +from typing import List, Optional, Tuple + +import numpy as np +from datasets import Dataset as HfMapDataset +from datasets import IterableDataset as HfIterableDataset # noqa: F401 +from datasets import concatenate_datasets as hf_concat +from datasets import interleave_datasets as hf_interleave + +from .schema import DATASET_TYPE + + +def sample_dataset( + dataset: HfMapDataset, + n: int, + shuffle: bool = True, + seed: Optional[int] = None, +) -> HfMapDataset: + """Sample n rows from dataset. Supports oversampling (n > len(dataset)). + + Args: + dataset: must be a Map-style Dataset (not IterableDataset) + n: target number of rows. If > len(dataset), repeats the dataset. + shuffle: whether to shuffle the sampled indices + seed: random seed + """ + if not isinstance(dataset, HfMapDataset): + raise TypeError("sample_dataset only supports Map-style Dataset, not IterableDataset.") + if n <= 0: + raise ValueError("Cannot sample non-positive number of rows") + if not isinstance(n, int): + raise ValueError("Cannot sample non-integer number of rows") + + rng = np.random.RandomState(seed) + length = len(dataset) + + if n <= length: + if shuffle: + idx = rng.permutation(length)[:n].tolist() + else: + idx = list(range(n)) + else: + # Oversampling: tile full dataset + remainder + repeats = n // length + remainder = n % length + idx = np.tile(np.arange(length), repeats) + if remainder > 0: + extra = rng.permutation(length)[:remainder] if shuffle else np.arange(remainder) + idx = np.concatenate([idx, extra]) + if shuffle: + rng.shuffle(idx) + idx = idx.tolist() + + return dataset.select(idx) + + +def split_dataset( + dataset: DATASET_TYPE, + test_ratio: float = 0.1, + shuffle: bool = True, + seed: Optional[int] = None, +) -> Tuple[DATASET_TYPE, DATASET_TYPE]: + """Split dataset into train and validation sets. + + Args: + dataset: HF Dataset or IterableDataset + test_ratio: fraction for validation (0.0 ~ 1.0) + shuffle: whether to shuffle before splitting (Map-style only) + seed: random seed + + Returns: + (train_dataset, val_dataset) + """ + if isinstance(dataset, HfMapDataset): + split = dataset.train_test_split(test_size=test_ratio, shuffle=shuffle, seed=seed) + return split["train"], split["test"] + else: + # IterableDataset: use take/skip (requires known length or estimate) + # For streaming, caller should provide dataset_sample to know the size + raise ValueError( + "split_dataset on IterableDataset requires known size. " + "Use sample first or manually call .take()/.skip()." + ) + + +def concat_datasets(datasets: List[DATASET_TYPE]) -> Optional[DATASET_TYPE]: + """Concatenate multiple datasets into one.""" + if not datasets: + return None + if len(datasets) == 1: + return datasets[0] + return hf_concat(datasets) + + +def interleave( + datasets: List[DATASET_TYPE], + probabilities: Optional[List[float]] = None, + seed: Optional[int] = None, + stopping_strategy: str = "first_exhausted", +) -> Optional[DATASET_TYPE]: + """Interleave multiple datasets with optional probability weights. + + Args: + datasets: list of datasets to interleave + probabilities: sampling probability for each dataset (must sum to 1) + seed: random seed + stopping_strategy: 'first_exhausted' or 'all_exhausted' + """ + if not datasets: + return None + if len(datasets) == 1: + return datasets[0] + return hf_interleave( + datasets, + probabilities=probabilities, + seed=seed, + stopping_strategy=stopping_strategy, + ) + + +def shuffle_dataset( + dataset: DATASET_TYPE, + seed: int = 42, + buffer_size: int = 1000, +) -> DATASET_TYPE: + """Shuffle a dataset. + + For Map-style Dataset: full in-memory shuffle. + For IterableDataset: streaming shuffle with buffer. + """ + if isinstance(dataset, HfMapDataset): + return dataset.shuffle(seed=seed) + else: + return dataset.shuffle(seed=seed, buffer_size=buffer_size) diff --git a/paddleformers/datasets_v2/preprocessors/__init__.py b/paddleformers/datasets_v2/preprocessors/__init__.py new file mode 100644 index 00000000000..428890fd9c5 --- /dev/null +++ b/paddleformers/datasets_v2/preprocessors/__init__.py @@ -0,0 +1,19 @@ +# 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 .auto import AutoPreprocessor +from .base import BasePreprocessor +from .extra import AlpacaPreprocessor +from .messages import MessagesPreprocessor +from .response import ResponsePreprocessor, history_to_messages diff --git a/paddleformers/datasets_v2/preprocessors/auto.py b/paddleformers/datasets_v2/preprocessors/auto.py new file mode 100644 index 00000000000..9b6f54e14d4 --- /dev/null +++ b/paddleformers/datasets_v2/preprocessors/auto.py @@ -0,0 +1,61 @@ +# 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. + +"""AutoPreprocessor: automatically selects the right preprocessor based on dataset columns. + +Corresponds to ms-swift's AutoPreprocessor (preprocessor/core.py:531). +""" + +from typing import Dict, Optional + +from ..schema import DATASET_TYPE +from .base import BasePreprocessor +from .extra import AlpacaPreprocessor +from .messages import MessagesPreprocessor +from .response import ResponsePreprocessor + + +class AutoPreprocessor: + """Auto-detect dataset format and dispatch to the appropriate preprocessor. + + Detection logic: + 1. Has 'messages' / 'conversation' / 'conversations' column → MessagesPreprocessor + 2. Has 'instruction' + 'input' columns → AlpacaPreprocessor + 3. Otherwise → ResponsePreprocessor + """ + + def __init__(self, *, columns: Optional[Dict[str, str]] = None, **kwargs) -> None: + self.columns = columns or {} + self.kwargs = kwargs + + def _get_preprocessor(self, dataset: DATASET_TYPE) -> BasePreprocessor: + col_names = dataset.column_names + for key in ["messages", "conversation", "conversations"]: + if key in col_names: + return MessagesPreprocessor(columns=self.columns, **self.kwargs) + if "instruction" in col_names and "input" in col_names: + return AlpacaPreprocessor(columns=self.columns, **self.kwargs) + return ResponsePreprocessor(columns=self.columns, **self.kwargs) + + def __call__( + self, + dataset: DATASET_TYPE, + *, + num_proc: int = 1, + batch_size: int = 1000, + strict: bool = False, + ) -> DATASET_TYPE: + dataset = BasePreprocessor._rename_columns(dataset, self.columns) + preprocessor = self._get_preprocessor(dataset) + return preprocessor(dataset, num_proc=num_proc, batch_size=batch_size, strict=strict) diff --git a/paddleformers/datasets_v2/preprocessors/base.py b/paddleformers/datasets_v2/preprocessors/base.py new file mode 100644 index 00000000000..e8ba5b1a9f6 --- /dev/null +++ b/paddleformers/datasets_v2/preprocessors/base.py @@ -0,0 +1,193 @@ +# 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. + +"""Base preprocessor for datasets_v2. + +Corresponds to ms-swift's RowPreprocessor (preprocessor/core.py:25). +Designed to work with HuggingFace Dataset.map(). +""" + +import logging +import traceback +from typing import Any, Dict, List, Optional, Union + +from datasets import Dataset as HfMapDataset + +from ..schema import ( + DATASET_TYPE, + STANDARD_KEYS, + cast_images, + cast_media_list, + check_messages, +) + +logger = logging.getLogger(__name__) + + +class BasePreprocessor: + """Base class for all dataset preprocessors. + + Subclasses must implement `preprocess(row) -> row | list[row] | None`. + + Usage: + preprocessor = MyPreprocessor(columns={'instruction': 'query'}) + dataset = preprocessor(dataset) + """ + + DEFAULT_COLUMN_ALIASES = { + "image": "images", + "video": "videos", + "audio": "audios", + } + + def __init__( + self, + *, + columns: Optional[Dict[str, str]] = None, + traceback_limit: int = 10, + ) -> None: + self.columns = {**self.DEFAULT_COLUMN_ALIASES, **(columns or {})} + self._traceback_counter = 0 + self.traceback_limit = traceback_limit + + def preprocess(self, row: Dict[str, Any]) -> Optional[Union[Dict[str, Any], List[Dict[str, Any]]]]: + """Process a single row. Must be implemented by subclasses. + + Returns: + - A dict: the processed row + - A list of dicts: one row expanded into multiple + - None: row is skipped + """ + raise NotImplementedError + + def prepare_dataset(self, dataset: DATASET_TYPE) -> DATASET_TYPE: + """Hook for subclasses to do dataset-level setup (e.g. download media).""" + return dataset + + # ================================================================ + # Batched processing (for Dataset.map) + # ================================================================ + + def _batched_preprocess(self, batched_row: Dict[str, Any], strict: bool = False) -> Dict[str, Any]: + """Process a batch of rows. Used as the fn for Dataset.map(batched=True).""" + rows = self._batched_to_rows(batched_row) + new_rows = [] + + for row in rows: + try: + result = self.preprocess(row) + if result is None: + result = [] + if isinstance(result, dict): + result = [result] + for r in result: + self._check_and_cast(r) + new_rows += result + except Exception: + if strict: + raise + if self._traceback_counter < self.traceback_limit: + logger.warning(f"Row preprocessing error (row will be skipped):\n{traceback.format_exc()}") + self._traceback_counter += 1 + + res = self._rows_to_batched(new_rows) + if len(res) == 0: + res["messages"] = [] + return res + + # ================================================================ + # __call__: the main entry point + # ================================================================ + + def __call__( + self, + dataset: DATASET_TYPE, + *, + num_proc: int = 1, + batch_size: int = 1000, + strict: bool = False, + ) -> DATASET_TYPE: + """Apply this preprocessor to a dataset. + + Args: + dataset: HuggingFace Dataset or IterableDataset + num_proc: number of parallel workers for map + batch_size: rows per batch in map + strict: if True, raise on first error instead of skipping + """ + dataset = self._rename_columns(dataset, self.columns) + dataset = self.prepare_dataset(dataset) + + map_kwargs = { + "batched": True, + "batch_size": batch_size, + "fn_kwargs": {"strict": strict}, + "remove_columns": self._columns_to_remove(dataset), + } + if isinstance(dataset, HfMapDataset): + map_kwargs["num_proc"] = num_proc + + dataset = dataset.map(self._batched_preprocess, **map_kwargs) + return dataset + + # ================================================================ + # Checks and casting (per row) + # ================================================================ + + @staticmethod + def _check_and_cast(row: Dict[str, Any]) -> None: + """Validate and normalize a single processed row in-place.""" + if "messages" in row: + check_messages(row["messages"]) + + for key in ["images", "rejected_images"]: + if key in row and row[key] is not None: + row[key] = cast_images(row[key]) + + for key in ["videos", "audios"]: + if key in row and row[key] is not None: + row[key] = cast_media_list(row[key]) + + # ================================================================ + # Column utilities + # ================================================================ + + @staticmethod + def _rename_columns(dataset: DATASET_TYPE, columns: Dict[str, str]) -> DATASET_TYPE: + """Rename columns, skipping those not present in dataset.""" + col_renames = {k: v for k, v in columns.items() if k in dataset.column_names and k != v} + col_renames = {k: v for k, v in col_renames.items() if v not in dataset.column_names} + return dataset.rename_columns(col_renames) if col_renames else dataset + + @staticmethod + def _columns_to_remove(dataset: DATASET_TYPE) -> List[str]: + """Get list of columns to remove after map (non-standard columns).""" + return [c for c in dataset.column_names if c not in STANDARD_KEYS] + + # ================================================================ + # Batch <-> rows conversion + # ================================================================ + + @staticmethod + def _batched_to_rows(batched_row: Dict[str, Any]) -> List[Dict[str, Any]]: + """Convert a batched dict to a list of row dicts.""" + return [dict(zip(batched_row, vals)) for vals in zip(*batched_row.values())] + + @staticmethod + def _rows_to_batched(rows: List[Dict[str, Any]]) -> Dict[str, Any]: + """Convert a list of row dicts to a batched dict.""" + if not rows: + return {} + all_keys = set().union(*rows) + return {k: [row.get(k) for row in rows] for k in all_keys} diff --git a/paddleformers/datasets_v2/preprocessors/extra.py b/paddleformers/datasets_v2/preprocessors/extra.py new file mode 100644 index 00000000000..3e20cda7759 --- /dev/null +++ b/paddleformers/datasets_v2/preprocessors/extra.py @@ -0,0 +1,55 @@ +# 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. + +"""Extra preprocessors: thin wrappers over the base preprocessors. + +Contains AlpacaPreprocessor and other lightweight derived classes. +""" + +from typing import Any, Dict, Optional + +from .response import ResponsePreprocessor + + +class AlpacaPreprocessor(ResponsePreprocessor): + """Preprocessor for Alpaca-style datasets (instruction/input/output). + + Concatenates instruction + input into query, then delegates to ResponsePreprocessor. + + Expected input: + {"instruction": "...", "input": "...", "output": "..."} + """ + + def __init__(self, *, columns: Optional[Dict[str, str]] = None, **kwargs) -> None: + super().__init__(columns=columns, **kwargs) + # Remove these from column rename map — we handle them manually in preprocess + for key in ("instruction", "input", "output"): + self.columns.pop(key, None) + + @staticmethod + def concat_inst_input(instruction: Optional[str], input_: Optional[str]) -> str: + if instruction and input_: + return f"{instruction}\n{input_}" + query = instruction or input_ + assert isinstance(query, str), f"query must be str, got: {query}" + return query + + def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]: + instruction = row.pop("instruction", None) + input_ = row.pop("input", None) + output = row.pop("output", None) + if output is not None: + row["response"] = output + row["query"] = self.concat_inst_input(instruction, input_) + return super().preprocess(row) diff --git a/paddleformers/datasets_v2/preprocessors/messages.py b/paddleformers/datasets_v2/preprocessors/messages.py new file mode 100644 index 00000000000..6268bfdae77 --- /dev/null +++ b/paddleformers/datasets_v2/preprocessors/messages.py @@ -0,0 +1,150 @@ +# 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. + +"""MessagesPreprocessor: handles datasets already in conversation format. + +Corresponds to ms-swift's MessagesPreprocessor (preprocessor/core.py:421). +Supports: standard messages, ShareGPT format, various key naming conventions. +""" + +import ast +from typing import Any, Callable, Dict, List, Optional, Union + +from .base import BasePreprocessor + + +def default_repair_messages(s: Union[str, Any]) -> Any: + """Parse stringified messages list back to Python objects.""" + if isinstance(s, str): + return ast.literal_eval(s) + return s + + +class MessagesPreprocessor(BasePreprocessor): + """Preprocessor for datasets that already have a conversation structure. + + Handles various naming conventions: + - role keys: 'role', 'from' + - content keys: 'content', 'value' + - user roles: 'user', 'human' + - assistant roles: 'assistant', 'gpt', 'bot' + + Also handles ShareGPT paired format: + [{'human': '...', 'gpt': '...'}, ...] + + Output: standard messages format. + """ + + def __init__( + self, + *, + role_key: Optional[str] = None, + content_key: Optional[str] = None, + user_role: Optional[str] = None, + assistant_role: Optional[str] = None, + system_role: str = "system", + repair_messages: Callable = default_repair_messages, + inner_key: Optional[str] = None, + columns: Optional[Dict[str, str]] = None, + **kwargs, + ) -> None: + super().__init__(columns=columns, **kwargs) + + self.role_keys = ["role", "from"] if role_key is None else [role_key] + self.content_keys = ["content", "value"] if content_key is None else [content_key] + self.user_roles = ["user", "human"] if user_role is None else [user_role] + self.assistant_roles = ["assistant", "gpt", "bot"] if assistant_role is None else [assistant_role] + self.tool_call_roles = ["function_call"] + self.tool_response_roles = ["function_response", "observation", "observations"] + + self.system_role = system_role + self.repair_messages = repair_messages + self.inner_key = inner_key + + for key in ["messages", "conversation", "conversations"]: + self.columns[key] = "messages" + for key in ["system", "system_prompt"]: + self.columns[key] = "system" + + def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]: + if "rejected_messages" in row and row["rejected_messages"] is not None: + sub = self.preprocess({"messages": row["rejected_messages"]}) + row["rejected_messages"] = sub["messages"] if sub else None + + messages = row["messages"] + if self.inner_key is not None: + messages = messages[self.inner_key] + + messages = self.repair_messages(messages) + if not messages or isinstance(messages, str): + return None + + self._to_std_key(messages, "role", self.role_keys) + self._to_std_key(messages, "content", self.content_keys) + + system = row.pop("system", None) + + if self._is_sharegpt_format(messages[0]): + messages = self._sharegpt_to_messages(messages, system) + else: + self._to_std_messages(messages, system) + + row["messages"] = messages + return row + + # ================================================================ + # Format detection and conversion + # ================================================================ + + @staticmethod + def _is_sharegpt_format(message: Dict[str, str]) -> bool: + return "role" not in message and "content" not in message + + def _sharegpt_to_messages(self, messages: List[Dict[str, str]], system: Optional[str]) -> List[Dict[str, str]]: + """Convert ShareGPT paired format to standard messages.""" + self._to_std_key(messages, "user", self.user_roles) + self._to_std_key(messages, "assistant", self.assistant_roles) + new_messages = [] + if system is not None: + new_messages.append({"role": "system", "content": system}) + for message in messages: + new_messages.append({"role": "user", "content": message["user"]}) + new_messages.append({"role": "assistant", "content": message["assistant"]}) + return new_messages + + def _to_std_messages(self, messages: List[Dict[str, str]], system: Optional[str]) -> None: + """Normalize role names in-place for standard messages format.""" + if messages[0]["role"] == self.system_role: + messages[0]["role"] = "system" + elif system is not None: + messages.insert(0, {"role": "system", "content": system}) + + for message in messages: + role = message["role"] + if role in self.user_roles: + message["role"] = "user" + elif role in self.assistant_roles: + message["role"] = "assistant" + elif role.replace("-", "_") in self.tool_call_roles: + message["role"] = "tool_call" + elif role.replace("-", "_") in self.tool_response_roles: + message["role"] = "tool_response" + + @staticmethod + def _to_std_key(messages: List[Dict[str, str]], std_key: str, optional_keys: List[str]) -> None: + """Rename variant keys to a standard key in each message dict.""" + for message in messages: + for key in optional_keys: + if key in message: + message[std_key] = message.pop(key) diff --git a/paddleformers/datasets_v2/preprocessors/response.py b/paddleformers/datasets_v2/preprocessors/response.py new file mode 100644 index 00000000000..852c4f4728a --- /dev/null +++ b/paddleformers/datasets_v2/preprocessors/response.py @@ -0,0 +1,98 @@ +# 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. + +"""ResponsePreprocessor: handles query/response format datasets. + +Corresponds to ms-swift's ResponsePreprocessor (preprocessor/core.py:356). +Supports: query + response + optional history + optional system. +""" + +import ast +from typing import Any, Dict, List, Optional + +from .base import BasePreprocessor + + +def history_to_messages(history: List[List[Optional[str]]], system: Optional[str] = None) -> List[Dict[str, str]]: + """Convert history pairs to messages list. + + Args: + history: [['query1', 'response1'], ['query2', 'response2'], ...] + Either element can be None (will be skipped). + system: optional system prompt + """ + messages = [] + if system is not None: + messages.append({"role": "system", "content": system}) + for pair in history: + assert isinstance(pair, (list, tuple)) + if pair[0] is not None: + messages.append({"role": "user", "content": pair[0]}) + if pair[1] is not None: + messages.append({"role": "assistant", "content": pair[1]}) + return messages + + +class ResponsePreprocessor(BasePreprocessor): + """Preprocessor for query/response format datasets. + + Handles datasets with columns like: + - query (or: prompt, input, instruction, question, problem) + - response (or: answer, output, targets, text, completion, ...) + - system (optional) + - history (optional): [['q1','r1'], ['q2','r2'], ...] + + Output: standard messages format. + """ + + SYSTEM_KEYS = ["system", "system_prompt"] + QUERY_KEYS = ["query", "prompt", "input", "instruction", "question", "problem"] + RESPONSE_KEYS = [ + "response", + "answer", + "output", + "targets", + "target", + "answer_key", + "answers", + "solution", + "text", + "completion", + "content", + ] + + def __init__(self, *, columns: Optional[Dict[str, str]] = None, **kwargs) -> None: + super().__init__(columns=columns, **kwargs) + for key in self.SYSTEM_KEYS: + self.columns[key] = "system" + for key in self.QUERY_KEYS: + self.columns[key] = "query" + for key in self.RESPONSE_KEYS: + self.columns[key] = "response" + + def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]: + response = row.pop("response", None) + if isinstance(response, (list, tuple)): + response = response[0] + + history = row.pop("history", None) or [] + query = row.pop("query", None) + system = row.pop("system", None) + + if isinstance(history, str): + history = ast.literal_eval(history) + + history.append([query, response]) + row["messages"] = history_to_messages(history, system) + return row diff --git a/paddleformers/datasets_v2/registry.py b/paddleformers/datasets_v2/registry.py new file mode 100644 index 00000000000..507e5b751cd --- /dev/null +++ b/paddleformers/datasets_v2/registry.py @@ -0,0 +1,246 @@ +# 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 registry: metadata definitions and registration. + +Provides a central registry for known datasets, enabling lookup by name +and bulk registration from JSON configuration files. + +Corresponds to ms-swift's dataset/register.py + dataset/dataset_meta.py. +""" + +import json +import logging +import os +from dataclasses import dataclass, field +from typing import Callable, Dict, List, Optional, Union + +from .schema import DATASET_TYPE + +logger = logging.getLogger(__name__) + + +@dataclass +class DatasetMeta: + """Metadata descriptor for a registered dataset. + + Attributes: + name: Unique identifier for the dataset in the registry. + path: Local file/directory path or HuggingFace Hub repo ID. + If None, `name` is used as the Hub ID when loading. + subset: HuggingFace dataset config/subset name. + split: Default split to load. + preprocessor: Preprocessor to apply after loading. + - "auto": use AutoPreprocessor (default) + - "messages": use MessagesPreprocessor + - "response": use ResponsePreprocessor + - "alpaca": use AlpacaPreprocessor + - None: skip preprocessing + - A callable: used directly + columns: Column rename mapping {src_name: dst_name}. + tags: Metadata tags for categorization. + """ + + name: str + path: Optional[str] = None + subset: Optional[str] = None + split: str = "train" + preprocessor: Union[str, Callable[..., DATASET_TYPE], None] = "auto" + columns: Optional[Dict[str, str]] = None + tags: List[str] = field(default_factory=list) + + +@dataclass +class DatasetSpec: + """Parsed result of a dataset string specification. + + Syntax: "dataset_name_or_path#sample_count" + + Examples: + "alpaca" -> name="alpaca", sample=None + "alpaca#500" -> name="alpaca", sample=500 + "/path/to/data.jsonl#1000" -> name="/path/to/data.jsonl", sample=1000 + """ + + name: str + sample: Optional[int] = None + + +# ============================================================ +# Module-level registry +# ============================================================ + +_DATASET_REGISTRY: Dict[str, DatasetMeta] = {} + + +# ============================================================ +# Registration functions +# ============================================================ + + +def register_dataset(meta: DatasetMeta, *, exist_ok: bool = False) -> None: + """Register a DatasetMeta into the global registry. + + Args: + meta: The dataset metadata to register. + exist_ok: If True, silently overwrites existing entries. + + Raises: + ValueError: If the name is already registered and exist_ok=False. + """ + if not exist_ok and meta.name in _DATASET_REGISTRY: + raise ValueError(f"Dataset '{meta.name}' is already registered. " f"Use exist_ok=True to overwrite.") + _DATASET_REGISTRY[meta.name] = meta + + +def get_dataset_meta(name: str) -> Optional[DatasetMeta]: + """Look up a registered dataset by name. + + Returns: + The DatasetMeta if found, otherwise None. + """ + return _DATASET_REGISTRY.get(name) + + +def list_datasets(tag: Optional[str] = None) -> List[str]: + """List all registered dataset names, optionally filtered by tag. + + Args: + tag: If provided, only return datasets containing this tag. + + Returns: + Sorted list of dataset names. + """ + if tag is None: + return sorted(_DATASET_REGISTRY.keys()) + return sorted(name for name, meta in _DATASET_REGISTRY.items() if tag in meta.tags) + + +def register_dataset_info(json_path: str) -> List[DatasetMeta]: + """Bulk-register datasets from a JSON file. + + The JSON file should contain a list of objects with fields matching + DatasetMeta: name, path, subset, split, preprocessor, columns, tags. + Only "name" is required; others use defaults. + + Relative paths in "path" are resolved relative to the JSON file's directory. + + Args: + json_path: Path to the JSON file. + + Returns: + List of registered DatasetMeta objects. + + Raises: + FileNotFoundError: If json_path does not exist. + ValueError: If an entry is missing the "name" field. + """ + if not os.path.isfile(json_path): + raise FileNotFoundError(f"Dataset info file not found: {json_path}") + + base_dir = os.path.dirname(os.path.abspath(json_path)) + + with open(json_path, "r", encoding="utf-8") as f: + entries = json.load(f) + + if not isinstance(entries, list): + raise ValueError(f"Expected a JSON list in {json_path}, got {type(entries).__name__}") + + registered = [] + for entry in entries: + if "name" not in entry: + raise ValueError(f"Dataset entry missing 'name' field: {entry}") + + # Resolve relative paths + path = entry.get("path") + if path and not os.path.isabs(path): + # Hub IDs contain "/" but are not relative paths (e.g. "org/repo") + # Only resolve paths that don't contain "/" or start with "./" + is_hub_id = "/" in path and not path.startswith(("./", "../")) + if not is_hub_id: + path = os.path.join(base_dir, path) + + meta = DatasetMeta( + name=entry["name"], + path=path, + subset=entry.get("subset"), + split=entry.get("split", "train"), + preprocessor=entry.get("preprocessor", "auto"), + columns=entry.get("columns"), + tags=entry.get("tags", []), + ) + register_dataset(meta, exist_ok=True) + registered.append(meta) + + logger.info(f"Registered {len(registered)} datasets from {json_path}") + return registered + + +# ============================================================ +# String syntax parsing +# ============================================================ + + +def parse_dataset_string(dataset_str: str) -> DatasetSpec: + """Parse a dataset specification string. + + Supports syntax: "name_or_path#N" where #N is optional sample count. + + Args: + dataset_str: The dataset string to parse. + + Returns: + A DatasetSpec with parsed name and optional sample count. + + Raises: + ValueError: If the sample count is not a positive integer. + """ + dataset_str = dataset_str.strip() + if "#" in dataset_str: + name, sample_str = dataset_str.rsplit("#", 1) + try: + sample = int(sample_str) + except ValueError: + raise ValueError( + f"Invalid sample count '{sample_str}' in '{dataset_str}'. " f"Expected a positive integer after '#'." + ) + if sample <= 0: + raise ValueError(f"Sample count must be positive, got {sample} in '{dataset_str}'.") + return DatasetSpec(name=name.strip(), sample=sample) + return DatasetSpec(name=dataset_str) + + +# ============================================================ +# Built-in registry loading +# ============================================================ + +_BUILTIN_LOADED = False + + +def _load_builtin_registry() -> None: + """Load the built-in dataset_info.json shipped with this package. + + Called once on first import. Safe to call multiple times (idempotent). + """ + global _BUILTIN_LOADED + if _BUILTIN_LOADED: + return + _BUILTIN_LOADED = True + + builtin_path = os.path.join(os.path.dirname(__file__), "dataset_info.json") + if os.path.isfile(builtin_path): + register_dataset_info(builtin_path) + logger.debug(f"Loaded built-in dataset registry from {builtin_path}") + else: + logger.debug("No built-in dataset_info.json found, skipping.") diff --git a/paddleformers/datasets_v2/schema.py b/paddleformers/datasets_v2/schema.py new file mode 100644 index 00000000000..1ae9a19d051 --- /dev/null +++ b/paddleformers/datasets_v2/schema.py @@ -0,0 +1,192 @@ +# 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. + +"""Standard schema definitions for datasets. + +All preprocessors normalize raw data into this schema. +After preprocessing, only fields defined here are retained. +""" + +from dataclasses import dataclass, field +from itertools import chain +from typing import Any, Dict, List, Literal, Optional, Union + +from datasets import Dataset as HfMapDataset +from datasets import IterableDataset as HfIterableDataset + +DATASET_TYPE = Union[HfMapDataset, HfIterableDataset] + +# ============================================================ +# Constants +# ============================================================ + +ROLES = ("system", "user", "assistant", "tool_call", "tool_response", "tool") +Role = Literal["system", "user", "assistant", "tool_call", "tool_response", "tool"] +# Core fields of a single sample +PAIR_KEYS = ["messages", "images", "videos", "audios", "tools", "objects"] +# Prefixed pair fields for preference learning +# Expands to: ['rejected_messages', 'rejected_images', ..., 'positive_videos', ...] +PREFIXED_PAIR_KEYS = list( + chain.from_iterable([f"{prefix}_{k}" for k in PAIR_KEYS] for prefix in ["rejected", "positive", "negative"]) +) +# Scalar fields +SCALAR_KEYS = ["rejected_response", "label", "channel", "margin", "teacher_prompt"] +STANDARD_KEYS = PAIR_KEYS + PREFIXED_PAIR_KEYS + SCALAR_KEYS + +# ============================================================ +# Type definitions +# ============================================================ + + +class Message(dict): + """A single message in a conversation. + + Keys: + role: one of ROLES + content: str + loss: Optional[bool] — whether to compute loss on this turn + """ + + pass + + +class ImageMedia(dict): + """Image media representation. + + Keys: + bytes: Optional[bytes] — raw image bytes, None if using path + path: str — file path to the image + """ + + pass + + +class GroundingObjects(dict): + """Grounding annotation for object detection / referring. + + Keys: + ref: List[str] + bbox: List[List[float]] — each is [x1, y1, x2, y2] or [x, y] + bbox_type: str + image_id: int + """ + + pass + + +# ============================================================ +# Standard row dataclass +# ============================================================ + + +@dataclass +class StandardRow: + """The canonical representation of a single sample after preprocessing.""" + + # Core conversation + messages: List[Dict[str, Any]] = field(default_factory=list) + + # Multimodal + images: Optional[List[Dict[str, Any]]] = None + videos: Optional[List[str]] = None + audios: Optional[List[str]] = None + + # Tool use + tools: Optional[List[Dict[str, Any]]] = None + + # Grounding + objects: Optional[Dict[str, Any]] = None + + # Preference (DPO / RLHF) + rejected_response: Optional[str] = None + rejected_messages: Optional[List[Dict[str, Any]]] = None + rejected_images: Optional[List[Dict[str, Any]]] = None + rejected_videos: Optional[List[str]] = None + rejected_audios: Optional[List[str]] = None + rejected_tools: Optional[List[Dict[str, Any]]] = None + rejected_objects: Optional[Dict[str, Any]] = None + + # Embedding / Reranking + positive_messages: Optional[List[Dict[str, Any]]] = None + positive_images: Optional[List[Dict[str, Any]]] = None + positive_videos: Optional[List[str]] = None + positive_audios: Optional[List[str]] = None + positive_tools: Optional[List[Dict[str, Any]]] = None + positive_objects: Optional[Dict[str, Any]] = None + + negative_messages: Optional[List[Dict[str, Any]]] = None + negative_images: Optional[List[Dict[str, Any]]] = None + negative_videos: Optional[List[str]] = None + negative_audios: Optional[List[str]] = None + negative_tools: Optional[List[Dict[str, Any]]] = None + negative_objects: Optional[Dict[str, Any]] = None + + # Classification / Reward + label: Optional[int] = None + channel: Optional[str] = None + margin: Optional[float] = None + + # Distillation + teacher_prompt: Optional[str] = None + + +# ============================================================ +# Checks +# ============================================================ + + +def check_messages(messages: List[Dict[str, Any]]) -> None: + """Check that a messages list conforms to schema.""" + assert len(messages) > 0, "empty messages" + for msg in messages: + assert "role" in msg, f'message missing "role": {msg}' + assert "content" in msg, f'message missing "content": {msg}' + assert msg["role"] in ROLES, f'invalid role "{msg["role"]}", must be one of {ROLES}' + assert msg["content"] is not None, f"message content is None: {msg}" + extra_keys = set(msg.keys()) - {"role", "content", "loss"} + assert not extra_keys, f"unexpected keys in message: {extra_keys}" + + +# ============================================================ +# Type casting utilities +# ============================================================ + + +def cast_images(images: Union[str, Dict, List]) -> List[Dict[str, Any]]: + """Normalize image input to List[ImageMedia] format.""" + if isinstance(images, str): + return [{"bytes": None, "path": images}] + if isinstance(images, dict): + return [images] + if isinstance(images, list): + result = [] + for img in images: + if isinstance(img, str): + result.append({"bytes": None, "path": img}) + else: + result.append(img) + return result + raise TypeError(f"unsupported images type: {type(images)}") + + +def cast_media_list(media: Union[str, List[str]]) -> List[str]: + """Normalize videos/audios to List[str] format.""" + if isinstance(media, str): + return [media] + return media + + +def remove_non_standard_keys(row: Dict[str, Any]) -> Dict[str, Any]: + """Keep only standard keys in a row dict.""" + return {k: v for k, v in row.items() if k in STANDARD_KEYS} diff --git a/tests/datasets_v2/__init__.py b/tests/datasets_v2/__init__.py new file mode 100644 index 00000000000..290f972cf31 --- /dev/null +++ b/tests/datasets_v2/__init__.py @@ -0,0 +1,13 @@ +# 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. diff --git a/tests/datasets_v2/debug_pipeline.py b/tests/datasets_v2/debug_pipeline.py new file mode 100644 index 00000000000..6b547e5174f --- /dev/null +++ b/tests/datasets_v2/debug_pipeline.py @@ -0,0 +1,142 @@ +# 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. + +"""Debug script: observe the full dataset loading pipeline step by step. + +Usage: + 1. Set breakpoints at the print() lines below + 2. Run via VSCode debugger using "Debug datasets_v2 Pipeline" config + 3. Or run directly: python tests/datasets_v2/debug_pipeline.py +""" + +import importlib +import json + +# Workaround for broken torchcodec +_original_find_spec = importlib.util.find_spec + + +def _patched_find_spec(name, *args, **kwargs): + if name == "torchcodec": + return None + return _original_find_spec(name, *args, **kwargs) + + +importlib.util.find_spec = _patched_find_spec + +# ============================================================ +# Step 0: Prepare test data +# ============================================================ + +test_file = "/tmp/debug_pipeline_data.jsonl" +test_data = [ + {"query": "什么是机器学习?", "response": "机器学习是人工智能的一个分支。"}, + {"query": "1+1=?", "response": "2", "system": "你是一个数学助手"}, + {"query": "hello", "response": "你好"}, +] + +with open(test_file, "w", encoding="utf-8") as f: + for row in test_data: + f.write(json.dumps(row, ensure_ascii=False) + "\n") + +print(f"=== Test data written to: {test_file} ===") +print(f" Rows: {len(test_data)}") +print() + +# ============================================================ +# Step 1: Raw loading (no preprocessing) +# ============================================================ + +from paddleformers.datasets_v2.loaders import _load_local_file + +raw_ds = _load_local_file(test_file) +print("=== Step 1: Raw loaded dataset ===") +print(f" Columns: {raw_ds.column_names}") +print(f" Num rows: {len(raw_ds)}") +for i, row in enumerate(raw_ds): + print(f" Row {i}: {row}") +print() + +# ============================================================ +# Step 2: AutoPreprocessor detects format +# ============================================================ + +from paddleformers.datasets_v2.preprocessors import AutoPreprocessor + +auto = AutoPreprocessor() + +# See what preprocessor it picks +from paddleformers.datasets_v2.preprocessors.base import BasePreprocessor + +renamed_ds = BasePreprocessor._rename_columns(raw_ds, auto.columns) +print("=== Step 2: After column rename ===") +print(f" Columns: {renamed_ds.column_names}") +for i, row in enumerate(renamed_ds): + print(f" Row {i}: {row}") +print() + +# Check which preprocessor AutoPreprocessor would choose +preprocessor = auto._get_preprocessor(renamed_ds) +print(f" Selected preprocessor: {type(preprocessor).__name__}") +print() + +# ============================================================ +# Step 3: Apply preprocessing row by row (manual simulation) +# ============================================================ + +print("=== Step 3: Row-by-row preprocessing ===") +for i in range(len(renamed_ds)): + row = dict(renamed_ds[i]) # Make a mutable copy + print(f" [Before] Row {i}: {row}") + result = preprocessor.preprocess(row) + print(f" [After] Row {i}: {result}") + print() + +# ============================================================ +# Step 4: Full pipeline via load_dataset +# ============================================================ + +from paddleformers.datasets_v2 import load_dataset + +final_ds = load_dataset(test_file) +print("=== Step 4: Final dataset (full pipeline) ===") +print(f" Columns: {final_ds.column_names}") +print(f" Num rows: {len(final_ds)}") +for i, row in enumerate(final_ds): + print(f" Row {i}: {row}") +print() + +# ============================================================ +# Step 5: With registered dataset + sampling +# ============================================================ + +from paddleformers.datasets_v2 import DatasetMeta, register_dataset + +register_dataset( + DatasetMeta( + name="debug_test", + path=test_file, + preprocessor="response", + tags=["debug"], + ) +) + +sampled_ds = load_dataset("debug_test#2") +print("=== Step 5: Registered + sampled (2 rows) ===") +print(f" Columns: {sampled_ds.column_names}") +print(f" Num rows: {len(sampled_ds)}") +for i, row in enumerate(sampled_ds): + print(f" Row {i}: {row}") + +print("\n=== Pipeline debug complete ===") diff --git a/tests/datasets_v2/test_datapipe.py b/tests/datasets_v2/test_datapipe.py new file mode 100644 index 00000000000..04b9764a530 --- /dev/null +++ b/tests/datasets_v2/test_datapipe.py @@ -0,0 +1,585 @@ +# 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. + +"""Tests for datasets_v2/datapipe and datasets_v2/dataset.""" + +import pytest + +# ============================================================ +# Mock tokenizer for testing (no real model needed) +# ============================================================ + + +class MockTokenizer: + """Minimal tokenizer mock that maps chars to IDs.""" + + def __init__(self): + self.bos_token_id = 1 + self.eos_token_id = 2 + self.pad_token_id = 0 + self._vocab = {"<|im_start|>": 100, "<|im_end|>": 101, "<|endoftext|>": 102} + + def encode(self, text, add_special_tokens=False): + # Simple: each character becomes an ID (ord + 1000) + return [ord(c) + 1000 for c in text] + + def convert_tokens_to_ids(self, token): + return self._vocab.get(token, 999) + + +# ============================================================ +# Test template.py +# ============================================================ + + +class TestTemplate: + def test_register_and_get(self): + from paddleformers.datasets_v2.datapipe.template import ( + get_template, + list_templates, + register_template, + ) + + register_template( + "test_tpl", + user=["USER:{{content}}\nBOT:"], + assistant=["{{content}}\n"], + exist_ok=True, + ) + tpl = get_template("test_tpl") + assert tpl.name == "test_tpl" + assert "test_tpl" in list_templates() + + def test_get_nonexistent_raises(self): + from paddleformers.datasets_v2.datapipe.template import get_template + + with pytest.raises(KeyError): + get_template("nonexistent_xyz") + + def test_builtin_templates_exist(self): + from paddleformers.datasets_v2.datapipe.template import list_templates + + templates = list_templates() + assert "chatml" in templates + assert "llama3" in templates + assert "deepseek3" in templates + + def test_encode_multiturn_basic(self): + from paddleformers.datasets_v2.datapipe.template import ( + encode_multiturn, + get_template, + ) + + tpl = get_template("chatml") + tokenizer = MockTokenizer() + messages = [ + {"role": "user", "content": "Hi"}, + {"role": "assistant", "content": "Hello"}, + ] + pairs = encode_multiturn(tpl, tokenizer, messages) + assert len(pairs) == 1 + prompt_ids, response_ids = pairs[0] + assert len(prompt_ids) > 0 + assert len(response_ids) > 0 + + def test_encode_multiturn_with_system(self): + from paddleformers.datasets_v2.datapipe.template import ( + encode_multiturn, + get_template, + ) + + tpl = get_template("chatml") + tokenizer = MockTokenizer() + messages = [ + {"role": "system", "content": "Be helpful"}, + {"role": "user", "content": "Hi"}, + {"role": "assistant", "content": "Hello"}, + ] + pairs = encode_multiturn(tpl, tokenizer, messages) + assert len(pairs) == 1 + # System is prepended to first prompt + prompt_ids, _ = pairs[0] + assert len(prompt_ids) > 0 + + def test_encode_multiturn_multi_turn(self): + from paddleformers.datasets_v2.datapipe.template import ( + encode_multiturn, + get_template, + ) + + tpl = get_template("chatml") + tokenizer = MockTokenizer() + messages = [ + {"role": "user", "content": "Hi"}, + {"role": "assistant", "content": "Hello"}, + {"role": "user", "content": "Bye"}, + {"role": "assistant", "content": "See ya"}, + ] + pairs = encode_multiturn(tpl, tokenizer, messages) + assert len(pairs) == 2 + + def test_slot_with_dict_token(self): + from paddleformers.datasets_v2.datapipe.template import ( + encode_multiturn, + get_template, + register_template, + ) + + register_template( + "test_dict_slot", + prefix=[{"token": "<|im_start|>"}], + user=["{{content}}"], + assistant=["{{content}}"], + exist_ok=True, + ) + tpl = get_template("test_dict_slot") + tokenizer = MockTokenizer() + messages = [ + {"role": "user", "content": "A"}, + {"role": "assistant", "content": "B"}, + ] + pairs = encode_multiturn(tpl, tokenizer, messages) + prompt_ids, _ = pairs[0] + # First token should be the dict-token ID (100 for <|im_start|>) + assert prompt_ids[0] == 100 + + def test_slot_with_set_bos(self): + from paddleformers.datasets_v2.datapipe.template import ( + encode_multiturn, + get_template, + register_template, + ) + + register_template( + "test_set_slot", + prefix=[{"bos_token"}], + user=["{{content}}"], + assistant=["{{content}}"], + exist_ok=True, + ) + tpl = get_template("test_set_slot") + tokenizer = MockTokenizer() + messages = [ + {"role": "user", "content": "X"}, + {"role": "assistant", "content": "Y"}, + ] + pairs = encode_multiturn(tpl, tokenizer, messages) + prompt_ids, _ = pairs[0] + assert prompt_ids[0] == tokenizer.bos_token_id + + +# ============================================================ +# Test encode.py +# ============================================================ + + +class TestEncode: + def test_encode_sft_basic(self): + from paddleformers.datasets_v2.datapipe import ( + EncodeConfig, + encode_sft, + get_template, + ) + + tokenizer = MockTokenizer() + tpl = get_template("chatml") + config = EncodeConfig(max_seq_len=4096) + example = { + "messages": [ + {"role": "user", "content": "Hi"}, + {"role": "assistant", "content": "Hello"}, + ] + } + sample = encode_sft(example, tokenizer, tpl, config) + assert sample is not None + assert len(sample.input_ids) == len(sample.labels) + assert sample.seq_len == len(sample.input_ids) + # Prompt positions should have -100 labels + assert -100 in sample.labels + + def test_encode_sft_truncation(self): + from paddleformers.datasets_v2.datapipe import ( + EncodeConfig, + encode_sft, + get_template, + ) + + tokenizer = MockTokenizer() + tpl = get_template("chatml") + config = EncodeConfig(max_seq_len=10) + example = { + "messages": [ + {"role": "user", "content": "This is a long message"}, + {"role": "assistant", "content": "This is also long"}, + ] + } + sample = encode_sft(example, tokenizer, tpl, config) + assert sample is not None + assert sample.seq_len <= 10 + + def test_encode_sft_loss_mask(self): + from paddleformers.datasets_v2.datapipe import EncodeConfig, encode_sft + from paddleformers.datasets_v2.datapipe.template import ( + get_template, + register_template, + ) + + register_template( + "test_no_eos", + user=["{{content}}"], + assistant=["{{content}}"], + efficient_eos=False, + exist_ok=True, + ) + tokenizer = MockTokenizer() + tpl = get_template("test_no_eos") + config = EncodeConfig(max_seq_len=4096, label_shift=False) + example = { + "messages": [ + {"role": "user", "content": "Q"}, + {"role": "assistant", "content": "A", "loss": False}, + ] + } + sample = encode_sft(example, tokenizer, tpl, config) + assert sample is not None + # All labels should be -100 (user has no loss, assistant has loss=False) + assert all(l == -100 for l in sample.labels) + + def test_encode_sft_empty_messages_returns_none(self): + from paddleformers.datasets_v2.datapipe import ( + EncodeConfig, + encode_sft, + get_template, + ) + + tokenizer = MockTokenizer() + tpl = get_template("chatml") + config = EncodeConfig(max_seq_len=4096) + assert encode_sft({"messages": []}, tokenizer, tpl, config) is None + assert encode_sft({"messages": [{"role": "user", "content": "hi"}]}, tokenizer, tpl, config) is None + + def test_encode_sft_eos_appended(self): + from paddleformers.datasets_v2.datapipe import ( + EncodeConfig, + encode_sft, + get_template, + ) + + tokenizer = MockTokenizer() + tpl = get_template("chatml") # efficient_eos=True + config = EncodeConfig(max_seq_len=4096, label_shift=False) + example = { + "messages": [ + {"role": "user", "content": "A"}, + {"role": "assistant", "content": "B"}, + ] + } + sample = encode_sft(example, tokenizer, tpl, config) + # Last token should be EOS + assert sample.input_ids[-1] == tokenizer.eos_token_id + + +# ============================================================ +# Test packing.py +# ============================================================ + + +class TestPacking: + def test_greedy_pack_basic(self): + from paddleformers.datasets_v2.datapipe import EncodedSample, greedy_pack + + samples = [ + EncodedSample(input_ids=[1, 2, 3], labels=[1, 2, 3], seq_len=3), + EncodedSample(input_ids=[4, 5], labels=[4, 5], seq_len=2), + EncodedSample(input_ids=[6, 7, 8], labels=[6, 7, 8], seq_len=3), + ] + groups = greedy_pack(samples, max_seq_len=5) + # Sample 0 (3) + Sample 1 (2) = 5, fits in one bin + # Sample 2 (3) in another bin + assert len(groups) == 2 + total_samples = sum(len(g) for g in groups) + assert total_samples == 3 + + def test_greedy_pack_all_fit(self): + from paddleformers.datasets_v2.datapipe import EncodedSample, greedy_pack + + samples = [ + EncodedSample(input_ids=[1], labels=[1], seq_len=1), + EncodedSample(input_ids=[2], labels=[2], seq_len=1), + EncodedSample(input_ids=[3], labels=[3], seq_len=1), + ] + groups = greedy_pack(samples, max_seq_len=10) + # All fit in one bin + assert len(groups) == 1 + assert len(groups[0]) == 3 + + def test_greedy_pack_each_alone(self): + from paddleformers.datasets_v2.datapipe import EncodedSample, greedy_pack + + samples = [ + EncodedSample(input_ids=[1, 2, 3, 4, 5], labels=[0] * 5, seq_len=5), + EncodedSample(input_ids=[1, 2, 3, 4, 5], labels=[0] * 5, seq_len=5), + ] + groups = greedy_pack(samples, max_seq_len=5) + # Each sample fills a bin exactly + assert len(groups) == 2 + + def test_greedy_pack_empty(self): + from paddleformers.datasets_v2.datapipe import greedy_pack + + assert greedy_pack([], max_seq_len=10) == [] + + +# ============================================================ +# Test collate.py +# ============================================================ + + +class TestCollate: + def test_collate_simple(self): + from paddleformers.datasets_v2.datapipe import EncodedSample, collate_sft + + batch = [ + EncodedSample(input_ids=[1, 2, 3], labels=[-100, 2, 3], seq_len=3), + EncodedSample(input_ids=[4, 5], labels=[-100, 5], seq_len=2), + ] + out = collate_sft(batch, pad_token_id=0, max_seq_len=10, packing=False) + assert out["input_ids"].shape == (2, 3) # padded to max in batch + assert out["labels"].shape == (2, 3) + assert out["position_ids"].shape == (2, 3) + assert out["attention_mask"].shape == (2, 1, 3, 3) + # Check padding + assert out["input_ids"][1, 2] == 0 # pad_token_id + assert out["labels"][1, 2] == -100 + + def test_collate_packed(self): + from paddleformers.datasets_v2.datapipe import EncodedSample, collate_sft + + batch = [ + EncodedSample(input_ids=[1, 2, 3], labels=[-100, 2, 3], seq_len=3), + EncodedSample(input_ids=[4, 5], labels=[-100, 5], seq_len=2), + EncodedSample(input_ids=[6, 7, 8], labels=[-100, 7, 8], seq_len=3), + ] + out = collate_sft(batch, pad_token_id=0, max_seq_len=8, packing=True) + # With max_seq_len=8: samples 0(3)+1(2)=5 fit together, sample 2(3) alone + assert out["input_ids"].shape[1] == 8 + assert out["attention_mask"].shape == (out["input_ids"].shape[0], 1, 8, 8) + + def test_collate_packed_block_diagonal(self): + from paddleformers.datasets_v2.datapipe import EncodedSample, collate_sft + + batch = [ + EncodedSample(input_ids=[1, 2], labels=[1, 2], seq_len=2), + EncodedSample(input_ids=[3, 4], labels=[3, 4], seq_len=2), + ] + out = collate_sft(batch, pad_token_id=0, max_seq_len=4, packing=True) + # Both samples packed into one group (2+2=4) + mask = out["attention_mask"][0, 0] # [4, 4] + # Position [2,0] should be 0 (second sub-seq can't attend to first) + assert mask[2, 0] == 0.0 + assert mask[2, 1] == 0.0 + # Position [2,2] and [3,2] should be 1 (within second block) + assert mask[2, 2] == 1.0 + assert mask[3, 2] == 1.0 + + def test_collate_position_ids_reset(self): + from paddleformers.datasets_v2.datapipe import EncodedSample, collate_sft + + batch = [ + EncodedSample(input_ids=[1, 2, 3], labels=[1, 2, 3], seq_len=3), + EncodedSample(input_ids=[4, 5], labels=[4, 5], seq_len=2), + ] + out = collate_sft(batch, pad_token_id=0, max_seq_len=5, packing=True) + pos = out["position_ids"][0] # packed: [0,1,2, 0,1] + assert pos[0] == 0 + assert pos[1] == 1 + assert pos[2] == 2 + assert pos[3] == 0 # resets for second sub-sequence + assert pos[4] == 1 + + +# ============================================================ +# Test dataset/lazy_dataset.py +# ============================================================ + + +class TestLazyDataset: + def test_basic_getitem(self): + from paddleformers.datasets_v2.datapipe import EncodedSample + from paddleformers.datasets_v2.dataset import LazyEncodeDataset + + class FakeDataset: + def __len__(self): + return 3 + + def __getitem__(self, idx): + return { + "messages": [ + {"role": "user", "content": f"Q{idx}"}, + {"role": "assistant", "content": f"A{idx}"}, + ] + } + + def fake_encode(row): + return EncodedSample(input_ids=[1, 2, 3], labels=[-100, 2, 3], seq_len=3) + + ds = LazyEncodeDataset(FakeDataset(), fake_encode) + assert len(ds) == 3 + sample = ds[0] + assert sample.seq_len == 3 + + def test_retry_on_failure(self): + from paddleformers.datasets_v2.datapipe import EncodedSample + from paddleformers.datasets_v2.dataset import LazyEncodeDataset + + call_count = [0] + + class FakeDataset: + def __len__(self): + return 5 + + def __getitem__(self, idx): + return {"idx": idx} + + def flaky_encode(row): + call_count[0] += 1 + if call_count[0] <= 1: # first call fails + raise ValueError("bad sample") + return EncodedSample(input_ids=[1], labels=[1], seq_len=1) + + ds = LazyEncodeDataset(FakeDataset(), flaky_encode, seed=42) + sample = ds[0] + assert sample is not None + assert call_count[0] >= 2 + + def test_all_retries_fail(self): + from paddleformers.datasets_v2.dataset import LazyEncodeDataset + + class FakeDataset: + def __len__(self): + return 3 + + def __getitem__(self, idx): + return {} + + def always_fail(row): + raise ValueError("always fails") + + ds = LazyEncodeDataset(FakeDataset(), always_fail, n_try_fetch=3) + with pytest.raises(RuntimeError, match="Failed to encode"): + ds[0] + + +# ============================================================ +# Integration test +# ============================================================ + + +class TestEndToEnd: + def test_full_pipeline(self): + """Test the complete flow: messages → encode → collate.""" + from functools import partial + + from paddleformers.datasets_v2.datapipe import ( + EncodeConfig, + collate_sft, + encode_sft, + get_template, + ) + + tokenizer = MockTokenizer() + tpl = get_template("chatml") + config = EncodeConfig(max_seq_len=256) + encode_fn = partial(encode_sft, tokenizer=tokenizer, template=tpl, config=config) + + examples = [ + { + "messages": [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi there!"}, + ] + }, + { + "messages": [ + {"role": "user", "content": "What is 2+2?"}, + {"role": "assistant", "content": "4"}, + ] + }, + { + "messages": [ + {"role": "user", "content": "Bye"}, + {"role": "assistant", "content": "Goodbye"}, + ] + }, + ] + + samples = [encode_fn(ex) for ex in examples] + assert all(s is not None for s in samples) + + # Test without packing + out = collate_sft(samples, pad_token_id=0, max_seq_len=256, packing=False) + assert out["input_ids"].shape[0] == 3 + assert out["input_ids"].shape[1] <= 256 + assert out["attention_mask"].shape == (3, 1, out["input_ids"].shape[1], out["input_ids"].shape[1]) + + # Test with packing + out_packed = collate_sft(samples, pad_token_id=0, max_seq_len=256, packing=True) + assert out_packed["input_ids"].shape[1] == 256 + assert "attention_mask" in out_packed + + def test_with_lazy_dataset(self): + """Test LazyEncodeDataset integration.""" + from functools import partial + + from paddleformers.datasets_v2.datapipe import ( + EncodeConfig, + collate_sft, + encode_sft, + get_template, + ) + from paddleformers.datasets_v2.dataset import LazyEncodeDataset + + tokenizer = MockTokenizer() + tpl = get_template("chatml") + config = EncodeConfig(max_seq_len=256) + encode_fn = partial(encode_sft, tokenizer=tokenizer, template=tpl, config=config) + + class FakeHFDataset: + def __init__(self): + self.data = [ + { + "messages": [ + {"role": "user", "content": f"Q{i}"}, + {"role": "assistant", "content": f"A{i}"}, + ] + } + for i in range(10) + ] + + def __len__(self): + return len(self.data) + + def __getitem__(self, idx): + return self.data[idx] + + ds = LazyEncodeDataset(FakeHFDataset(), encode_fn) + assert len(ds) == 10 + + # Simulate a mini batch + batch = [ds[i] for i in range(4)] + out = collate_sft(batch, pad_token_id=0, max_seq_len=256, packing=True) + assert out["input_ids"].ndim == 2 + assert out["labels"].ndim == 2 + assert out["position_ids"].ndim == 2 + assert out["attention_mask"].ndim == 4 diff --git a/tests/datasets_v2/test_loaders.py b/tests/datasets_v2/test_loaders.py new file mode 100644 index 00000000000..60d762300f7 --- /dev/null +++ b/tests/datasets_v2/test_loaders.py @@ -0,0 +1,365 @@ +# 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. + +"""Tests for datasets_v2 registry and loaders. + +Run with: python -m pytest tests/datasets_v2/test_loaders.py -v +""" + +import csv +import importlib +import json + +# Workaround: broken torchcodec residual in env +_original_find_spec = importlib.util.find_spec + + +def _patched_find_spec(name, *args, **kwargs): + if name == "torchcodec": + return None + return _original_find_spec(name, *args, **kwargs) + + +importlib.util.find_spec = _patched_find_spec + +import pytest +from datasets import Dataset # noqa: F401 + +from paddleformers.datasets_v2.loaders import ( + _detect_file_format, + _resolve_source, + load_dataset, + load_datasets, +) +from paddleformers.datasets_v2.registry import ( + _DATASET_REGISTRY, + DatasetMeta, + get_dataset_meta, + list_datasets, + parse_dataset_string, + register_dataset, + register_dataset_info, +) + +# ============================================================ +# Fixtures +# ============================================================ + + +@pytest.fixture(autouse=True) +def clean_registry(): + """Clear the registry before and after each test.""" + _DATASET_REGISTRY.clear() + yield + _DATASET_REGISTRY.clear() + + +@pytest.fixture +def jsonl_file(tmp_path): + """Create a temporary jsonl file with query/response data.""" + data = [ + {"query": "hello", "response": "hi there"}, + {"query": "what is 1+1", "response": "2"}, + {"query": "tell me a joke", "response": "why did the chicken cross the road?"}, + ] + path = tmp_path / "train.jsonl" + with open(path, "w") as f: + for row in data: + f.write(json.dumps(row) + "\n") + return str(path) + + +@pytest.fixture +def csv_file(tmp_path): + """Create a temporary CSV file with alpaca-style data.""" + path = tmp_path / "train.csv" + with open(path, "w", newline="") as f: + writer = csv.writer(f) + writer.writerow(["instruction", "input", "output"]) + writer.writerow(["translate", "hello", "hola"]) + writer.writerow(["summarize", "long text", "short"]) + return str(path) + + +@pytest.fixture +def messages_jsonl(tmp_path): + """Create a temporary jsonl file with messages format data.""" + data = [ + {"messages": [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hello"}]}, + {"messages": [{"role": "user", "content": "bye"}, {"role": "assistant", "content": "goodbye"}]}, + ] + path = tmp_path / "messages.jsonl" + with open(path, "w") as f: + for row in data: + f.write(json.dumps(row) + "\n") + return str(path) + + +@pytest.fixture +def dataset_info_json(tmp_path, jsonl_file): + """Create a temporary dataset_info.json for bulk registration.""" + info = [ + {"name": "ds_alpha", "path": jsonl_file, "preprocessor": "response", "tags": ["sft"]}, + {"name": "ds_beta", "path": jsonl_file, "preprocessor": "auto", "tags": ["sft", "zh"]}, + ] + path = tmp_path / "dataset_info.json" + with open(path, "w") as f: + json.dump(info, f) + return str(path) + + +# ============================================================ +# Registry Tests +# ============================================================ + + +class TestParseDatasetString: + def test_simple_name(self): + spec = parse_dataset_string("alpaca") + assert spec.name == "alpaca" + assert spec.sample is None + + def test_name_with_sample(self): + spec = parse_dataset_string("alpaca#500") + assert spec.name == "alpaca" + assert spec.sample == 500 + + def test_path_with_sample(self): + spec = parse_dataset_string("/data/train.jsonl#1000") + assert spec.name == "/data/train.jsonl" + assert spec.sample == 1000 + + def test_hub_id(self): + spec = parse_dataset_string("tatsu-lab/alpaca") + assert spec.name == "tatsu-lab/alpaca" + assert spec.sample is None + + def test_invalid_sample_not_int(self): + with pytest.raises(ValueError, match="Invalid sample count"): + parse_dataset_string("alpaca#abc") + + def test_invalid_sample_zero(self): + with pytest.raises(ValueError, match="positive"): + parse_dataset_string("alpaca#0") + + def test_invalid_sample_negative(self): + with pytest.raises(ValueError, match="positive"): + parse_dataset_string("alpaca#-1") + + def test_whitespace_stripped(self): + spec = parse_dataset_string(" alpaca#100 ") + assert spec.name == "alpaca" + assert spec.sample == 100 + + +class TestRegisterDataset: + def test_register_and_get(self): + meta = DatasetMeta(name="test_ds", path="/tmp/test.jsonl") + register_dataset(meta) + result = get_dataset_meta("test_ds") + assert result is meta + assert result.path == "/tmp/test.jsonl" + + def test_register_duplicate_raises(self): + register_dataset(DatasetMeta(name="dup", path="/a")) + with pytest.raises(ValueError, match="already registered"): + register_dataset(DatasetMeta(name="dup", path="/b")) + + def test_register_duplicate_exist_ok(self): + register_dataset(DatasetMeta(name="dup", path="/a")) + register_dataset(DatasetMeta(name="dup", path="/b"), exist_ok=True) + assert get_dataset_meta("dup").path == "/b" + + def test_get_nonexistent(self): + assert get_dataset_meta("no_such_ds") is None + + def test_list_all(self): + register_dataset(DatasetMeta(name="bb")) + register_dataset(DatasetMeta(name="aa")) + assert list_datasets() == ["aa", "bb"] + + def test_list_by_tag(self): + register_dataset(DatasetMeta(name="a", tags=["sft", "zh"])) + register_dataset(DatasetMeta(name="b", tags=["dpo"])) + register_dataset(DatasetMeta(name="c", tags=["sft"])) + assert list_datasets(tag="sft") == ["a", "c"] + assert list_datasets(tag="dpo") == ["b"] + assert list_datasets(tag="nonexist") == [] + + +class TestRegisterDatasetInfo: + def test_bulk_register(self, dataset_info_json): + metas = register_dataset_info(dataset_info_json) + assert len(metas) == 2 + assert get_dataset_meta("ds_alpha") is not None + assert get_dataset_meta("ds_beta") is not None + assert "sft" in get_dataset_meta("ds_alpha").tags + + def test_missing_file(self): + with pytest.raises(FileNotFoundError): + register_dataset_info("/nonexistent/path.json") + + def test_missing_name_field(self, tmp_path): + path = tmp_path / "bad.json" + with open(path, "w") as f: + json.dump([{"path": "/tmp/a.jsonl"}], f) + with pytest.raises(ValueError, match="missing 'name'"): + register_dataset_info(str(path)) + + +# ============================================================ +# Loader Tests +# ============================================================ + + +class TestDetectFileFormat: + def test_json(self): + assert _detect_file_format("/a/b.json") == "json" + + def test_jsonl(self): + assert _detect_file_format("/a/b.jsonl") == "json" + + def test_csv(self): + assert _detect_file_format("/a/b.csv") == "csv" + + def test_tsv(self): + assert _detect_file_format("/a/b.tsv") == "csv" + + def test_parquet(self): + assert _detect_file_format("/a/b.parquet") == "parquet" + + def test_txt(self): + assert _detect_file_format("/a/b.txt") == "text" + + def test_unsupported(self): + with pytest.raises(ValueError, match="Unsupported file format"): + _detect_file_format("/a/b.xyz") + + +class TestResolveSource: + def test_file(self, jsonl_file): + assert _resolve_source(jsonl_file) == "file" + + def test_directory(self, tmp_path): + assert _resolve_source(str(tmp_path)) == "directory" + + def test_hub(self): + assert _resolve_source("tatsu-lab/alpaca") == "hub" + + +class TestLoadDataset: + def test_local_jsonl_no_preprocess(self, jsonl_file): + ds = load_dataset(jsonl_file, preprocess=False) + assert len(ds) == 3 + assert "query" in ds.column_names + assert "response" in ds.column_names + + def test_local_jsonl_auto_preprocess(self, jsonl_file): + ds = load_dataset(jsonl_file) + assert len(ds) == 3 + assert ds.column_names == ["messages"] + row = ds[0] + assert row["messages"][0]["role"] == "user" + assert row["messages"][1]["role"] == "assistant" + + def test_local_csv_alpaca_preprocess(self, csv_file): + ds = load_dataset(csv_file) + assert len(ds) == 2 + assert ds.column_names == ["messages"] + msg = ds[0]["messages"] + assert msg[0]["role"] == "user" + assert "translate" in msg[0]["content"] + assert msg[1]["content"] == "hola" + + def test_local_messages_preprocess(self, messages_jsonl): + ds = load_dataset(messages_jsonl) + assert len(ds) == 2 + assert ds.column_names == ["messages"] + assert ds[0]["messages"][0]["content"] == "hi" + + def test_sampling_syntax(self, jsonl_file): + ds = load_dataset(jsonl_file + "#2", preprocess=False) + assert len(ds) == 2 + + def test_oversampling_syntax(self, jsonl_file): + ds = load_dataset(jsonl_file + "#10", preprocess=False) + assert len(ds) == 10 + + def test_registered_dataset(self, jsonl_file): + register_dataset( + DatasetMeta( + name="my_registered", + path=jsonl_file, + preprocessor="response", + tags=["test"], + ) + ) + ds = load_dataset("my_registered") + assert len(ds) == 3 + assert ds.column_names == ["messages"] + + def test_registered_with_none_preprocessor(self, jsonl_file): + register_dataset( + DatasetMeta( + name="raw_ds", + path=jsonl_file, + preprocessor=None, + ) + ) + ds = load_dataset("raw_ds") + assert "query" in ds.column_names + + def test_file_not_found(self): + with pytest.raises(FileNotFoundError): + load_dataset("/nonexistent/path.jsonl", preprocess=False) + + def test_streaming_mode(self, jsonl_file): + ds = load_dataset(jsonl_file, streaming=True, preprocess=False) + from datasets import IterableDataset + + assert isinstance(ds, IterableDataset) + rows = list(ds.take(2)) + assert len(rows) == 2 + + def test_streaming_with_sample(self, jsonl_file): + ds = load_dataset(jsonl_file + "#2", streaming=True, preprocess=False) + rows = list(ds) + assert len(rows) == 2 + + +class TestLoadDatasets: + def test_single_dataset(self, jsonl_file): + ds = load_datasets(jsonl_file, preprocess=False) + assert len(ds) == 3 + + def test_multiple_concat(self, jsonl_file, csv_file): + ds = load_datasets([jsonl_file, csv_file], preprocess=False) + assert len(ds) == 5 # 3 jsonl + 2 csv + + def test_multiple_with_preprocess(self, jsonl_file, messages_jsonl): + ds = load_datasets([jsonl_file, messages_jsonl]) + assert len(ds) == 5 # 3 + 2 + assert ds.column_names == ["messages"] + + +class TestLoadDirectory: + def test_directory_with_jsonl_files(self, tmp_path): + data = [{"query": "q1", "response": "r1"}, {"query": "q2", "response": "r2"}] + for i in range(2): + path = tmp_path / f"shard_{i}.jsonl" + with open(path, "w") as f: + for row in data: + f.write(json.dumps(row) + "\n") + ds = load_dataset(str(tmp_path), preprocess=False) + assert len(ds) == 4 # 2 files x 2 rows diff --git a/tests/datasets_v2/test_preprocessors.py b/tests/datasets_v2/test_preprocessors.py new file mode 100644 index 00000000000..310e3f87edd --- /dev/null +++ b/tests/datasets_v2/test_preprocessors.py @@ -0,0 +1,316 @@ +# 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. + +"""Basic tests for datasets_v2 preprocessors. + +Run with: python -m pytest tests/datasets_v2/test_preprocessors.py -v +Or debug via VSCode launch configuration "Test Preprocessors". +""" + +import importlib + +# Workaround: broken torchcodec residual in env triggers datasets import error. +# Patch find_spec to return None for torchcodec so datasets thinks it's not installed. +_original_find_spec = importlib.util.find_spec + + +def _patched_find_spec(name, *args, **kwargs): + if name == "torchcodec": + return None + return _original_find_spec(name, *args, **kwargs) + + +importlib.util.find_spec = _patched_find_spec + +from datasets import Dataset + +from paddleformers.datasets_v2.preprocessors import ( + AlpacaPreprocessor, + AutoPreprocessor, + MessagesPreprocessor, + ResponsePreprocessor, +) + + +def make_dataset(data: dict) -> Dataset: + """Helper: create a HF Dataset from a dict of columns.""" + return Dataset.from_dict(data) + + +# ============================================================ +# ResponsePreprocessor +# ============================================================ + + +def test_response_basic(): + """Basic query/response → messages conversion.""" + ds = make_dataset( + { + "query": ["What is 2+2?", "Hello"], + "response": ["4", "Hi there"], + } + ) + preprocessor = ResponsePreprocessor() + result = preprocessor(ds) + + assert "messages" in result.column_names + row0 = result[0] + assert row0["messages"] == [ + {"role": "user", "content": "What is 2+2?"}, + {"role": "assistant", "content": "4"}, + ] + + +def test_response_with_system_and_history(): + """Query/response with system prompt and history.""" + ds = make_dataset( + { + "query": ["Current question"], + "response": ["Current answer"], + "system": ["You are a helpful assistant."], + "history": ["[['q1', 'r1']]"], + } + ) + preprocessor = ResponsePreprocessor() + result = preprocessor(ds) + + row = result[0] + assert row["messages"] == [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "q1"}, + {"role": "assistant", "content": "r1"}, + {"role": "user", "content": "Current question"}, + {"role": "assistant", "content": "Current answer"}, + ] + + +def test_response_column_aliases(): + """Columns like 'instruction' or 'answer' should auto-map.""" + ds = make_dataset( + { + "instruction": ["Translate hello to French"], + "answer": ["Bonjour"], + } + ) + preprocessor = ResponsePreprocessor() + result = preprocessor(ds) + + row = result[0] + assert row["messages"][0]["content"] == "Translate hello to French" + assert row["messages"][1]["content"] == "Bonjour" + + +# ============================================================ +# MessagesPreprocessor +# ============================================================ + + +def test_messages_standard(): + """Standard messages format passthrough.""" + ds = make_dataset( + { + "messages": [ + [ + {"role": "user", "content": "Hi"}, + {"role": "assistant", "content": "Hello!"}, + ] + ], + } + ) + preprocessor = MessagesPreprocessor() + result = preprocessor(ds) + + row = result[0] + assert row["messages"] == [ + {"role": "user", "content": "Hi"}, + {"role": "assistant", "content": "Hello!"}, + ] + + +def test_messages_sharegpt_format(): + """ShareGPT paired format conversion.""" + ds = make_dataset( + { + "conversations": [ + [ + {"human": "What is AI?", "gpt": "Artificial Intelligence."}, + {"human": "Thanks", "gpt": "You are welcome."}, + ] + ], + } + ) + preprocessor = MessagesPreprocessor() + result = preprocessor(ds) + + row = result[0] + assert row["messages"] == [ + {"role": "user", "content": "What is AI?"}, + {"role": "assistant", "content": "Artificial Intelligence."}, + {"role": "user", "content": "Thanks"}, + {"role": "assistant", "content": "You are welcome."}, + ] + + +def test_messages_from_value_keys(): + """Messages with 'from'/'value' keys instead of 'role'/'content'.""" + ds = make_dataset( + { + "conversations": [ + [ + {"from": "human", "value": "Hello"}, + {"from": "gpt", "value": "Hi!"}, + ] + ], + } + ) + preprocessor = MessagesPreprocessor() + result = preprocessor(ds) + + row = result[0] + assert row["messages"] == [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi!"}, + ] + + +# ============================================================ +# AlpacaPreprocessor +# ============================================================ + + +def test_alpaca_basic(): + """Alpaca format: instruction + input + output.""" + ds = make_dataset( + { + "instruction": ["Translate the following sentence."], + "input": ["Hello, how are you?"], + "output": ["Bonjour, comment allez-vous?"], + } + ) + preprocessor = AlpacaPreprocessor() + result = preprocessor(ds) + + row = result[0] + assert row["messages"][0]["content"] == "Translate the following sentence.\nHello, how are you?" + assert row["messages"][1]["content"] == "Bonjour, comment allez-vous?" + + +def test_alpaca_no_input(): + """Alpaca format with empty input.""" + ds = make_dataset( + { + "instruction": ["Tell me a joke."], + "input": [""], + "output": ["Why did the chicken cross the road?"], + } + ) + preprocessor = AlpacaPreprocessor() + result = preprocessor(ds) + + row = result[0] + assert row["messages"][0]["content"] == "Tell me a joke." + + +# ============================================================ +# AutoPreprocessor +# ============================================================ + + +def test_auto_detects_messages(): + """Auto should pick MessagesPreprocessor when 'messages' column exists.""" + ds = make_dataset( + { + "messages": [ + [ + {"role": "user", "content": "Hi"}, + {"role": "assistant", "content": "Hello!"}, + ] + ], + } + ) + result = AutoPreprocessor()(ds) + assert result[0]["messages"][0]["role"] == "user" + + +def test_auto_detects_alpaca(): + """Auto should pick AlpacaPreprocessor when instruction+input columns exist.""" + ds = make_dataset( + { + "instruction": ["Say hi"], + "input": [""], + "output": ["Hi!"], + } + ) + result = AutoPreprocessor()(ds) + assert result[0]["messages"][1]["content"] == "Hi!" + + +def test_auto_detects_response(): + """Auto should pick ResponsePreprocessor as fallback.""" + ds = make_dataset( + { + "query": ["Hello"], + "response": ["Hi"], + } + ) + result = AutoPreprocessor()(ds) + assert result[0]["messages"][0] == {"role": "user", "content": "Hello"} + + +# ============================================================ +# Multimodal (images/videos preserved alongside messages) +# ============================================================ + + +def test_response_with_images(): + """Images column should be preserved and normalized.""" + ds = make_dataset( + { + "query": ["Describe this image"], + "response": ["A cat sitting on a mat"], + "images": [["path/to/cat.jpg"]], + } + ) + preprocessor = ResponsePreprocessor() + result = preprocessor(ds) + + row = result[0] + assert row["messages"][0]["content"] == "Describe this image" + assert row["images"] == [{"bytes": None, "path": "path/to/cat.jpg"}] + + +def test_messages_with_videos(): + """Videos column should be preserved and normalized.""" + ds = make_dataset( + { + "messages": [ + [ + {"role": "user", "content": "What happens in this video?"}, + {"role": "assistant", "content": "A dog is running."}, + ] + ], + "videos": [["path/to/dog.mp4"]], + } + ) + preprocessor = MessagesPreprocessor() + result = preprocessor(ds) + + row = result[0] + assert row["videos"] == ["path/to/dog.mp4"] + + +if __name__ == "__main__": + import pytest + + pytest.main([__file__, "-v"]) diff --git a/tests/datasets_v2/test_schema.py b/tests/datasets_v2/test_schema.py new file mode 100644 index 00000000000..290f972cf31 --- /dev/null +++ b/tests/datasets_v2/test_schema.py @@ -0,0 +1,13 @@ +# 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 6bded3551b0f4d225657e6c1ca6a926636857431 Mon Sep 17 00:00:00 2001 From: weiyixuanxx <864644273@qq.com> Date: Fri, 15 May 2026 11:18:31 +0800 Subject: [PATCH 02/17] First version of changes: the YAML config file requires setting stage: SFT_v2. Not all features are supported yet, and the overall pipeline still needs to be reviewed. Co-Authored-By: Claude Opus 4.6 --- paddleformers/cli/hparams/data_args.py | 4 + paddleformers/cli/train/sft/__init__.py | 3 +- paddleformers/cli/train/sft/workflow2.py | 26 +- paddleformers/cli/train/sft/workflow_vl_v2.py | 432 ++++++++++++++++++ paddleformers/cli/train/tuner.py | 5 + paddleformers/datasets_v2/__init__.py | 4 + .../datasets_v2/datapipe/__init__.py | 11 +- paddleformers/datasets_v2/datapipe/collate.py | 109 ++++- paddleformers/datasets_v2/datapipe/encode.py | 142 ++++++ .../datasets_v2/datapipe/template.py | 1 + paddleformers/datasets_v2/loaders.py | 2 - paddleformers/datasets_v2/ops.py | 2 - .../datasets_v2/preprocessors/auto.py | 5 +- .../datasets_v2/preprocessors/base.py | 1 - .../datasets_v2/preprocessors/messages.py | 1 - .../datasets_v2/preprocessors/response.py | 1 - paddleformers/datasets_v2/registry.py | 2 - 17 files changed, 728 insertions(+), 23 deletions(-) create mode 100644 paddleformers/cli/train/sft/workflow_vl_v2.py diff --git a/paddleformers/cli/hparams/data_args.py b/paddleformers/cli/hparams/data_args.py index e7a134634b7..e66f77c9e9c 100644 --- a/paddleformers/cli/hparams/data_args.py +++ b/paddleformers/cli/hparams/data_args.py @@ -190,3 +190,7 @@ class DataArguments: ) }, ) + mm_plugin: str = field( + default=None, + metadata={"help": "Multimodal plugin name (e.g. 'qwen2_vl'). Auto-detect from model type if None."}, + ) diff --git a/paddleformers/cli/train/sft/__init__.py b/paddleformers/cli/train/sft/__init__.py index 9efce437d02..9465692206f 100644 --- a/paddleformers/cli/train/sft/__init__.py +++ b/paddleformers/cli/train/sft/__init__.py @@ -14,5 +14,6 @@ from .workflow import run_sft from .workflow2 import run_sft_v2 +from .workflow_vl_v2 import run_vl_sft_v2 -__all__ = ["run_sft", "run_sft_v2"] +__all__ = ["run_sft", "run_sft_v2", "run_vl_sft_v2"] diff --git a/paddleformers/cli/train/sft/workflow2.py b/paddleformers/cli/train/sft/workflow2.py index 086c7deb487..414287c93b1 100644 --- a/paddleformers/cli/train/sft/workflow2.py +++ b/paddleformers/cli/train/sft/workflow2.py @@ -38,6 +38,7 @@ from paddleformers.datasets_v2 import ( EncodeConfig, LazyEncodeDataset, + encode_pt, encode_sft, get_template, ) @@ -65,6 +66,9 @@ 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 @@ -217,12 +221,22 @@ def run_sft_v2( ) # Build encode function - encode_fn = partial( - encode_sft, - tokenizer=tokenizer, - template=template, - config=encode_config, - ) + 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 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..026bf29d829 --- /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.template.mm_plugin import get_mm_plugin +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.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 training_args.benchmark: + 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 f3fb3273fde..f3a147d57bb 100644 --- a/paddleformers/cli/train/tuner.py +++ b/paddleformers/cli/train/tuner.py @@ -54,6 +54,11 @@ def _training_function(config: dict[str, Any]) -> None: elif model_args.stage in ["SFT-V2", "PT-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) diff --git a/paddleformers/datasets_v2/__init__.py b/paddleformers/datasets_v2/__init__.py index 99342041045..60db3528bee 100644 --- a/paddleformers/datasets_v2/__init__.py +++ b/paddleformers/datasets_v2/__init__.py @@ -59,9 +59,13 @@ EncodeConfig, EncodedSample, TemplateMeta, + VLEncodedSample, collate_sft, + collate_vl_sft, encode_multiturn, + encode_pt, encode_sft, + encode_vl_sft, get_template, greedy_pack, list_templates, diff --git a/paddleformers/datasets_v2/datapipe/__init__.py b/paddleformers/datasets_v2/datapipe/__init__.py index 03443e22817..7230886f786 100644 --- a/paddleformers/datasets_v2/datapipe/__init__.py +++ b/paddleformers/datasets_v2/datapipe/__init__.py @@ -14,8 +14,15 @@ """datapipe: template encoding + packing + collation for datasets_v2.""" -from .collate import collate_sft -from .encode import EncodeConfig, EncodedSample, encode_sft +from .collate import collate_sft, collate_vl_sft +from .encode import ( + EncodeConfig, + EncodedSample, + VLEncodedSample, + encode_pt, + encode_sft, + encode_vl_sft, +) from .packing import greedy_pack from .template import ( Slot, diff --git a/paddleformers/datasets_v2/datapipe/collate.py b/paddleformers/datasets_v2/datapipe/collate.py index 820f8ba7dda..2808a46c45f 100644 --- a/paddleformers/datasets_v2/datapipe/collate.py +++ b/paddleformers/datasets_v2/datapipe/collate.py @@ -20,7 +20,7 @@ Output format is compatible with PaddleFormers' existing Trainer. """ -from typing import Dict, List +from typing import Any, Callable, Dict, List, Optional import numpy as np @@ -197,3 +197,110 @@ def _build_startend_packed( if offset < max_seq_len: indices[i, 0, offset:, 0] = np.arange(offset, max_seq_len) return indices + + +# --------------------------------------------------------------------------- +# 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 diff --git a/paddleformers/datasets_v2/datapipe/encode.py b/paddleformers/datasets_v2/datapipe/encode.py index 55ff9f4b060..ddb85277d9f 100644 --- a/paddleformers/datasets_v2/datapipe/encode.py +++ b/paddleformers/datasets_v2/datapipe/encode.py @@ -127,6 +127,39 @@ def encode_sft( return EncodedSample(input_ids=token_ids, labels=labels, seq_len=seq_len) +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. + + Matches the old pipeline's _process_pretraining_sequence behavior: + 1. Extract messages[0]["content"] + 2. Tokenize + append EOS + 3. Truncate to max_seq_len + 1 + 4. input_ids = tokens[:-1], labels = tokens[1:] + """ + 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)) + tokens = tokens + [tokenizer.eos_token_id] + + if len(tokens) > config.max_seq_len + 1: + tokens = tokens[: config.max_seq_len + 1] + + input_ids = tokens[:-1] + labels = tokens[1:] + + return EncodedSample(input_ids=input_ids, labels=labels, seq_len=len(input_ids)) + + def _extract_loss_mask(messages: List[Dict[str, Any]]) -> List[bool]: """Extract per-assistant-turn loss mask from messages. @@ -138,3 +171,112 @@ def _extract_loss_mask(messages: List[Dict[str, Any]]) -> List[bool]: if msg.get("role") == "assistant": mask.append(msg.get("loss", True)) return mask + + +# --------------------------------------------------------------------------- +# VL (Vision-Language) encoding +# --------------------------------------------------------------------------- + + +@dataclass +class VLEncodedSample: + """Output of encode_vl_sft. Extends EncodedSample with vision fields.""" + + input_ids: List[int] + labels: List[int] + seq_len: int + mm_inputs: Dict[str, Any] + + +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. + + Args: + example: Dict with 'messages' and 'images' keys. + tokenizer: Tokenizer instance. + template: TemplateMeta for encoding. If None, uses jinja path. + config: EncodeConfig. + processor: AutoProcessor (contains image_processor). + mm_plugin: Multimodal plugin instance (from datasets/template/mm_plugin). + + Returns: + VLEncodedSample or None if the sample is invalid. + """ + messages = example.get("messages") + if not messages or len(messages) < 2: + return None + + images = example.get("images", []) + if not images: + return None + + # Normalize HF Image feature dicts back to path strings + 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, + ) + + loss_mask = _extract_loss_mask(messages) + + if template is not None: + pairs = encode_multiturn(template, tokenizer, messages) + else: + pairs = encode_multiturn_jinja(tokenizer, messages) + + if not pairs: + return None + + token_ids: List[int] = [] + labels: List[int] = [] + + 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: + labels += response_ids + + if not token_ids: + return None + + if template and template.efficient_eos: + eos_id = tokenizer.eos_token_id + if eos_id is not None: + token_ids.append(eos_id) + labels.append(eos_id) + + labels = mm_plugin.process_tokens(token_ids, processor) + + if config.label_shift: + labels = labels[1:] + [-100] + + if len(token_ids) > config.max_seq_len: + if config.truncation == "right": + token_ids = token_ids[: config.max_seq_len] + labels = labels[: config.max_seq_len] + elif config.truncation == "left": + token_ids = token_ids[-config.max_seq_len :] + labels = labels[-config.max_seq_len :] + else: + return None + + seq_len = len(token_ids) + return VLEncodedSample(input_ids=token_ids, labels=labels, seq_len=seq_len, mm_inputs=mm_inputs) diff --git a/paddleformers/datasets_v2/datapipe/template.py b/paddleformers/datasets_v2/datapipe/template.py index 6a8ccc9817d..8072cbf1d2f 100644 --- a/paddleformers/datasets_v2/datapipe/template.py +++ b/paddleformers/datasets_v2/datapipe/template.py @@ -423,4 +423,5 @@ def _encode_jinja_simple( "empty", user=["{{content}}"], assistant=["{{content}}"], + efficient_eos=False, ) diff --git a/paddleformers/datasets_v2/loaders.py b/paddleformers/datasets_v2/loaders.py index 2644049fa13..8df89e83d33 100644 --- a/paddleformers/datasets_v2/loaders.py +++ b/paddleformers/datasets_v2/loaders.py @@ -20,8 +20,6 @@ - HuggingFace Hub datasets (by repo_id) - Streaming mode (IterableDataset) - Automatic preprocessing via registry metadata - -Corresponds to ms-swift's dataset/loader.py. """ import logging diff --git a/paddleformers/datasets_v2/ops.py b/paddleformers/datasets_v2/ops.py index d96b0047105..e16bd33b5d3 100644 --- a/paddleformers/datasets_v2/ops.py +++ b/paddleformers/datasets_v2/ops.py @@ -16,8 +16,6 @@ Thin wrappers over HuggingFace datasets utilities with added support for oversampling, streaming (IterableDataset), and unified API. - -Corresponds to ms-swift's dataset/utils.py + dataset/dataset_meta.py. """ from typing import List, Optional, Tuple diff --git a/paddleformers/datasets_v2/preprocessors/auto.py b/paddleformers/datasets_v2/preprocessors/auto.py index 9b6f54e14d4..987942772f5 100644 --- a/paddleformers/datasets_v2/preprocessors/auto.py +++ b/paddleformers/datasets_v2/preprocessors/auto.py @@ -12,10 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""AutoPreprocessor: automatically selects the right preprocessor based on dataset columns. - -Corresponds to ms-swift's AutoPreprocessor (preprocessor/core.py:531). -""" +"""AutoPreprocessor: automatically selects the right preprocessor based on dataset columns.""" from typing import Dict, Optional diff --git a/paddleformers/datasets_v2/preprocessors/base.py b/paddleformers/datasets_v2/preprocessors/base.py index e8ba5b1a9f6..06e8a05057c 100644 --- a/paddleformers/datasets_v2/preprocessors/base.py +++ b/paddleformers/datasets_v2/preprocessors/base.py @@ -14,7 +14,6 @@ """Base preprocessor for datasets_v2. -Corresponds to ms-swift's RowPreprocessor (preprocessor/core.py:25). Designed to work with HuggingFace Dataset.map(). """ diff --git a/paddleformers/datasets_v2/preprocessors/messages.py b/paddleformers/datasets_v2/preprocessors/messages.py index 6268bfdae77..3acac76d1c6 100644 --- a/paddleformers/datasets_v2/preprocessors/messages.py +++ b/paddleformers/datasets_v2/preprocessors/messages.py @@ -14,7 +14,6 @@ """MessagesPreprocessor: handles datasets already in conversation format. -Corresponds to ms-swift's MessagesPreprocessor (preprocessor/core.py:421). Supports: standard messages, ShareGPT format, various key naming conventions. """ diff --git a/paddleformers/datasets_v2/preprocessors/response.py b/paddleformers/datasets_v2/preprocessors/response.py index 852c4f4728a..8790fe42f5d 100644 --- a/paddleformers/datasets_v2/preprocessors/response.py +++ b/paddleformers/datasets_v2/preprocessors/response.py @@ -14,7 +14,6 @@ """ResponsePreprocessor: handles query/response format datasets. -Corresponds to ms-swift's ResponsePreprocessor (preprocessor/core.py:356). Supports: query + response + optional history + optional system. """ diff --git a/paddleformers/datasets_v2/registry.py b/paddleformers/datasets_v2/registry.py index 507e5b751cd..63a91907c6b 100644 --- a/paddleformers/datasets_v2/registry.py +++ b/paddleformers/datasets_v2/registry.py @@ -16,8 +16,6 @@ Provides a central registry for known datasets, enabling lookup by name and bulk registration from JSON configuration files. - -Corresponds to ms-swift's dataset/register.py + dataset/dataset_meta.py. """ import json From 21ff444aa3d5ad2d158967cb20793f636a863ec6 Mon Sep 17 00:00:00 2001 From: weiyixuanxx <864644273@qq.com> Date: Mon, 18 May 2026 11:11:50 +0800 Subject: [PATCH 03/17] fix .gitignore --- .gitignore | 2 ++ 1 file changed, 2 insertions(+) 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/ From 9a196767bc40b6213105cfa56caf717fadaf548b Mon Sep 17 00:00:00 2001 From: weiyixuanxx <864644273@qq.com> Date: Mon, 18 May 2026 11:23:58 +0800 Subject: [PATCH 04/17] Add datasets_v2/dataset module (LazyEncodeDataset) Previously blocked by .gitignore global `dataset/` rule. This fixes the CI ModuleNotFoundError for paddleformers.datasets_v2.dataset. Co-Authored-By: Claude Opus 4.6 --- paddleformers/datasets_v2/dataset/__init__.py | 17 ++++ .../datasets_v2/dataset/lazy_dataset.py | 90 +++++++++++++++++++ 2 files changed, 107 insertions(+) create mode 100644 paddleformers/datasets_v2/dataset/__init__.py create mode 100644 paddleformers/datasets_v2/dataset/lazy_dataset.py diff --git a/paddleformers/datasets_v2/dataset/__init__.py b/paddleformers/datasets_v2/dataset/__init__.py new file mode 100644 index 00000000000..6f9ea218219 --- /dev/null +++ b/paddleformers/datasets_v2/dataset/__init__.py @@ -0,0 +1,17 @@ +# 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 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 From 9214b9bdaa1a01c128ea36d259f623fa654f60a8 Mon Sep 17 00:00:00 2001 From: weiyixuanxx <864644273@qq.com> Date: Thu, 21 May 2026 16:15:22 +0800 Subject: [PATCH 05/17] feat(datasets_v2): complete template features migration and pipeline enhancements Migrate all missing template features from old datasets/ pipeline into datasets_v2/ as independent code (no cross-imports): - ReasoningTemplate (encode_multiturn_reasoning, thought tag management) - Tool calling (tool_utils.py with 9 model-specific formatters) - Function/Observation role support in encode_multiturn - fix_special_tokens and parse_template utilities - Grounding plugin (grounding_plugin.py) - mm_plugin for VL-SFT support - 25 registered templates (qwen3/3.5/vl, glm4/moe/v, ernie/vl, etc.) Also includes: streaming dataset, packing improvements, collate enhancements, workflow2 updates, and comprehensive tests. Co-Authored-By: Claude Opus 4.6 --- paddleformers/cli/hparams/data_args.py | 10 + paddleformers/cli/train/sft/workflow2.py | 119 +- paddleformers/cli/train/sft/workflow_vl_v2.py | 2 +- paddleformers/datasets_v2/__init__.py | 15 +- paddleformers/datasets_v2/augment_utils.py | 103 ++ .../datasets_v2/datapipe/__init__.py | 7 +- paddleformers/datasets_v2/datapipe/collate.py | 313 +++- paddleformers/datasets_v2/datapipe/encode.py | 6 +- paddleformers/datasets_v2/datapipe/packing.py | 59 +- .../datasets_v2/datapipe/template.py | 640 ++++++- .../datasets_v2/datapipe/tool_utils.py | 357 ++++ paddleformers/datasets_v2/dataset/__init__.py | 1 + .../datasets_v2/dataset/streaming_dataset.py | 37 + paddleformers/datasets_v2/loaders.py | 25 +- paddleformers/datasets_v2/mm_plugin.py | 1529 +++++++++++++++++ paddleformers/datasets_v2/ops.py | 7 +- .../datasets_v2/preprocessors/__init__.py | 2 +- .../datasets_v2/preprocessors/auto.py | 18 +- .../datasets_v2/preprocessors/base.py | 14 +- .../datasets_v2/preprocessors/extra.py | 19 + paddleformers/datasets_v2/registry.py | 33 +- tests/datasets_v2/test_extended.py | 583 +++++++ tests/datasets_v2/test_ops.py | 270 +++ tests/datasets_v2/test_schema.py | 216 +++ 24 files changed, 4318 insertions(+), 67 deletions(-) create mode 100644 paddleformers/datasets_v2/augment_utils.py create mode 100644 paddleformers/datasets_v2/datapipe/tool_utils.py create mode 100644 paddleformers/datasets_v2/dataset/streaming_dataset.py create mode 100644 paddleformers/datasets_v2/mm_plugin.py create mode 100644 tests/datasets_v2/test_extended.py create mode 100644 tests/datasets_v2/test_ops.py diff --git a/paddleformers/cli/hparams/data_args.py b/paddleformers/cli/hparams/data_args.py index e66f77c9e9c..87b6d9d467c 100644 --- a/paddleformers/cli/hparams/data_args.py +++ b/paddleformers/cli/hparams/data_args.py @@ -194,3 +194,13 @@ class DataArguments: default=None, metadata={"help": "Multimodal plugin name (e.g. 'qwen2_vl'). Auto-detect from model type if None."}, ) + online: bool = field( + default=False, + metadata={ + "help": ( + "Whether to load dataset from online/network source (e.g. HuggingFace Hub) " + "instead of local files. When True, data is fetched from network. " + "When False (default), data is loaded from local paths specified in train_dataset_path." + ) + }, + ) diff --git a/paddleformers/cli/train/sft/workflow2.py b/paddleformers/cli/train/sft/workflow2.py index 414287c93b1..43f3ce6a654 100644 --- a/paddleformers/cli/train/sft/workflow2.py +++ b/paddleformers/cli/train/sft/workflow2.py @@ -38,6 +38,7 @@ from paddleformers.datasets_v2 import ( EncodeConfig, LazyEncodeDataset, + StreamingDataset, encode_pt, encode_sft, get_template, @@ -175,6 +176,27 @@ def run_sft_v2( 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...") @@ -242,30 +264,83 @@ def run_sft_v2( train_dataset = None eval_dataset = None num_proc = getattr(data_args, "dataset_num_proc", 1) + # dataset_type 控制数据集加载模式: + # iterable → streaming=True, 不下载到本地, 逐条从网络拉取(IterableDataset) + # map → streaming=False, 下载到HF cache后全量加载(MapDataset, 支持随机访问) + is_iterable = getattr(data_args, "dataset_type", "map") == "iterable" + + if is_iterable and getattr(data_args, "packing", False): + raise ValueError( + "packing/binpacking is not supported with dataset_type='iterable'. " + "Bin-packing requires knowing all sample lengths upfront, which is incompatible " + "with streaming (iterable) datasets. Please set dataset_type='map' to use packing, " + "or disable packing for streaming mode." + ) 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) + if is_iterable: + if training_args.max_steps <= 0: + raise ValueError("dataset_type=iterable requires --max_steps to be explicitly set (no len available).") + + # dataset_type=iterable → HF streaming=True，无论在线离线都返回 IterableDataset + hf_ds = v2_load_dataset(data_args.train_dataset_path, streaming=True, num_proc=1) + + # Use HF IterableDataset.map() for lazy encoding + 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 []) + # Filter out empty samples + hf_ds = hf_ds.filter(lambda x: x["seq_len"] > 0) + # Shuffle with buffer + hf_ds = hf_ds.shuffle(buffer_size=getattr(data_args, "buffer_size", 500)) + + train_dataset = StreamingDataset(hf_ds) + # Force single-thread DataLoader for iterable mode + training_args.dataloader_num_workers = 0 + logger.info("[datasets_v2] Train dataset ready (iterable mode, streaming=True)") + else: + # dataset_type=map → HF streaming=False，全量加载（在线则先下载到缓存） + hf_ds = v2_load_dataset(data_args.train_dataset_path, streaming=False, num_proc=num_proc) - # Auto-split: if eval path is same as train or not set, split from train - 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 - ) + # Auto-split: if eval path is same as train or not set, split from train + 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 + 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)}") + 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)}") - train_dataset = LazyEncodeDataset(hf_ds, encode_fn, seed=training_args.seed) - logger.info(f"[datasets_v2] Train dataset loaded: {len(train_dataset)} samples") + train_dataset = LazyEncodeDataset(hf_ds, encode_fn, seed=training_args.seed) + logger.info(f"[datasets_v2] Train dataset loaded: {len(train_dataset)} samples") # Load eval dataset (only if not auto-split above) 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"[datasets_v2] Eval dataset loaded: {len(eval_dataset)} samples") + if is_iterable: + # eval 也走流式，不下载到本地 + hf_eval_ds = v2_load_dataset(data_args.eval_dataset_path, streaming=True, num_proc=1) + + def _streaming_eval_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_eval_ds = hf_eval_ds.map(_streaming_eval_encode, remove_columns=hf_eval_ds.column_names or []) + hf_eval_ds = hf_eval_ds.filter(lambda x: x["seq_len"] > 0) + eval_dataset = StreamingDataset(hf_eval_ds) + logger.info("[datasets_v2] Eval dataset ready (iterable mode, streaming=True)") + else: + 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"[datasets_v2] Eval dataset loaded: {len(eval_dataset)} samples") # ====== Collate Function ====== max_seq_len = ( @@ -277,13 +352,19 @@ def run_sft_v2( # Determine whether to use flashmask compact format 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) ====== @@ -309,6 +390,8 @@ def run_sft_v2( # ====== 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 @@ -354,6 +437,12 @@ def run_sft_v2( 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) ] diff --git a/paddleformers/cli/train/sft/workflow_vl_v2.py b/paddleformers/cli/train/sft/workflow_vl_v2.py index 026bf29d829..c0a5587ffce 100644 --- a/paddleformers/cli/train/sft/workflow_vl_v2.py +++ b/paddleformers/cli/train/sft/workflow_vl_v2.py @@ -35,11 +35,11 @@ ModelArguments, ) from paddleformers.cli.utils.process import add_new_special_tokens -from paddleformers.datasets.template.mm_plugin import get_mm_plugin 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 ( diff --git a/paddleformers/datasets_v2/__init__.py b/paddleformers/datasets_v2/__init__.py index 60db3528bee..f2b63d36e90 100644 --- a/paddleformers/datasets_v2/__init__.py +++ b/paddleformers/datasets_v2/__init__.py @@ -26,6 +26,7 @@ BasePreprocessor, MessagesPreprocessor, ResponsePreprocessor, + TextPreprocessor, ) # Load built-in dataset registry on import @@ -58,17 +59,29 @@ from .datapipe import ( EncodeConfig, EncodedSample, + FunctionCall, TemplateMeta, VLEncodedSample, collate_sft, collate_vl_sft, 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 +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 index 7230886f786..f112c1507eb 100644 --- a/paddleformers/datasets_v2/datapipe/__init__.py +++ b/paddleformers/datasets_v2/datapipe/__init__.py @@ -23,13 +23,18 @@ encode_sft, encode_vl_sft, ) -from .packing import greedy_pack +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 index 2808a46c45f..ca90eabadd3 100644 --- a/paddleformers/datasets_v2/datapipe/collate.py +++ b/paddleformers/datasets_v2/datapipe/collate.py @@ -14,8 +14,11 @@ """Collation: batch of EncodedSamples → padded numpy dict for training. -Supports both non-packed (simple padding) and packed (greedy packing with -block-diagonal attention mask) modes. +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. """ @@ -25,7 +28,7 @@ import numpy as np from .encode import EncodedSample -from .packing import greedy_pack +from .packing import binpack_ffd, greedy_pack def collate_sft( @@ -33,7 +36,12 @@ def collate_sft( 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. @@ -43,22 +51,73 @@ def collate_sft( 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 greedy packing to combine short sequences. + 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 use_attn_mask_startend_row_indices=False) - - attn_mask_startend_row_indices: [B, 1, S, 1] int32 (when use_attn_mask_startend_row_indices=True) + - 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: - return _collate_packed(batch, pad_token_id, max_seq_len, use_attn_mask_startend_row_indices) + # 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, max_seq_len, use_attn_mask_startend_row_indices) + 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( @@ -66,6 +125,9 @@ def _collate_simple( 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) @@ -96,25 +158,56 @@ def _collate_simple( 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( - batch: List[EncodedSample], + 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: bin samples together, block-diagonal attention mask.""" - # Pack samples into groups - groups = greedy_pack(batch, max_seq_len) + """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 (needed for startend format) + # Track sub-sequence lengths for each group group_seq_lens: List[List[int]] = [] for i, group in enumerate(groups): @@ -151,9 +244,46 @@ def _collate_packed( 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, @@ -199,6 +329,163 @@ def _build_startend_packed( 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 # --------------------------------------------------------------------------- diff --git a/paddleformers/datasets_v2/datapipe/encode.py b/paddleformers/datasets_v2/datapipe/encode.py index ddb85277d9f..7f14c62064f 100644 --- a/paddleformers/datasets_v2/datapipe/encode.py +++ b/paddleformers/datasets_v2/datapipe/encode.py @@ -149,7 +149,11 @@ def encode_pt( return None tokens = tokenizer.convert_tokens_to_ids(tokenizer.tokenize(content)) - tokens = tokens + [tokenizer.eos_token_id] + 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] diff --git a/paddleformers/datasets_v2/datapipe/packing.py b/paddleformers/datasets_v2/datapipe/packing.py index 3642744bed0..c733ccd04ee 100644 --- a/paddleformers/datasets_v2/datapipe/packing.py +++ b/paddleformers/datasets_v2/datapipe/packing.py @@ -12,10 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Greedy packing: bin multiple short sequences into max_seq_len slots. +"""Packing algorithms: bin multiple short sequences into max_seq_len slots. -Algorithm matches PaddleFormers' existing greedy_intokens approach: -largest-remaining-capacity first-fit. +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 @@ -72,3 +73,55 @@ def greedy_pack( 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 index 8072cbf1d2f..8cbfb238156 100644 --- a/paddleformers/datasets_v2/datapipe/template.py +++ b/paddleformers/datasets_v2/datapipe/template.py @@ -19,12 +19,24 @@ 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: @@ -43,11 +55,20 @@ class TemplateMeta: 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 # ============================================================ @@ -64,11 +85,16 @@ def register_template( 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. @@ -91,11 +117,16 @@ def register_template( 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, ) @@ -161,6 +192,75 @@ def _slots_to_ids(tokenizer: Any, slots: List[Slot]) -> List[int]: return token_ids +# ============================================================ +# Tool call formatting helpers +# ============================================================ + + +def _format_function_content(content: str, 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. + """ + thought = None + if 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: + tool_calls = json.loads(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: + raise RuntimeError(f"Invalid JSON 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: str, tool_format: str) -> str: + """Format tools description using tool_utils.""" + tool_utils = get_tool_utils(tool_format) + try: + tools = json.loads(content) + return tool_utils.tool_formatter(tools) if len(tools) != 0 else "" + except json.JSONDecodeError: + raise RuntimeError(f"Invalid JSON 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") + + +def _get_thought_word_ids(tokenizer: Any, thought_words: Tuple[str, str]) -> List[int]: + """Get token IDs for empty thought (open + close tags).""" + return tokenizer.encode(f"{thought_words[0]}{thought_words[1]}", add_special_tokens=False) + + # ============================================================ # Core encoding # ============================================================ @@ -171,14 +271,16 @@ def encode_multiturn( tokenizer: Any, messages: List[Dict[str, str]], system: Optional[str] = None, + tools: Optional[str] = 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", "content": str}. + 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. @@ -201,9 +303,13 @@ def encode_multiturn( # Prefix (BOS, etc.) if template.prefix: elements += template.prefix - # System message - if system: - elements += _substitute_slots(template.system, content=system) + # System message (+ 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", "") @@ -211,9 +317,31 @@ def encode_multiturn( if role == "user": elements += _substitute_slots(template.user, content=content) elif role == "assistant": - elements += _substitute_slots(template.assistant, content=content) + # 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 + ) + 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 response + obs_slots = template.observation or template.user + elements += _substitute_slots(obs_slots, content=content) else: # Fallback: treat unknown roles as user elements += _substitute_slots(template.user, content=content) @@ -230,6 +358,250 @@ def encode_multiturn( return pairs +def encode_multiturn_reasoning( + template: TemplateMeta, + tokenizer: Any, + messages: List[Dict[str, str]], + system: Optional[str] = None, + tools: Optional[str] = None, +) -> List[Tuple[List[int], List[int]]]: + """Encode with reasoning/thinking support. + + Like encode_multiturn but handles thought tags: + - Removes thought from intermediate assistant messages + - For the last assistant message: if no thought present, adds empty + - If enable_thinking is False: thought goes into prompt (not trained on) + - If enable_thinking is True: thought goes into response (trained on) + + Args: + template: Template with thought_words and enable_thinking set. + tokenizer: Tokenizer. + messages: Conversation messages. + system: Override system message. + tools: Tools JSON string. + + Returns: + List of (prompt_ids, response_ids) tuples. + """ + 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:] + + # Remove thought from all intermediate assistant messages (keep only last) + for i in range(1, len(actual_messages) - 2, 2): + if actual_messages[i].get("role") == "assistant": + actual_messages[i]["content"] = _remove_thought(actual_messages[i]["content"], thought_words) + + # Handle last assistant message + if template.enable_thinking is False: + # Remove all CoT from the last response too + if actual_messages and actual_messages[-1].get("role") == "assistant": + actual_messages[-1]["content"] = _remove_thought(actual_messages[-1]["content"], thought_words) + + # Rebuild messages with system + if system: + rebuild = [{"role": "system", "content": system}] + actual_messages + else: + rebuild = actual_messages + + # Use normal encoding first + pairs = encode_multiturn(template, tokenizer, rebuild, system=system, tools=tools) + + # Now handle thought insertion for the last turn + if pairs: + last_prompt, last_response = pairs[-1] + last_msg = actual_messages[-1] if actual_messages else {} + last_content = last_msg.get("content", "") + + has_thought = thought_words[0].strip() in last_content and thought_words[1].strip() in last_content + + if not has_thought and last_msg.get("role") == "assistant": + thought_ids = _get_thought_word_ids(tokenizer, thought_words) + if not template.enable_thinking: + # Don't compute loss on thought → prepend to prompt + last_prompt = last_prompt + thought_ids + else: + # Compute loss on thought → prepend to response + last_response = thought_ids + last_response + pairs[-1] = (last_prompt, last_response) + + 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 = "" + + 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=True, + ) + + +# ============================================================ +# 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]], @@ -282,7 +654,6 @@ def encode_multiturn_jinja( i += 1 # Simpler approach: encode full conversation, then use incremental to split - # For v1, provide a straightforward implementation if not pairs: pairs = _encode_jinja_simple(tokenizer, messages) @@ -345,7 +716,114 @@ def _encode_jinja_simple( 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"], @@ -365,8 +843,15 @@ def _encode_jinja_simple( ], 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) @@ -377,6 +862,7 @@ def _encode_jinja_simple( system=["{{content}}\n\n"], prefix=[{"token": "<｜begin▁of▁sentence｜>"}], stop_tokens=["<｜end▁of▁sentence｜>"], + chat_sep="<｜end▁of▁sentence｜>", ) # GLM4 format @@ -385,18 +871,79 @@ def _encode_jinja_simple( 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", ) -# Gemma format +# GLM4 MOE register_template( - "gemma", - user=["user\n{{content}}\nmodel\n"], + "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=["{{content}}\n"], - chat_sep="\n", - suffix=[""], - stop_tokens=[""], + 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 @@ -405,9 +952,76 @@ def _encode_jinja_simple( 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 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 index 6f9ea218219..24531f4de9d 100644 --- a/paddleformers/datasets_v2/dataset/__init__.py +++ b/paddleformers/datasets_v2/dataset/__init__.py @@ -15,3 +15,4 @@ """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/streaming_dataset.py b/paddleformers/datasets_v2/dataset/streaming_dataset.py new file mode 100644 index 00000000000..e7ffa9b6af8 --- /dev/null +++ b/paddleformers/datasets_v2/dataset/streaming_dataset.py @@ -0,0 +1,37 @@ +# 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 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). + + Args: + hf_iterable: A HuggingFace IterableDataset instance (e.g. from + load_dataset(..., streaming=True) with .map()/.filter() applied). + """ + + def __init__(self, hf_iterable): + super().__init__() + self._hf = hf_iterable + + def __iter__(self): + for item in self._hf: + yield item diff --git a/paddleformers/datasets_v2/loaders.py b/paddleformers/datasets_v2/loaders.py index 8df89e83d33..84a0558d895 100644 --- a/paddleformers/datasets_v2/loaders.py +++ b/paddleformers/datasets_v2/loaders.py @@ -27,7 +27,6 @@ from typing import Dict, List, Optional, Union from datasets import Dataset as HfMapDataset -from datasets import IterableDataset as HfIterableDataset # noqa: F401 from datasets import load_dataset as hf_load_dataset from .registry import DatasetMeta, get_dataset_meta, parse_dataset_string @@ -155,6 +154,23 @@ def _load_local_directory( 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, *, @@ -165,10 +181,11 @@ def _load_hub_dataset( **kwargs, ) -> DATASET_TYPE: """Load a dataset from HuggingFace Hub.""" + _ensure_proxy() + load_kwargs = { "split": split, "streaming": streaming, - "trust_remote_code": True, **kwargs, } if subset is not None: @@ -198,6 +215,7 @@ def _get_preprocessor(meta: DatasetMeta): AutoPreprocessor, MessagesPreprocessor, ResponsePreprocessor, + TextPreprocessor, ) _PREPROCESSOR_MAP = { @@ -205,6 +223,7 @@ def _get_preprocessor(meta: DatasetMeta): "messages": MessagesPreprocessor, "response": ResponsePreprocessor, "alpaca": AlpacaPreprocessor, + "text": TextPreprocessor, } if preprocessor not in _PREPROCESSOR_MAP: @@ -263,7 +282,7 @@ def load_dataset( # 3. Resolve effective parameters effective_split = split or (meta.split if meta else "train") - effective_subset = subset or (meta.subset if meta else None) + 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) 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", "