chugr 0.2.0.dev0__py3-none-any.whl

chugr/__init__.py

@@ -0,0 +1,46 @@
+from .common import (
+    ImageInputCfg,
+    ImageAugCfg,
+    LoaderBundle,
+    ImageFeatureInfo,
+    FeatureInfo,
+    ShardSpec,
+    SourceSpec,
+    DataArg,
+    DataCfg,
+    DistributedCfg,
+)
+from .hfds import create_loader_hf
+from .image import (
+    build_image_transforms,
+    build_transforms_image_basic,
+    build_transforms_image_timm,
+    build_transforms_doc_basic,
+    build_transforms_doc_better,
+    build_transforms_doc_nougat,
+    create_image_preprocessor,
+)
+from .loader import create_loader, create_loader_from_config_hf, create_loader_from_config_wds
+from .task_pipeline import (
+    create_task_pipeline,
+    build_task_pipeline_doc_read,
+    build_task_pipeline_doc_vqa,
+    build_task_pipeline_gtparse,
+    build_task_pipeline_image_text,
+    build_task_pipeline_manual,
+    DataTaskDocReadCfg,
+    DataTaskDocVqaCfg,
+    DataTaskImageTextCfg,
+    DataTaskManualCfg,
+)
+from .text import tokenize, text_input_to_target, prepare_text_input, create_text_preprocessor
+from .version import __version__
+from .wds import (
+    create_loader_wds,
+    build_data_pipeline,
+    decode_image_pages,
+    decode_pdf_pages,
+    create_image_decoder,
+    DecodeDoc,
+)
+

chugr/app/test.py

@@ -0,0 +1,105 @@
+""" Iterate a document loader for throughput benchmarking & sample inspection.
+
+Example:
+    python -m chugr.app.test --data.source "idl-train-0000{0..3}.tar" --data.batch-size 8 --data.num-samples 1000 \\
+        --input.transform-type doc_naflex --input.patch-size 16 --input.max-seq-len 4096 --task doc_read
+"""
+import logging
+import os
+import sys
+import time
+from dataclasses import dataclass, field
+from typing import Optional
+
+import torch
+
+import chugr
+from chugr.common import DataArg, DataCfg, ImageInputCfg, Serializable, parse_args
+
+_logger = logging.getLogger('chugr.test')
+
+
+@dataclass
+class TestArgs(Serializable):
+    data: DataArg = field(default_factory=lambda: DataArg(source=''))
+    input: ImageInputCfg = field(default_factory=lambda: ImageInputCfg(transform_type='doc_better', size=(1024, 768)))
+    task: str = 'doc_read'  # one of 'doc_read', 'doc_vqa'
+    tokenizer: Optional[str] = None  # HF tokenizer name, text left as strings if not set
+    text_max_length: int = 1024
+    is_training: bool = True
+    num_batches: int = 50
+    dump_dir: Optional[str] = None  # write first batch images (un-normalized) here
+
+
+def _dump_image(x, input_cfg: ImageInputCfg, path: str):
+    from torchvision.utils import save_image
+    if isinstance(x, dict):
+        # reassemble naflex patches into an image
+        valid = x['patch_valid']
+        patches = x['patches'][valid]
+        coords = x['patch_coord'][valid]
+        p = input_cfg.patch_size
+        c = patches.shape[-1] // (p * p)
+        h, w = (coords.max(0).values + 1).tolist()
+        img = torch.zeros(c, h * p, w * p)
+        for patch, (py, px) in zip(patches, coords.tolist()):
+            img[:, py * p:(py + 1) * p, px * p:(px + 1) * p] = patch.reshape(p, p, c).permute(2, 0, 1)
+        x = img
+    mean = torch.tensor(input_cfg.mean).view(-1, 1, 1)
+    std = torch.tensor(input_cfg.std).view(-1, 1, 1)
+    save_image((x * std + mean).clamp(0, 1), path)
+
+
+def main(argv=None):
+    logging.basicConfig(level=logging.INFO)
+    args: TestArgs = parse_args(TestArgs, argv=argv, description='chugr loader test')
+    print(args.dumps_yaml())
+
+    img_fn = chugr.create_image_preprocessor(args.input, is_training=args.is_training)
+    txt_fn = None
+    if args.tokenizer:
+        txt_fn = chugr.create_text_preprocessor(args.tokenizer, max_length=args.text_max_length)
+    if args.task == 'doc_read':
+        task_cfg = chugr.DataTaskDocReadCfg(
+            image_process_fn=img_fn,
+            text_process_fn=txt_fn,
+            text_target_feat=chugr.FeatureInfo('text_target', input_key=None) if txt_fn else None,
+            error_handler='dump_and_reraise',
+        )
+    elif args.task == 'doc_vqa':
+        task_cfg = chugr.DataTaskDocVqaCfg(
+            image_process_fn=img_fn,
+            text_process_fn=txt_fn,
+            text_target_feat=chugr.FeatureInfo('text_target', input_key=None) if txt_fn else None,
+        )
+    else:
+        assert False, f'Unknown task {args.task}'
+
+    loader = chugr.create_loader(DataCfg.from_arg(args.data), task_cfg, is_training=args.is_training)
+    start = time.time()
+    num_samples = 0
+    for i, batch in enumerate(loader):
+        image = batch['image_input']
+        bs = (image['patches'] if isinstance(image, dict) else image).shape[0]
+        num_samples += bs
+        if i == 0:
+            if isinstance(image, dict):
+                _logger.info(f"patches {tuple(image['patches'].shape)}, valid {image['patch_valid'].sum(1).tolist()}")
+            else:
+                _logger.info(f'image {tuple(image.shape)}')
+            if args.dump_dir:
+                os.makedirs(args.dump_dir, exist_ok=True)
+                for j in range(bs):
+                    x = {k: v[j] for k, v in image.items()} if isinstance(image, dict) else image[j]
+                    _dump_image(x, args.input, os.path.join(args.dump_dir, f'sample_{j}.png'))
+                    text = batch['text_input'][j]
+                    with open(os.path.join(args.dump_dir, f'sample_{j}.txt'), 'w') as f:
+                        f.write(text if isinstance(text, str) else str(text))
+        if i + 1 >= args.num_batches:
+            break
+    elapsed = time.time() - start
+    _logger.info(f'{num_samples} samples in {elapsed:.1f}s, {num_samples / elapsed:.1f} samples/s')
+
+
+if __name__ == '__main__':
+    sys.exit(main())

