chugr 0.2.0.dev0__py3-none-any.whl

chugr/naflex.py

@@ -0,0 +1,357 @@
+""" NaFlex (variable resolution patch sequence) batching for document pipelines.
+
+Adapted from the open_clip NaFlex streaming data pipeline:
+
+* NaFlexBatchScheduler: a deterministic (seq_len, batch_size) schedule, identical across ranks (seeded), sampled
+  per batch from a set of patch sequence length buckets, w/ batch size from a token budget (patches + an optional
+  per-row text token cost). Shuffled per interval, padded and stride-split across dataloader workers.
+* NaFlexBatcher: webdataset stage that pulls samples, runs the task processor for the batch's seq_len bucket
+  (decode / render + image transform + text processing), skips & replaces samples that fail to process so every
+  batch is full (equal step shapes across ranks), and collates patches padded to the bucket seq_len.
+* NaFlexMapDataset: the same schedule driven batching over a map-style dataset (e.g. non-streaming HF datasets),
+  adapted from the open_clip NaFlexMapDatasetWrapper, w/ rank sharded shuffled indices in place of a sample stream.
+* LengthBucketer: reorders samples so similar (text) lengths batch together, reduces decoder padding.
+
+NOTE schedule driven batching of a stream needs an effectively infinite sample stream so every rank / worker can
+always fill its schedule, the wds loader uses resampled shards or (w/o replacement) a repeating, per-pass reshuffled
+shard list.
+"""
+import json
+import logging
+import math
+import os
+import random
+from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple
+
+import torch
+
+from chugr.common import collate, collate_naflex, shuffled_interval_indices
+
+_logger = logging.getLogger(__name__)
+
+_MAX_CONSECUTIVE_FAILURES = int(os.environ.get('CHUGR_MAX_CONSECUTIVE_FAILURES', '1000'))
+
+
+def calculate_batch_size(tokens_per_batch: int, row_tokens: int, divisor: int = 1, max_size: Optional[int] = None):
+    batch_size = tokens_per_batch // max(row_tokens, 1)
+    if divisor > 1:
+        batch_size = (batch_size // divisor) * divisor
+    batch_size = max(batch_size, 1)
+    if max_size is not None:
+        batch_size = min(batch_size, max_size)
+    return batch_size
+
+
+def collate_naflex_batch(samples: List[Dict[str, Any]], image_key: str, seq_len: int) -> Dict[str, Any]:
+    """ Collate processed samples, NaFlex image patches padded to the batch's bucket `seq_len`. """
+    images = [s.pop(image_key) for s in samples]
+    out = collate(samples)
+    out[image_key] = collate_naflex(images, pad_to=seq_len)
+    return out
+
+
+def _interval_value(interval) -> int:
+    if hasattr(interval, 'get_value'):
+        return int(interval.get_value())
+    return int(interval or 0)
+
+
+class NaFlexBatchScheduler:
+    """ Deterministic, rank identical NaFlex (seq_len, batch_size) schedule for one interval (epoch). """
+
+    def __init__(
+            self,
+            seq_lens: Sequence[int],
+            max_tokens_per_batch: int,
+            num_samples: Optional[int] = None,
+            num_tokens: Optional[int] = None,
+            seq_len_probs: Optional[Sequence[float]] = None,
+            per_row_text_tokens: int = 0,
+            batch_divisor: int = 8,
+            seed: int = 42,
+            shuffle: bool = True,
+            world_size: int = 1,
+    ):
+        if (num_samples is None) == (num_tokens is None):
+            raise ValueError('Specify exactly one of `num_samples` or `num_tokens` for NaFlex batching.')
+        if seq_len_probs is not None:
+            if len(seq_len_probs) != len(seq_lens):
+                raise ValueError('`seq_len_probs` must match `seq_lens` length.')
+            pairs = sorted({int(s): float(p) for s, p in zip(seq_lens, seq_len_probs)}.items())
+            self.seq_lens = [s for s, _ in pairs]
+            total = sum(p for _, p in pairs)
+            assert total > 0
+            self.seq_len_probs = [p / total for _, p in pairs]
+        else:
+            self.seq_lens = sorted(set(int(s) for s in seq_lens))
+            self.seq_len_probs = None
+        assert self.seq_lens and all(s > 0 for s in self.seq_lens)
+        self.max_tokens_per_batch = int(max_tokens_per_batch)
+        self.per_row_text_tokens = int(per_row_text_tokens or 0)
+        self.batch_divisor = int(batch_divisor)
+        self.seed = int(seed)
+        self.shuffle = shuffle
+        self.world_size = max(1, int(world_size))
+
+        generator = torch.Generator()
+        generator.manual_seed(self.seed)
+        schedule = []
+        if num_samples is not None:
+            remaining = math.ceil(int(num_samples) / self.world_size)
+            while remaining > 0:
+                seq_len = self._next_seq_len(generator)
+                bs = calculate_batch_size(
+                    self.max_tokens_per_batch, seq_len + self.per_row_text_tokens, self.batch_divisor, remaining)
+                schedule.append((seq_len, bs))
+                remaining -= bs
+        else:
+            remaining = math.ceil(int(num_tokens) / self.world_size)
+            while remaining > 0:
+                seq_len = self._next_seq_len(generator)
+                bs = calculate_batch_size(
+                    min(self.max_tokens_per_batch, remaining), seq_len + self.per_row_text_tokens, self.batch_divisor)
+                schedule.append((seq_len, bs))
+                remaining -= bs * seq_len
+        self._schedule = schedule
+
+    def _next_seq_len(self, generator: torch.Generator) -> int:
+        if self.seq_len_probs is None:
+            idx = int(torch.randint(0, len(self.seq_lens), (1,), generator=generator).item())
+        else:
+            idx = int(torch.multinomial(torch.tensor(self.seq_len_probs), 1, generator=generator).item())
+        return self.seq_lens[idx]
+
+    @staticmethod
+    def pad_schedule_for_workers(schedule: List[Tuple[int, int]], num_workers: int) -> List[Tuple[int, int]]:
+        if num_workers <= 1 or not schedule:
+            return schedule
+        target = math.ceil(len(schedule) / num_workers) * num_workers
+        pad = target - len(schedule)
+        if pad > 0:
+            schedule = schedule + (schedule * math.ceil(pad / len(schedule)))[:pad]
+        return schedule
+
+    def interval_schedule(self, interval: int, num_workers: int = 1) -> List[Tuple[int, int]]:
+        schedule = self.pad_schedule_for_workers(list(self._schedule), max(1, num_workers))
+        if self.shuffle:
+            random.Random(self.seed + interval).shuffle(schedule)
+        return schedule
+
+    def worker_schedule(self, interval: int) -> List[Tuple[int, int]]:
+        info = torch.utils.data.get_worker_info()
+        num_workers = info.num_workers if info is not None else 1
+        worker_id = info.id if info is not None else 0
+        return self.interval_schedule(interval, num_workers)[worker_id::num_workers]
+
+    def num_batches_for_workers(self, num_workers: int) -> int:
+        """ Batches per rank per interval. """
+        return len(self.pad_schedule_for_workers(list(self._schedule), max(1, num_workers)))
+
+    def num_samples_for_workers(self, num_workers: int) -> int:
+        """ Samples per interval (all ranks). """
+        sched = self.pad_schedule_for_workers(list(self._schedule), max(1, num_workers))
+        return sum(bs for _, bs in sched) * self.world_size
+
+    def mean_batch_size_for_workers(self, num_workers: int) -> float:
+        sched = self.pad_schedule_for_workers(list(self._schedule), max(1, num_workers))
+        return sum(bs for _, bs in sched) / len(sched)
+
+
+class NaFlexBatcher:
+    """ webdataset stage, schedule driven NaFlex batching w/ per seq_len bucket task processors. """
+
+    def __init__(
+            self,
+            scheduler: NaFlexBatchScheduler,
+            processors: Dict[int, Callable[[Dict[str, Any]], Any]],
+            image_key: str = 'image_input',
+            interval=0,
+            handler: Optional[Callable[[Exception], bool]] = None,
+            drop_last: bool = True,
+    ):
+        missing = set(scheduler.seq_lens) - set(processors.keys())
+        assert not missing, f'Missing processors for seq_lens {missing}'
+        self.scheduler = scheduler
+        self.processors = processors
+        self.image_key = image_key
+        self.interval = interval
+        self.handler = handler
+        self.drop_last = drop_last
+
+    def _collate(self, samples: List[Dict[str, Any]], seq_len: int) -> Dict[str, Any]:
+        return collate_naflex_batch(samples, self.image_key, seq_len)
+
+    def __call__(self, src: Iterable[Dict[str, Any]]):
+        interval = _interval_value(self.interval)
+        samples_iter = iter(src)
+        for seq_len, batch_size in self.scheduler.worker_schedule(interval):
+            processor = self.processors[seq_len]
+            batch = []
+            failures = 0
+            while len(batch) < batch_size:
+                try:
+                    sample = next(samples_iter)
+                except StopIteration:
+                    break
+                try:
+                    processed = processor(sample)
+                except Exception as ex:
+                    # skip & replace, a falsy handler return (e.g. reraise) stops
+                    if self.handler is not None and not self.handler(ex):
+                        raise
+                    failures += 1
+                    if failures > _MAX_CONSECUTIVE_FAILURES:
+                        raise
+                    continue
+                failures = 0
+                batch.append(processed)
+            if len(batch) == batch_size or (batch and not self.drop_last):
+                yield self._collate(batch, seq_len)
+
+
+class NaFlexMapDataset(torch.utils.data.IterableDataset):
+    """ Schedule driven NaFlex batching over a map-style dataset (e.g. a non-streaming HF dataset), yields batches.
+
+    Adapted from the open_clip NaFlexMapDatasetWrapper. Every rank / worker derives the same shuffled index list
+    and (seq_len, batch_size) schedule for the interval, a rank takes indices [rank::world_size], a worker every
+    num_workers-th batch of the schedule. Use w/ a DataLoader w/ batch_size=None.
+
+    Differences from the open_clip wrapper:
+      * Indices continue through the dataset permutation across intervals (see `shuffled_interval_indices`), so
+        intervals shorter than the dataset visit every sample before repeating (vs the head of a new permutation).
+      * Samples that fail to load / process are replaced w/ a random other sample (subject to `handler`, as
+        NaFlexBatcher), so batches are always full and step shapes equal across ranks.
+    """
+
+    def __init__(
+            self,
+            dataset,
+            scheduler: NaFlexBatchScheduler,
+            processors: Dict[int, Callable[[Dict[str, Any]], Any]],
+            image_key: str = 'image_input',
+            interval=0,
+            rank: int = 0,
+            world_size: int = 1,
+            seed: int = 42,
+            handler: Optional[Callable[[Exception], bool]] = None,
+    ):
+        if not hasattr(dataset, '__len__') or not hasattr(dataset, '__getitem__'):
+            raise TypeError('NaFlex map batching requires a map-style dataset.')
+        missing = set(scheduler.seq_lens) - set(processors.keys())
+        assert not missing, f'Missing processors for seq_lens {missing}'
+        assert 0 <= rank < world_size
+        assert world_size == scheduler.world_size, 'scheduler world_size must match (per rank schedule)'
+        self.dataset = dataset
+        self.scheduler = scheduler
+        self.processors = processors
+        self.image_key = image_key
+        self.interval = interval
+        self.rank = int(rank)
+        self.world_size = int(world_size)
+        self.seed = int(seed)
+        self.handler = handler
+
+    def __iter__(self):
+        interval = _interval_value(self.interval)
+        info = torch.utils.data.get_worker_info()
+        num_workers = info.num_workers if info is not None else 1
+        worker_id = info.id if info is not None else 0
+        dataset_len = len(self.dataset)
+
+        # same (padded, shuffled) schedule and index list on every rank / worker
+        schedule = self.scheduler.interval_schedule(interval, num_workers)
+        samples_per_rank = sum(batch_size for _, batch_size in schedule)
+        indices = shuffled_interval_indices(
+            dataset_len, samples_per_rank * self.world_size, interval, self.seed)[self.rank::self.world_size]
+        # replacement indices for failed samples, distinct per interval / rank / worker
+        rng = random.Random(self.seed + 7919 * interval + 104729 * self.rank + 1299709 * worker_id)
+
+        offset = 0
+        for batch_idx, (seq_len, batch_size) in enumerate(schedule):
+            batch_indices = indices[offset:offset + batch_size]
+            offset += batch_size
+            if batch_idx % num_workers != worker_id:
+                continue
+            processor = self.processors[seq_len]
+            batch = []
+            failures = 0
+            pos = 0
+            while len(batch) < batch_size:
+                if pos < len(batch_indices):
+                    idx = batch_indices[pos]
+                    pos += 1
+                else:
+                    idx = rng.randrange(dataset_len)
+                try:
+                    processed = processor(self.dataset[idx])
+                except Exception as ex:
+                    # skip & replace, a falsy handler return (e.g. reraise) stops
+                    if self.handler is not None and not self.handler(ex):
+                        raise
+                    failures += 1
+                    if failures > _MAX_CONSECUTIVE_FAILURES:
+                        raise
+                    continue
+                failures = 0
+                batch.append(processed)
+            yield collate_naflex_batch(batch, self.image_key, seq_len)
+
+
+def annotation_text_length(sample: Dict[str, Any], keys: Sequence[str] = ('json',)) -> int:
+    """ Approximate text length (characters) of a raw document sample's annotation (for length bucketing). """
+    for k in keys:
+        if k not in sample:
+            continue
+        anno = sample[k]
+        if isinstance(anno, bytes):
+            anno = json.loads(anno)
+            sample[k] = anno  # keep decoded
+        pages = anno.get('pages', anno) if isinstance(anno, dict) else anno
+        total = 0
+        try:
+            for page in pages:
+                page = page.get('lines', page)
+                total += sum(len(t) for t in page.get('text', []))
+        except (AttributeError, TypeError):
+            return 0
+        n = max(len(pages), 1)
+        return total // n  # mean per page (pages are sampled)
+    return 0
+
+
+class LengthBucketer:
+    """ webdataset stage, reorder samples in pools sorted by length (chunks shuffled), reorder only. """
+
+    def __init__(
+            self,
+            length_fn: Callable[[Dict[str, Any]], int] = annotation_text_length,
+            pool: int = 1024,
+            chunk: int = 64,
+            seed: int = 42,
+            interval=0,
+    ):
+        self.length_fn = length_fn
+        self.pool = max(1, int(pool))
+        self.chunk = max(1, int(chunk))
+        self.seed = int(seed)
+        self.interval = interval
+
+    def _flush(self, buffer, rng):
+        buffer.sort(key=self.length_fn)
+        chunks = [buffer[i:i + self.chunk] for i in range(0, len(buffer), self.chunk)]
+        rng.shuffle(chunks)
+        for c in chunks:
+            yield from c
+
+    def __call__(self, src):
+        info = torch.utils.data.get_worker_info()
+        worker_id = info.id if info is not None else 0
+        rng = random.Random(self.seed + _interval_value(self.interval) * 131 + worker_id)
+        buffer = []
+        for sample in src:
+            buffer.append(sample)
+            if len(buffer) >= self.pool:
+                yield from self._flush(buffer, rng)
+                buffer = []
+        if buffer:
+            yield from self._flush(buffer, rng)

chugr/task_pipeline/__init__.py

@@ -0,0 +1,7 @@
+from .pipeline_doc_read import build_task_pipeline_doc_read, DataTaskDocReadCfg
+from .pipeline_doc_vqa import build_task_pipeline_doc_vqa, DataTaskDocVqaCfg
+from .pipeline_gtparse import build_task_pipeline_gtparse
+from .pipeline_image_text import build_task_pipeline_image_text, DataTaskImageTextCfg
+from .pipeline_manual import build_task_pipeline_manual, DataTaskManualCfg
+
+from .pipeline_factory import create_task_pipeline, create_task_processor

chugr/task_pipeline/pipeline_doc_read.py

@@ -0,0 +1,66 @@
+from dataclasses import dataclass, field
+from typing import Callable, Optional
+
+import webdataset as wds
+
+from chugr.common import DataTaskCfg, FeatureInfo, ImageFeatureInfo
+from chugr.doc import DocReadProcessor, DEFAULT_DOC_FEAT
+from chugr.doc.synth_text import SynthTextCfg
+from chugr.wds import get_error_handler
+
+
+@dataclass
+class DataTaskDocReadCfg(DataTaskCfg):
+    image_input_feat: ImageFeatureInfo = DEFAULT_DOC_FEAT
+    text_input_feat: FeatureInfo = FeatureInfo('text_input', input_key='pages')
+    text_target_feat: Optional[FeatureInfo] = FeatureInfo('text_target', input_key=None)
+    page_sampling: str = 'random'
+    render_dpi: int = 150
+    synth_prob: float = 0.  # probability of a synthetic text rendering (of the sample's text) replacing the page
+    synth_cfg: Optional[SynthTextCfg] = None
+
+
+def create_processor_doc_read(cfg: DataTaskDocReadCfg, image_process_fn: Optional[Callable] = None):
+    return DocReadProcessor(
+        image_process_fn=image_process_fn if image_process_fn is not None else cfg.image_process_fn,
+        text_process_fn=cfg.text_process_fn,
+        image_input_feat=cfg.image_input_feat,
+        text_input_feat=cfg.text_input_feat,
+        text_target_feat=cfg.text_target_feat,
+        page_sampling=cfg.page_sampling,
+        render_dpi=cfg.render_dpi,
+        flatten_json=cfg.flatten_json,
+        synth_prob=cfg.synth_prob,
+        synth_cfg=cfg.synth_cfg,
+    )
+
+
+def build_task_pipeline_doc_read(
+        cfg: DataTaskDocReadCfg,
+):
+    handler = get_error_handler(cfg.error_handler)
+    pipe = []
+
+    # document decoding & pre-processing done together, there is coupling in random page
+    # selection and in the future, masking of image and/or text
+    pipe += [wds.map(create_processor_doc_read(cfg), handler=handler)]
+
+    if cfg.output_tuple:
+        # NOTE in this mode we lose '_parse' key and would need to derive from target
+        # Unless we add support for parse as the last tuple element?
+        if cfg.text_target_feat is not None:
+            pipe += [
+                wds.to_tuple(
+                    cfg.image_input_feat.output_name,
+                    cfg.text_input_feat.output_name,
+                    cfg.text_target_feat.output_name,
+                )
+            ]
+        else:
+            pipe += [
+                wds.to_tuple(
+                    cfg.image_input_feat.output_name,
+                    cfg.text_input_feat.output_name,
+                )
+            ]
+    return pipe

chugr/task_pipeline/pipeline_doc_vqa.py

@@ -0,0 +1,124 @@
+from dataclasses import dataclass, field
+from functools import partial
+from typing import Callable, Dict, Optional, Union
+
+import webdataset as wds
+
+from chugr.common import DataTaskCfg, FeatureInfo, ImageFeatureInfo
+from chugr.wds import get_error_handler
+from chugr.doc import (
+    DocVqaProcessor,
+    DEFAULT_DOC_KEY,
+    DEFAULT_QUESTION_KEY,
+    DEFAULT_QUESTION_ID_KEY,
+    DEFAULT_ANSWER_KEY,
+    DEFAULT_DOC_KEY_TUPLE,
+    DEFAULT_QUESTION_KEY_TUPLE,
+    DEFAULT_ANSWER_KEY_TUPLE,
+    DEFAULT_DOC_FEAT,
+    DEFAULT_QUESTION_FEAT,
+    DEFAULT_QUESTION_ID_FEAT,
+    DEFAULT_ANSWER_FEAT
+)
+
+
+def filter_missing(
+        sample,
+        image_key=DEFAULT_DOC_KEY_TUPLE,
+        question_key=DEFAULT_QUESTION_KEY_TUPLE,
+        answer_key=DEFAULT_ANSWER_KEY_TUPLE,
+):
+    has_question = any(k in sample for k in question_key)
+    has_answer = any(k in sample for k in answer_key)
+    has_image = any(k in sample for k in image_key)
+    return has_question and has_answer and has_image
+
+
+# Currently assuming this schema as default, one set of question/answers per image, images possibly duplicated
+# sample = {
+#     'png': bytes,
+#     'question_id': 33,
+#     'doc_id': 55,  # optional
+#     'question': 'what is a trumpet?',
+#     'answers': ['an instrument', 'a brass instrument']
+# }
+#
+
+
+@dataclass
+class DataTaskDocVqaCfg(DataTaskCfg):
+    """
+    Attributes:
+        answer_feat:
+        question_feat:
+        question_id_feat:
+        image_input_feat:
+        text_input_feat:
+        text_target_feat:
+        question_prefix:
+        question_suffix:
+        answer_prefix:
+        answer_suffix:
+        render_dpi:
+    """
+    answer_feat: FeatureInfo = DEFAULT_ANSWER_FEAT
+    question_feat: FeatureInfo = DEFAULT_QUESTION_FEAT
+    question_id_feat: FeatureInfo = DEFAULT_QUESTION_ID_FEAT
+    image_input_feat: ImageFeatureInfo = DEFAULT_DOC_FEAT
+    text_input_feat: FeatureInfo = FeatureInfo('text_input', input_key=None)
+    text_target_feat: Optional[FeatureInfo] = FeatureInfo('text_target', input_key=None)
+    question_prefix: Optional[str] = '<s_question>'
+    question_suffix: Optional[str] = '</s_question>'
+    answer_prefix: Optional[str] = '<s_answer>'
+    answer_suffix: Optional[str] = '</s_answer>'
+    # FIXME prompt templates instead of prefix+suffix above?
+    render_dpi: int = 144
+
+
+def create_processor_doc_vqa(cfg: DataTaskDocVqaCfg, image_process_fn: Optional[Callable] = None):
+    return DocVqaProcessor(
+        image_process_fn=image_process_fn if image_process_fn is not None else cfg.image_process_fn,
+        text_process_fn=cfg.text_process_fn,
+        image_input_feat=cfg.image_input_feat,
+        text_input_feat=cfg.text_input_feat,
+        text_target_feat=cfg.text_target_feat,
+        question_feat=cfg.question_feat,
+        answer_feat=cfg.answer_feat,
+        question_id_feat=cfg.question_id_feat,
+        render_dpi=cfg.render_dpi,
+        question_prefix=cfg.question_prefix,
+        question_suffix=cfg.question_suffix,
+        answer_prefix=cfg.answer_prefix,
+        answer_suffix=cfg.answer_suffix,
+    )
+
+
+def build_task_pipeline_doc_vqa(
+        cfg: DataTaskDocVqaCfg,
+):
+    # document decoding & pre-processing done together, there is coupling in random page
+    # selection and in the future, masking of image and/or text
+    handler = get_error_handler(cfg.error_handler)
+
+    pipe = [wds.map(create_processor_doc_vqa(cfg), handler=handler)]
+
+    if cfg.output_tuple:
+        # NOTE in this mode we lose '_parse' key and would need to derive from target
+        # Unless we add support for parse as the last tuple element?
+        if cfg.text_target_feat is not None:
+            pipe += [
+                wds.to_tuple(
+                    cfg.image_input_feat.output_name,
+                    cfg.text_input_feat.output_name,
+                    cfg.text_target_feat.output_name,
+                )
+            ]
+        else:
+            pipe += [
+                wds.to_tuple(
+                    cfg.image_input_feat.output_name,
+                    cfg.text_input_feat.output_name,
+                )
+            ]
+    return pipe
+

chugr/task_pipeline/pipeline_factory.py

@@ -0,0 +1,32 @@
+from chugr.common import DataTaskCfg
+
+from .pipeline_doc_read import build_task_pipeline_doc_read, create_processor_doc_read, DataTaskDocReadCfg
+from .pipeline_doc_vqa import build_task_pipeline_doc_vqa, create_processor_doc_vqa, DataTaskDocVqaCfg
+from .pipeline_image_text import build_task_pipeline_image_text, DataTaskImageTextCfg
+# from .donut_gtparse_pipe import
+from .pipeline_manual import build_task_pipeline_manual, DataTaskManualCfg
+
+_cfg_to_create = {
+    DataTaskDocReadCfg: build_task_pipeline_doc_read,
+    DataTaskDocVqaCfg: build_task_pipeline_doc_vqa,
+    DataTaskImageTextCfg: build_task_pipeline_image_text,
+    DataTaskManualCfg: build_task_pipeline_manual,
+}
+
+
+def create_task_pipeline(cfg: DataTaskCfg):
+    create_fn = _cfg_to_create[type(cfg)]
+    return create_fn(cfg)
+
+
+_cfg_to_processor = {
+    DataTaskDocReadCfg: create_processor_doc_read,
+    DataTaskDocVqaCfg: create_processor_doc_vqa,
+}
+
+
+def create_task_processor(cfg: DataTaskCfg, image_process_fn=None):
+    """ Per-sample processor (decode + image / text processing) for tasks that support NaFlex bucketed batching. """
+    if type(cfg) not in _cfg_to_processor:
+        raise ValueError(f'NaFlex batching is not supported for task cfg {type(cfg).__name__}.')
+    return _cfg_to_processor[type(cfg)](cfg, image_process_fn=image_process_fn)

chugr/task_pipeline/pipeline_gtparse.py

@@ -0,0 +1,93 @@
+from dataclasses import dataclass, field
+from functools import partial
+from typing import Callable, Optional, Union
+
+import webdataset as wds
+
+from chugr.common import DataTaskCfg, ImageFeatureInfo, FeatureInfo
+from chugr.wds import get_error_handler, create_image_decoder
+
+_DEFAULT_IMG_KEY = "jpg;png;jpeg;webp;tif"
+_DEFAULT_TXT_KEY = "ground_truth"
+_DEFAULT_IMG_KEY_TUPLE = tuple(_DEFAULT_IMG_KEY.split(';'))
+_DEFAULT_TXT_KEY_TUPLE = tuple(_DEFAULT_TXT_KEY.split(';'))
+
+_DEFAULT_IMAGE_FEAT = ImageFeatureInfo('image_input', input_key=_DEFAULT_IMG_KEY, image_mode='L')
+_DEFAULT_TXT_FEAT = FeatureInfo('ground_truth', input_key=_DEFAULT_TXT_KEY)
+
+
+@dataclass
+class DataTaskImageTextCfg(DataTaskCfg):
+    image_input_feat: ImageFeatureInfo = _DEFAULT_IMAGE_FEAT
+    text_input_feat: FeatureInfo = _DEFAULT_TXT_FEAT
+
+
+def filter_no_caption_or_no_image(
+        sample,
+        image_key=_DEFAULT_IMG_KEY_TUPLE,
+        text_key=_DEFAULT_TXT_KEY_TUPLE
+):
+    has_caption = any(k in sample for k in text_key)
+    has_image = any(k in sample for k in image_key)
+    return has_caption and has_image
+
+
+def build_task_pipeline_gtparse(
+        cfg: DataTaskImageTextCfg,
+):
+    """ Create pipeline for dual image & text input pipelines.
+    FIXME add support for caption target for caption tasks or separate pipe?
+    """
+    handler = get_error_handler(cfg.error_handler)
+    pipe = []
+
+    if cfg.filter_valid:
+        filter_fn = partial(
+            filter_no_caption_or_no_image,
+            image_key=tuple(cfg.image_input_feat.input_key.split(';')),
+            text_key=tuple(cfg.text_input_feat.input_key.split(';')),
+        )
+        pipe += [
+            wds.select(filter_fn)
+        ]
+
+    if cfg.decode_and_process_fn:
+        pipe += [
+            wds.map(cfg.decode_and_process_fn)
+        ]
+    else:
+        decode_fn = create_image_decoder(
+            cfg.decode_fn,
+            image_mode=cfg.image_input_feat.image_mode,
+            handler=handler,
+        )
+
+        rename_dict = {
+            cfg.image_input_feat.output_name: cfg.text_input_feat.input_key,
+            cfg.text_input_feat.output_name: cfg.text_input_feat.input_key,
+        }
+        pipe += [
+            decode_fn,
+            wds.rename(**rename_dict),
+        ]
+
+        map_dict = {}
+        if cfg.image_process_fn is not None:
+            map_dict[cfg.image_input_feat.output_name] = cfg.image_process_fn
+        if cfg.text_process_fn is not None:
+            map_dict[cfg.text_input_feat.output_name] = cfg.text_process_fn
+
+        if map_dict:
+            pipe += [
+                wds.map_dict(**map_dict, handler=handler)
+            ]
+
+    if cfg.output_tuple:
+        pipe += [
+            wds.to_tuple(
+                cfg.image_input_feat.output_name,
+                cfg.text_input_feat.output_name,
+            )
+        ]
+
+    return pipe