数据灌输

一个用于多模态数据高效训练的库。最初专注于图像、文档和文本任务。

chug 目前利用 webdataset 和 Hugging Face 的 datasets。

webdataset tar 文件和数据集管道是可扩展预训练的首选。

支持 Hugging Face 的 datasets，非常适合探索、验证和微调用例。

chug 通过 pypdfium2 (https://github.com/pypdfium2-team/pypdfium2) 作为默认选项，或者如果您的用例可以接受 AGPL-3.0 许可，则使用 fitz/pymupdf (https://github.com/pymupdf/PyMuPDF) 提供即时 PDF 解码和渲染。必须手动启用 fitz 支持。PDF 处理在 webdataset 级别实现，因此您可以将其插入到其他 webdataset 管道中。这使得可以大规模分片流式传输原生 .pdf 文件，而无需预先渲染为 .png/.tiff 等格式。

状态

该库仍在开发中，将其视为 alpha 版本（预发布）。主要功能应该可以正常工作，该库已经通过我们即将公开的几个 PDF 数据集进行了测试。但是，请预期仍会有重大变更、大量改进等。

pip install --pre chug 将安装当前的开发版本。

待办事项

近期

清理和完善，代码库将发生变化
文档和单元测试
支持读取 .json/.yaml 信息文件，以自动解析 webdatasets 的分片信息（类似 timm）

中期

为利用布局的任务提供输出行（或单词+单词输出）的边界框注释选项
统一的图像和文本标记化预处理函数（图像+文本标记交错等）
类似 NaViT 的图像标记（补丁）打包。基于在线装箱算法，与图像预处理和管道集成。

长期

增加其他任务和建模需求的任务管道范围
支持额外的模态和目标（视频、音频、检测/密集像素目标、图像/视频/音频目标）
探索 .tar 分片的替代方案（array_record、arrow 等）

设计

子模块层次结构

该库的设计使得不同层级的函数、类可以独立使用。

如果想使用 JSON/YAML 可序列化配置构建加载器和管道，请使用 chug/loader.py 中的顶级 chug.create_loader()。根据数据集来源，可以轻松在 webdataset、HF datasets（未来还有其他来源）之间切换。

绕过最高级别，还可以调用 task_pipeline 中的 build_pipeline_* 方法，然后用一系列完整的参数调用 create_loader_wds，仅用于 wds 用例。

如果完全不想使用 chug 加载器和管道，image、text 和 wds（尤其是解码器）功能可能在其他项目中有用。

库模块（从高到低级别）

库内模块的依赖关系旨在遵循以下层次结构。例如，doc 依赖于 wds，但 wds 永远不应依赖于 doc。

app
|
loader (chug/loader.py)
|
task_pipeline
|
doc
|
wds, hfds, image, text
|
common

子模块

`common`

通用配置、结构（数据类）供整个库使用

`wds`

特定于 Webdataset（简称 wds）的代码。扩展和修改 webdataset 功能以适应涵盖的用例并提高鲁棒性。

chug 中的所有数据管道目前都利用 wds 管道，即使不使用 wds 数据集。

面向文档的解码（pdf 解码器）存在于 chug/wds/decode.py 中，它可以与任何 webdataset 管道一起用作解码器。例如 wds.decode(chug.wds.DecodeDoc('pill'), 'pill')

`hfds`

Hugging Face datasets 支持。最小包装器允许 datasets 与 chug 处理管道一起使用。

使用 datasets 时，处理管道仍然基于 webdataset，通过自定义整理类调用。

`image`

图像处理，基于 torchvision 和 albumentations 的转换构建代码。包括通用图像（imagenet、simclr）转换和特定文档转换的混合，包括基于 albumentations 的 nougat 转换实现。

`text`

文本处理、标记化代码。

`doc`

文档处理代码。目前专注于应用图像/PDF 解码器并处理文档 OCR 或 VQA 注释的处理器。

`task_pipeline`

特定任务的管道，数据集格式与建模需求在此处相遇。

任务管道的输入是基于数据集形式的样本字典，它们被解码然后处理成符合模型输入要求的输出。

处理数据 <--> 模型输入接口的特定任务管道被插入到一个包围性的数据管道中，该管道处理分片列表、随机化、包装、分布式工作、拆分、批处理等。

`chug.loader`

这个独立的顶级文件包含从配置数据类创建加载器及其关联管道的主要工厂方法。

`app`

使用 chug 的大多数应用程序将存在于库外的训练库等中。这里将包含一些内置的实用/探索应用程序。

概念

待完善

数据集

与该库配合良好的数据集可以在 Hugging Face Hub 上的 pixparse 组织（https://huggingface.co/pixparse）下找到。

我们会在了解到其他值得注意的可用数据集时添加相关链接。

使用 / 示例

文档阅读，使用 IDL 训练

import chug
img_cfg = chug.ImageInputCfg(size=(1024, 768), transform_type='doc_better')
img_fn = chug.create_image_preprocessor(input_cfg=img_cfg, is_training=True)
txt_fn = chug.create_text_preprocessor(
    'naver-clova-ix/donut-base',
    prompt_end_token='<s_idl>',
    task_start_token='<s_idl>',  # 注意：需要添加到分词器
)

task_cfg = chug.DataTaskDocReadCfg(
    image_process_fn=img_fn,
    text_process_fn=txt_fn,
    page_sampling='random',
    error_handler='dump_and_reraise',
)
data_cfg = chug.DataCfg(
    source='pipe:curl -s -f -L https://huggingface.co/datasets/pixparse/idl-wds/resolve/main/idl-train-0{0000..2999}.tar',
    batch_size=8,
    num_samples=3144726,
    format='wds',
)
lb = chug.create_loader(
    data_cfg,
    task_cfg,
    is_training=True,
)
ii = iter(lb)
sample = next(ii)

文档阅读，探索 IDL

import chug
task_cfg = chug.DataTaskDocReadCfg(page_sampling='all')
data_cfg = chug.DataCfg(
    source='pixparse/idl-wds',
    split='train',
    batch_size=None,
    format='hfids',
    num_workers=0,    
)
lb = chug.create_loader(
    data_cfg,
    task_cfg,
)
ii = iter(lb)
sample = next(ii)

文档阅读，使用 PDFA 训练

import chug
img_cfg = chug.ImageInputCfg(size=(1024, 768), transform_type='doc_nougat')
img_fn = chug.create_image_preprocessor(input_cfg=img_cfg, is_training=True)
txt_fn = chug.create_text_preprocessor(
    'naver-clova-ix/donut-base',
    prompt_end_token='<s_pdfa>',
    task_start_token='<s_pdfa>',  # 注意：需要添加到分词器
)
task_cfg = chug.DataTaskDocReadCfg(
    image_process_fn=img_fn,
    text_process_fn=txt_fn,
    page_sampling='random',
)
data_cfg = chug.DataCfg(
    source='pipe:curl -s -f -L https://huggingface.co/datasets/pixparse/pdfa-english-train/resolve/main/pdfa-eng-train-{000000..005000}.tar',
    batch_size=8,
    num_samples=1000000,  # 待修复：替换为实际值
    format='wds',   
)
lb = chug.create_loader(
    data_cfg,
    task_cfg,
    is_training=True,
)
ii = iter(lb)
sample = next(ii)

### 文档阅读，探索PDFA

import chug

task_cfg = chug.DataTaskDocReadCfg(
    page_sampling='all',
)
data_cfg = chug.DataCfg(
    source='pixparse/pdfa-eng-wds',
    split='train',
    batch_size=None,
    format='hfids',
    num_workers=0,
)
lb = chug.create_loader(
    data_cfg,
    task_cfg,
)
ii = iter(lb)
sample = next(ii)

### 图像 + 文本

### 训练

import chug
import transformers
from functools import partial
img_cfg = chug.ImageInputCfg(size=(512, 512), transform_type='image_timm')
img_fn = chug.create_image_preprocessor(input_cfg=img_cfg, is_training=True)
tokenizer = transformers.AutoTokenizer.from_pretrained('laion/CLIP-ViT-H-14-laion2B-s32B-b79K')
txt_fn = partial(chug.tokenize, max_length=1000, tokenizer=tokenizer)
task_cfg = chug.DataTaskImageTextCfg(
    image_process_fn=img_fn,
    text_process_fn=txt_fn,
)
data_cfg = chug.DataCfg(
    source='pipe:curl -s -f -L https://huggingface.co/datasets/pixparse/cc12m-wds/resolve/main/cc12m-train-{0000..2175}.tar',
    batch_size=8,
    num_samples=10968539,
    format='wds',   
)
lb = chug.create_loader(
    data_cfg,
    task_cfg,
    is_training=True,
)
ii = iter(lb)
sample = next(ii)

### 文档VQA

#### 训练、微调
import chug
from chug.task_pipeline import create_task_pipeline
img_cfg = chug.ImageInputCfg(size=(1024, 768), transform_type='doc_basic')
img_fn = chug.create_image_preprocessor(img_cfg, is_training=True)
txt_fn = chug.create_text_preprocessor(
    'naver-clova-ix/donut-base-finetuned-docvqa',
    prompt_end_token='<s_answer>',
    task_start_token='<s_docvqa>',
)

task_cfg = chug.DataTaskDocVqaCfg(
    image_process_fn=img_fn,
    text_process_fn=txt_fn,
)
data_cfg = chug.DataCfg(
    source='pipe:curl -s -f -L https://huggingface.co/datasets/pixparse/docvqa-wds/resolve/main/docvqa-train-{000..383}.tar',
    batch_size=8,
    format='wds',
    num_samples=39463,
)
lb = chug.create_loader(
    data_cfg,
    task_cfg,
    is_training=True,
)
ii = iter(lb)
sample = next(ii)

#### 探索

import chug
from chug.task_pipeline import create_task_pipeline
task_cfg = chug.DataTaskDocVqaCfg(
    question_prefix='问题：',
    question_suffix='',
    answer_prefix='答案：',
    answer_suffix=''
)
data_cfg = chug.DataCfg(
    source='pixparse/docvqa-single-page-questions',
    split='validation',
    batch_size=None,
    format='hfids',
    num_workers=0,
)
lb = chug.create_loader(
    data_cfg,
    task_cfg
)
ii = iter(lb)
sample = next(ii)

## 致谢

`chug`源自于在[OpenCLIP](https://github.com/mlfoundations/open_clip)项目中成功使用的`webdataset`数据管道。感谢该项目中的所有贡献者。未来的工作可能会涉及闭环并在OpenCLIP中利用`chug`以增强功能。

`chug`中的图像/文档增强依赖于多个外部影响。我们面向文档的`doc_better` torchvision增强受到`nougat`的影响，而`doc_nougat`则是对[`nougat`](https://github.com/facebookresearch/nougat)中[`albumentations`](https://albumentations.ai/) + `cv2`文档管道的直接改编。几个图像增强利用了`timm`库中的现有工作。

同时，非常感谢[`webdataset`](https://github.com/webdataset/webdataset)和Hugging Face [`datasets`](https://github.com/huggingface/datasets)的维护者们。