You signed in with another tab or window.Reload to refresh your session.You signed out in another tab or window.Reload to refresh your session.You switched accounts on another tab or window.Reload to refresh your session.Dismiss alert
A library to help w/ efficient training for multi-modal data. Initially focused on image & document + text tasks.
chug currently leverageswebdataset and Huggingfacedatasets.
webdataset tar files and dataset pipelines are preferred for scalable pretraining.
Huggingfacedatasets are supported and work great for exploration, validation, and fine-tune use cases.
Design
Submodule Hierarchy
The library has been designed so that functions, classes at different levels can be used independently.
If one wants to build a loader & pipeline with JSON/YAML serializable configs, use the top-levelchug.create_loader() inchug/loader.py. Depending on dataset sources, one can easily switch this between webdataset, HF datasets (in the future, other sources).
Bypassing the highest level, one can also callbuild_pipeline_* methods intask_pipeline and then callcreate_loader_wds with a full array of args forwds only use cases.
If one doesn't want to usechug loaders and pipelines at all,image,text, andwds (especially decoder) functionality may be useful in other projects.
Library modules (highest to lowest level)
The dependencies of modules within the library are intended to follow the hierarchy below. e.g. doc depends on wds, but wds should never depend on doc.
Configs, structures (dataclasses) for general use across the library
wds
Webdataset (wds for short) specific code. Extensions and alterations of webdataset functionality to fit covered use case and improve robustness.
All data pipelines inchug currently leveragewds pipelines, even when not usingwds datasets.
Document oriented decoding (pdf decoder) is present inchug/wds/decode.py, it can be used with any webdataset pipeline as a decoder. e.g.wds.decode(chug.wds.DecodeDoc('pill'), 'pill')
hfds
Huggingfacedatasets support. A minimal wrapper that allowsdatasets to be used with chug processing pipelines.
The processing pipelines remain webdataset based when usingdatasets, they are invoked by a custom collate class.
image
Image processing,torchvision andalbumentations based transform building code. A mix of generic image (imagenet, simclr) transforms and document specific transforms, including an implementation ofalbumentations basednougat transforms.
text
Text processing, tokenization code.
doc
Document processing code. Currently focused on processors that apply image/pdf decoders and process document OCR or VQA annotations.
task_pipeline
Task specific pipelines, where dataset formats meet modelling needs.
Inputs to task pipelines are sample dictionaries based on the dataset form, they are decoded and then processed into outputs that match model input requirements.
Task specific pipelines are inserted into the more generic data pipeline.
chug.loader
This lone top-level file includes the main factory methods for creating loaders w/ associated pipelines from config dataclasses.
app
Most applications usingchug will exist outside of the lib in training libraries, etc. Some builtin utility / exploration apps will be included here.
Concepts
WIP
TODOs
Nearish
Cleanup and refinement, codebase will change
Documentation & unit-tests
Support reading of info .json/.yaml files for automatic shard info resolution for webdatasets (like timm)
Support unified preprocessor functions for combined image + text tokenization (img+text token interleaving, etc.)
Longish
Increase range of task pipelines for other tasks, modelling needs
Explore alternatives to .tar shards (array_record, arrow, etc)
Usage / Examples
Document Reading, Training w/ IDL
importchugimg_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>',# NOTE needs to be added to tokenizer)task_cfg=chug.DataTaskDocReadCfg(image_process_fn=img_fn,text_process_fn=txt_fn,page_sampling='random',error_handler='dump_and_reraise',)task_pipe=chug.create_task_pipeline(task_cfg)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)
importchugimg_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>',# NOTE needs to be added to tokenizer)task_cfg=chug.DataTaskDocReadCfg(image_process_fn=img_fn,text_process_fn=txt_fn,page_sampling='random',)task_pipe=chug.create_task_pipeline(task_cfg)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,# FIXME replace with actualformat='wds', )lb=chug.create_loader(data_cfg,task_cfg,is_training=True,)ii=iter(lb)sample=next(ii)