chugr/common/__init__.py

@@ -0,0 +1,9 @@
+from .cli import parse_args, str2bool
+from .collate import collate, collate_naflex, is_naflex_sample
+from .config import ImageInputCfg, ImageAugCfg, PreprocessCfg, image_mode_to_chs
+from .config import DataArg, DataCfg, DistributedCfg, source_to_shard_spec
+from .serialize import Serializable, to_dict, from_dict
+from .random import get_pytorch_worker_seed, rank_seed, seed_worker, shuffled_interval_indices
+from .task_config import DataTaskCfg
+from .types import SourceSpec, SplitInfo, ShardSpec, SharedCount, LoaderBundle, FeatureInfo, ImageFeatureInfo
+from .urls import expand_urls

chugr/common/cli.py

@@ -0,0 +1,234 @@
+""" Nested dataclass <-> command line arguments
+
+A small replacement for the subset of `simple_parsing` used by chugr & pixparse apps.
+
+* nested dataclass fields map to dotted args, e.g. `--task.opt.learning-rate 1e-4`
+  (dashes and underscores are both accepted)
+* `subgroups` select the dataclass type of a field by name, e.g. `--task cruller_pretrain`
+* a `--config` yaml/json file provides defaults that command line args override
+"""
+import argparse
+import dataclasses
+import sys
+import typing
+from typing import Any, Dict, List, Optional, Sequence, Type, TypeVar
+
+from .serialize import _is_union, _unwrap_optional, from_dict, resolve_field_types, _load_file
+
+T = TypeVar('T')
+
+_NONE_STRINGS = ('none', 'null')
+
+
+def str2bool(v: str) -> bool:
+    if isinstance(v, bool):
+        return v
+    lv = v.lower()
+    if lv in ('true', 't', 'yes', 'y', '1'):
+        return True
+    if lv in ('false', 'f', 'no', 'n', '0'):
+        return False
+    raise argparse.ArgumentTypeError(f'Boolean value expected, got {v}.')
+
+
+def _scalar_parser(tp):
+    def _parse(v: str):
+        if v.lower() in _NONE_STRINGS:
+            return None
+        if tp is bool:
+            return str2bool(v)
+        if tp is int:
+            return int(v)
+        if tp is float:
+            return float(v)
+        if _is_union(tp):
+            # e.g. Union[int, str] -> try members in order
+            for a in typing.get_args(tp):
+                if a is type(None):
+                    continue
+                try:
+                    return _scalar_parser(a)(v)
+                except (ValueError, argparse.ArgumentTypeError):
+                    continue
+            return v
+        return v
+    return _parse
+
+
+def _arg_spec(tp):
+    """ Return (nargs, item parser) for a field type. """
+    tp = _unwrap_optional(tp)
+    origin = typing.get_origin(tp)
+    args = typing.get_args(tp)
+    if origin is tuple:
+        if len(args) == 2 and args[1] is Ellipsis:
+            return '+', _scalar_parser(args[0])
+        item_tp = args[0] if args else str
+        return '+', _scalar_parser(item_tp)  # length validated when building the dataclass
+    if origin is list:
+        return '+', _scalar_parser(args[0] if args else str)
+    if _is_union(tp):
+        # union containing a sequence type, accept one or more values
+        for a in args:
+            if typing.get_origin(a) in (list, tuple):
+                return '+', _scalar_parser(typing.get_args(a)[0] if typing.get_args(a) else str)
+    return None, _scalar_parser(tp)
+
+
+def _flatten_fields(cls, prefix: str = '') -> Dict[str, Any]:
+    """ dotted name -> field type for all leaf fields of a (nested) dataclass type """
+    out = {}
+    field_types = resolve_field_types(cls)
+    for f in dataclasses.fields(cls):
+        if not f.init:
+            continue
+        tp = field_types[f.name]
+        name = f'{prefix}{f.name}'
+        inner = _unwrap_optional(tp)
+        if dataclasses.is_dataclass(inner):
+            out.update(_flatten_fields(inner, prefix=name + '.'))
+        else:
+            out[name] = tp
+    return out
+
+
+def _set_nested(d: Dict[str, Any], dotted: str, value):
+    keys = dotted.split('.')
+    for k in keys[:-1]:
+        if d.get(k) is None:
+            d[k] = {}
+        d = d[k]
+    d[keys[-1]] = value
+
+
+def _norm(name: str) -> str:
+    return name.replace('-', '_')
+
+
+def parse_args(
+        cls: Type[T],
+        argv: Optional[Sequence[str]] = None,
+        subgroups: Optional[Dict[str, Dict[str, type]]] = None,
+        subgroup_defaults: Optional[Dict[str, str]] = None,
+        description: Optional[str] = None,
+) -> T:
+    """ Parse command line args into dataclass `cls`.
+
+    Args:
+        cls: The top level config dataclass.
+        argv: Arguments, defaults to sys.argv[1:].
+        subgroups: Mapping of dotted field name -> {choice name: dataclass type}. The choice is selected
+            with `--<field> <name>` and the selected type's fields become `--<field>.<subfield>` args.
+        subgroup_defaults: Default choice name per subgroup field.
+        description: Parser description.
+    """
+    argv = list(sys.argv[1:] if argv is None else argv)
+    subgroups = subgroups or {}
+    subgroup_defaults = subgroup_defaults or {}
+
+    # stage 1: config file + subgroup selection
+    pre = argparse.ArgumentParser(add_help=False)
+    pre.add_argument('--config', type=str, default=None)
+    for sg_name, choices in subgroups.items():
+        pre.add_argument(
+            f'--{sg_name.replace("_", "-")}', f'--{sg_name}', dest=sg_name, type=str,
+            default=None, choices=list(choices.keys()))
+    pre_args, _ = pre.parse_known_args(argv)
+
+    file_cfg = _load_file(pre_args.config) if pre_args.config else {}
+
+    selected = {}
+    for sg_name, choices in subgroups.items():
+        choice = getattr(pre_args, sg_name)
+        if choice is None:
+            # allow config file to select via a `_name` key in the subgroup dict
+            sg_file = file_cfg.get(sg_name)
+            if isinstance(sg_file, dict) and '_name' in sg_file:
+                choice = sg_file.pop('_name')
+            else:
+                choice = subgroup_defaults.get(sg_name)
+        assert choice is not None, f'A value for --{sg_name} is required (one of {list(choices.keys())}).'
+        selected[sg_name] = choice
+
+    # stage 2: build full parser over all leaf fields
+    leaf_types = {}
+    field_types = resolve_field_types(cls)
+    for f in dataclasses.fields(cls):
+        if not f.init:
+            continue
+        if f.name in subgroups:
+            sub_cls = subgroups[f.name][selected[f.name]]
+            leaf_types.update(_flatten_fields(sub_cls, prefix=f.name + '.'))
+            continue
+        tp = field_types[f.name]
+        inner = _unwrap_optional(tp)
+        if dataclasses.is_dataclass(inner):
+            leaf_types.update(_flatten_fields(inner, prefix=f.name + '.'))
+        else:
+            leaf_types[f.name] = tp
+
+    parser = argparse.ArgumentParser(description=description)
+    parser.add_argument('--config', type=str, default=None, help='yaml or json config file')
+    for sg_name, choices in subgroups.items():
+        parser.add_argument(
+            f'--{sg_name.replace("_", "-")}', f'--{sg_name}', dest=f'__sg_{sg_name}', type=str,
+            default=None, choices=list(choices.keys()))
+
+    for dotted, tp in leaf_types.items():
+        nargs, item_parser = _arg_spec(tp)
+        flags = {f'--{dotted}', f'--{dotted.replace("_", "-")}'}
+        kwargs = dict(dest=dotted, type=item_parser, default=argparse.SUPPRESS)
+        if nargs is not None:
+            kwargs['nargs'] = nargs
+        parser.add_argument(*sorted(flags), **kwargs)
+
+    args = parser.parse_args(argv)
+
+    # merge: file config < command line
+    merged: Dict[str, Any] = dict(file_cfg)
+    for dotted, value in vars(args).items():
+        if dotted == 'config' or dotted.startswith('__sg_'):
+            continue
+        if isinstance(value, list) and len(value) == 1 and value[0] is None:
+            value = None  # e.g. `--input.size none` for an Optional[Tuple]
+        _set_nested(merged, dotted, value)
+
+    # build nested dataclasses
+    kwargs = {}
+    for f in dataclasses.fields(cls):
+        if not f.init or f.name not in merged:
+            if f.name in subgroups:
+                kwargs[f.name] = subgroups[f.name][selected[f.name]]()
+            continue
+        value = merged[f.name]
+        if f.name in subgroups:
+            sub_cls = subgroups[f.name][selected[f.name]]
+            kwargs[f.name] = _build_with_defaults(sub_cls, value or {})
+            continue
+        tp = _unwrap_optional(field_types[f.name])
+        if dataclasses.is_dataclass(tp) and isinstance(value, dict):
+            kwargs[f.name] = _build_with_defaults(tp, value)
+        else:
+            from .serialize import from_value
+            kwargs[f.name] = from_value(field_types[f.name], value)
+    return cls(**kwargs)
+
+
+def _build_with_defaults(cls, overrides: Dict[str, Any]):
+    """ Build dataclass `cls` from partial (nested) overrides, keeping field defaults for the rest. """
+    from .serialize import from_value
+    field_types = resolve_field_types(cls)
+    kwargs = {}
+    for f in dataclasses.fields(cls):
+        if not f.init or f.name not in overrides:
+            continue
+        value = overrides[f.name]
+        tp = _unwrap_optional(field_types[f.name])
+        if dataclasses.is_dataclass(tp) and isinstance(value, dict):
+            kwargs[f.name] = _build_with_defaults(tp, value)
+        else:
+            kwargs[f.name] = from_value(field_types[f.name], value)
+    unknown = set(overrides.keys()) - {f.name for f in dataclasses.fields(cls)}
+    if unknown:
+        raise TypeError(f'Unknown fields {sorted(unknown)} for {cls.__name__}.')
+    return cls(**kwargs)

chugr/common/collate.py

@@ -0,0 +1,104 @@
+import collections
+from typing import Callable, Dict, List, Optional, Tuple, Type, Union
+
+import torch
+from torch.utils.data._utils.collate import default_collate_fn_map, default_collate_err_msg_format
+
+NAFLEX_KEYS = ('patches', 'patch_coord', 'patch_valid')
+
+
+def is_naflex_sample(elem) -> bool:
+    return isinstance(elem, collections.abc.Mapping) and all(k in elem for k in NAFLEX_KEYS)
+
+
+def collate_naflex(
+        batch: List[Dict[str, torch.Tensor]],
+        pad_multiple: int = 8,
+        pad_to: Optional[int] = None,
+) -> Dict[str, torch.Tensor]:
+    """ Pad variable length patch sequences (NaFlex) to the longest in the batch.
+
+    Args:
+        batch: list of dicts w/ 'patches' [N, D] (or [N, Ph, Pw, C]), 'patch_coord' [N, 2], 'patch_valid' [N]
+        pad_multiple: round padded length up to a multiple of this (kernel friendly shapes)
+        pad_to: pad to this fixed length (e.g. the NaFlex seq_len bucket) instead of the batch max
+    """
+    max_len = max(b['patches'].shape[0] for b in batch)
+    if pad_to is not None:
+        assert max_len <= pad_to, f'NaFlex sample length {max_len} exceeds bucket seq_len {pad_to}'
+        max_len = pad_to
+    elif pad_multiple > 1:
+        max_len = ((max_len + pad_multiple - 1) // pad_multiple) * pad_multiple
+    batch_size = len(batch)
+    patch_shape = batch[0]['patches'].shape[1:]
+    patches = torch.zeros((batch_size, max_len, *patch_shape), dtype=batch[0]['patches'].dtype)
+    patch_coord = torch.zeros((batch_size, max_len, 2), dtype=torch.int64)
+    patch_valid = torch.zeros((batch_size, max_len), dtype=torch.bool)
+    for i, b in enumerate(batch):
+        n = b['patches'].shape[0]
+        patches[i, :n] = b['patches']
+        patch_coord[i, :n] = b['patch_coord']
+        patch_valid[i, :n] = b['patch_valid']
+    return {'patches': patches, 'patch_coord': patch_coord, 'patch_valid': patch_valid}
+
+
+def collate(batch):
+    r"""
+        A customized collate function that handles collection type of element within each batch.
+
+        This collate function has been tweaked to provide different functionality when handling
+        dictionary samples. Certain keys are excluded or not tensorized.
+
+        Args:
+            batch: a single batch to be collated
+    """
+    elem = batch[0]
+    elem_type = type(elem)
+
+    if elem_type in default_collate_fn_map:
+        return default_collate_fn_map[elem_type](batch)
+
+    for collate_type in default_collate_fn_map:
+        if isinstance(elem, collate_type):
+            return default_collate_fn_map[collate_type](batch)
+
+    if is_naflex_sample(elem):
+        return collate_naflex(batch)
+
+    if isinstance(elem, collections.abc.Mapping):
+        try:
+            out = {}
+            for key in elem:
+                if key.startswith('__'):
+                    # skip keys starting with '__', e.g. '__key__',
+                    continue
+                elif key.startswith('_'):
+                    # do not recurse or tensorize values for keys starting with '_', e.g. '_parse'
+                    out[key] = [d[key] for d in batch]
+                else:
+                    out[key] =  collate([d[key] for d in batch])
+            out = elem_type(out)
+            return out
+        except TypeError:
+            # The mapping type may not support `__init__(iterable)`.
+            return {key: collate([d[key] for d in batch]) for key in elem}
+    elif isinstance(elem, tuple) and hasattr(elem, '_fields'):  # namedtuple
+        return elem_type(*(collate(samples) for samples in zip(*batch)))
+    elif isinstance(elem, collections.abc.Sequence):
+        # check to make sure that the elements in batch have consistent size
+        it = iter(batch)
+        elem_size = len(next(it))
+        if not all(len(elem) == elem_size for elem in it):
+            raise RuntimeError('each element in list of batch should be of equal size')
+        transposed = list(zip(*batch))  # It may be accessed twice, so we use a list.
+
+        if isinstance(elem, tuple):
+            return [collate(samples) for samples in transposed]  # Backwards compatibility.
+        else:
+            try:
+                return elem_type([collate(samples) for samples in transposed])
+            except TypeError:
+                # The sequence type may not support `__init__(iterable)` (e.g., `range`).
+                return [collate(samples) for samples in transposed]
+
+    raise TypeError(default_collate_err_msg_format.format(elem_type))