chugr 0.2.0.dev0__py3-none-any.whl

chugr/doc/doc_read_processor.py

@@ -0,0 +1,249 @@
+import json
+from typing import Optional, Callable
+
+from chugr import ImageFeatureInfo, FeatureInfo
+from chugr.doc import DocProcessor, DEFAULT_DOC_FEAT
+from chugr.doc.doc_processor import get_next_valid_page_index, _get_value, _logger
+from chugr.doc.synth_text import SynthTextCfg, SynthTextRenderer
+
+
+class DocReadProcessor(DocProcessor):
+    """ Process documents w/ OCR annotation for reading tasks.
+    """
+
+    def __init__(
+            self,
+            image_process_fn: Optional[Callable] = None,
+            text_process_fn: Optional[Callable] = None,
+            image_input_feat: ImageFeatureInfo = DEFAULT_DOC_FEAT,
+            text_input_feat: FeatureInfo = FeatureInfo('text_input', input_key='pages'),
+            text_target_feat: FeatureInfo = FeatureInfo('text_target', input_key=None),
+            line_break: str = '\n',
+            page_sampling: str = 'random',
+            render_dpi: int = 150,
+            squeeze_pages: bool = True,
+            expand_pages: bool = False,
+            flatten_json: bool = True,
+            seed: int = 0,
+            synth_prob: float = 0.,
+            synth_cfg: Optional[SynthTextCfg] = None,
+    ):
+        """
+        Args:
+            synth_prob: probability of replacing a sample's page image w/ a synthetic rendering of (part of) its
+                annotation text (target == rendered text, no layout annotations used, page image not decoded)
+            synth_cfg: synthetic text rendering config
+        """
+        super().__init__(
+            image_process_fn=image_process_fn,
+            text_process_fn=text_process_fn,
+            image_input_feat=image_input_feat,
+            text_input_feat=text_input_feat,
+            text_target_feat=text_target_feat,
+            render_dpi=render_dpi,
+            page_sampling=page_sampling,
+            squeeze_pages=squeeze_pages,
+            expand_pages=expand_pages,
+            flatten_json=flatten_json,
+            seed=seed,
+        )
+        self.line_break = line_break
+        assert page_sampling in ('random', 'first', 'all_valid', 'all')
+        self.synth_prob = synth_prob
+        self.synth = SynthTextRenderer(synth_cfg) if synth_prob > 0 else None
+
+    def __call__(self, sample):
+        if self.synth is not None and self.generator.random() < self.synth_prob:
+            return self._synth_sample(sample)
+        return super().__call__(sample)
+
+    def _process_text(self, text):
+        """ Tokenize text (if text_process_fn set), returns (text input, text target or None). """
+        if self.text_process_fn is None:
+            return text, None
+        processed = self.text_process_fn(text)
+        assert self.text_input_name in processed, \
+            f"Text input name '{self.text_input_name}' not found in processed sample."
+        target = processed.get(self.text_target_name) if self.text_target_name else None
+        if self.text_target_feat is not None and target is None:
+            assert False, f"Expected a text target named '{self.text_target_name}' in processed sample."
+        return processed[self.text_input_name], target
+
+    def _synth_sample(self, sample):
+        if 'json' in sample:
+            if isinstance(sample['json'], bytes):
+                sample['json'] = json.loads(sample['json'])
+            if self.flatten_json:
+                sample.update(sample.pop('json'))
+        anno = _get_value(self.text_input_key, sample)
+        assert anno, f"No annotation found with keys ({self.text_input_key})."
+        pages = [p['lines'] if 'lines' in p else p for p in anno]
+        valid = [i for i, p in enumerate(pages) if p.get('text')]
+        if not valid:
+            raise RuntimeError("No text on any page, skipping sample...")
+        page_index = self.generator.choice(valid)
+        image, lines = self.synth(pages[page_index]['text'], rng=self.generator)
+        if self.image_input_feat.image_mode and image.mode != self.image_input_feat.image_mode:
+            image = image.convert(self.image_input_feat.image_mode)
+        image = self._preprocess_image_pages([image])[0]
+        text_input, text_target = self._process_text(self.line_break.join(lines))
+
+        output = {
+            self.image_input_name: image,
+            self.text_input_name: text_input,
+            '_parse': {
+                'num_pages': len(anno),
+                'page_indices': [page_index],
+                'page_text': [lines],
+                'synthetic': True,
+            },
+        }
+        if text_target is not None:
+            output[self.text_target_name] = text_target
+        if not self.squeeze_pages:
+            for k in (self.image_input_name, self.text_input_name, self.text_target_name):
+                if k in output:
+                    output[k] = [output[k]]
+        return output
+
+    def _process_anno_pages(self, anno):
+        assert isinstance(anno, (list, tuple)), f"Annotation should be a list of pages"
+        num_pages = len(anno)
+        if not num_pages:
+            raise RuntimeError("Empty annotation. Skipping...")
+
+        # FIXME for initial behaviour we will randomly sample one of N pages
+        # TODO determine if we want to train in multi-page mode, use another sampling strategy?
+        page_indices = []
+        try:
+            if self.page_sampling == 'random':
+                n_wanted_pages = min(1, num_pages)  # TODO increase for multi-page processing, rand start+end?
+                current_index = self.generator.randrange(-1, num_pages - 1)
+                for _ in range(n_wanted_pages):
+                    current_index = get_next_valid_page_index(current_index, num_pages, anno)
+                    page_indices.append(current_index)
+            elif self.page_sampling == 'first':
+                current_index = get_next_valid_page_index(-1, num_pages, anno)
+                page_indices.append(current_index)
+            elif self.page_sampling == 'all_valid':
+                current_index = -1
+                for _ in range(num_pages):
+                    current_index = get_next_valid_page_index(current_index, num_pages, anno)
+                    page_indices.append(current_index)
+            elif self.page_sampling == 'all':
+                page_indices = list(range(num_pages))
+        except RuntimeError:
+            pass
+
+        if not page_indices:
+            raise RuntimeError("No valid annotated pages. Skipping...")
+
+        text_pages = []
+        tokenized_text_pages = []
+        target_pages = []
+        for current_index in page_indices:
+            # FIXME currently encoding each page separately with own start/end tokens.
+            #  For multi-age should consider encoding in one sequence w/ page-break tokens.
+            anno_page = anno[current_index]
+            if 'lines' in anno_page:
+                # Two supported formats right now
+                # {
+                #     'pages': [
+                #         {
+                #             'text': [],  # these are lines
+                #             'bbox': [],
+                #         }
+                #     ]
+                # }
+                #
+                # OR
+                #
+                # {
+                #     'pages': [
+                #         {
+                #             'lines': {
+                #                 'text': [],
+                #                 'bbox': [],
+                #             },
+                #             'words': {
+                #                 'text': [],
+                #                 'bbox': [],
+                #             }
+                #         }
+                #     ]
+                # }
+                #
+                #
+                anno_page = anno_page['lines']
+
+            # Currently page text is created by concatenating lines of text with a CR line break
+            # Additions could involve:
+            # * using different line-break tokens between lines
+            # * using word-level bbox anno information to mask works and construct partial lines
+            # * group lines into blocks (or use block annos) and treat blocks / paragraphs of text and
+            if not anno_page["text"]:
+                raise RuntimeError("No text on page, skipping sample...")
+
+            text = self.line_break.join(anno_page["text"])
+
+            # tokenize w/ and generate training target if enabled
+            text_input, text_target = self._process_text(text)
+            tokenized_text_pages.append(text_input)
+            if text_target is not None:
+                target_pages.append(text_target)
+
+            text_pages.append(anno_page["text"])  # unencoded text added as lines
+
+        gt_parse = {
+            'num_pages': num_pages,  # total # of pages in doc
+            'page_indices': page_indices,  # page indices sampled
+            'page_text': text_pages,  # text of sampled page indices pages[].lines[]
+        }
+
+        output = {
+            self.text_input_name: tokenized_text_pages,
+            '_parse': gt_parse,
+        }
+        if target_pages:
+            output[self.text_target_name] = target_pages
+
+        return output
+
+    def _decode_anno(self, sample):
+        anno = _get_value(self.text_input_key, sample)
+        assert anno is not None, f"No annotation found with keys ({self.text_input_key})."
+
+        try:
+            page_anno = self._process_anno_pages(anno)
+        except Exception as exn:
+            _logger.error(f'Issue processing annotation for {sample["__url__"]}, {sample["__key__"]}.')
+            #_logger.error(json.dumps(anno, indent=4))
+            raise exn
+
+        # extract info from the _parse
+        info = page_anno.get('_parse', {})
+        page_indices = info.get('page_indices', [0])  # the samples page indices
+        num_anno_pages = info.get('num_pages', 1)
+
+        # TODO support 'image info' to relay details such as text bbox, layout
+        # page_image_info = info.get('image_info', None)
+        # if page_image_info is not None:
+        #     assert len(page_image_info) == len(page_indices)
+
+        return page_anno, page_indices, num_anno_pages
+
+    def _expand_anno(self, anno, count: int):
+        expanded_annos = []
+        for i in range(count):
+            sample = {}
+            for k, v in anno.items():
+                if k == '_parse':
+                    gt_parse = {}
+                    gt_parse['num_pages'] = v['num_pages']
+                    gt_parse['page_indices'] = [v['page_indices'][i]]
+                    gt_parse['page_text'] = [v['page_text'][i]]
+                    sample[k] = gt_parse
+                else:
+                    sample[k] = v[i] if isinstance(v, (list, tuple)) else v
+            expanded_annos.append(sample)
+        return expanded_annos

chugr/doc/doc_vqa_processor.py

@@ -0,0 +1,113 @@
+import random
+from typing import Callable, Dict, List, Optional, Tuple
+
+from chugr import FeatureInfo, ImageFeatureInfo
+from chugr.doc import DocProcessor, DEFAULT_QUESTION_FEAT, DEFAULT_QUESTION_ID_FEAT, DEFAULT_ANSWER_FEAT, \
+    DEFAULT_DOC_FEAT
+from chugr.doc.doc_processor import _get_value
+
+
+class DocVqaProcessor(DocProcessor):
+    def __init__(
+            self,
+            image_process_fn: Optional[Callable] = None,
+            text_process_fn: Optional[Callable] = None,
+            question_feat: FeatureInfo = DEFAULT_QUESTION_FEAT,
+            question_id_feat: FeatureInfo = DEFAULT_QUESTION_ID_FEAT,
+            answer_feat: FeatureInfo = DEFAULT_ANSWER_FEAT,
+            multi_qa_feat: Optional[FeatureInfo] = None,
+            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>',
+            render_dpi: int = 150,
+            squeeze_pages: bool = True,
+            expand_pages: bool = False,
+            flatten_json: bool = True,
+            seed: int = 0,
+    ):
+        super().__init__(
+            image_process_fn=image_process_fn,
+            text_process_fn=text_process_fn,
+            image_input_feat=image_input_feat,
+            text_input_feat=text_input_feat,
+            text_target_feat=text_target_feat,
+            render_dpi=render_dpi,
+            squeeze_pages=squeeze_pages,
+            expand_pages=expand_pages,
+            flatten_json=flatten_json,
+            seed=seed,
+        )
+        self.question_feat = question_feat
+        self.question_key = question_feat.input_key.split(';')
+        self.question_id_feat = question_id_feat
+        self.question_id_key = question_id_feat.input_key.split(';')
+        self.answer_feat = answer_feat
+        self.answer_key = answer_feat.input_key.split(';')
+        if multi_qa_feat is not None:
+            self.expand_pages = True  # override
+            self.multi_qa_key = multi_qa_feat.input_key.split(';')
+        else:
+            # expand pages only used / supported for multi-qa expansion right now
+            self.expand_pages = False
+            self.multi_qa_key = None
+
+        # FIXME support flexible q/a prompting formats, do with prefix/suffix or template strings?
+        # Donut style: '<s_docvqa><s_question>{question}</s_question><s_answer>{answer}</answer><eos>'
+        # Common: 'Question: {question} Answer: {answer}<eos>'
+        self.question_prefix = question_prefix or ''
+        self.question_suffix = question_suffix or ''
+        self.answer_prefix = answer_prefix or ''
+        self.answer_suffix = answer_suffix or ''
+        #self.prompt_template = '<s_question>{question}</s_question><s_answer>'
+        #self.prompt_template_full = '<s_question>{question}</s_question><s_answer>{answer}</s_answer>'
+
+    def _decode_anno(self, sample) -> Tuple[Dict, List[int], int]:
+        if self.multi_qa_key:
+            # FIXME multi qa expansion is a WIP
+            qa_list = sample[self.multi_qa_key]
+            assert isinstance(qa_list, (list, tuple)), f'Expected a list or tuple, got {type(qa_list)}.'
+            assert False, 'WIP'
+        else:
+            question = _get_value(self.question_key, sample)
+            question_id = _get_value(self.question_id_key, sample)
+            answers = _get_value(self.answer_key, sample)
+
+            if isinstance(answers, str):
+                answers = [answers]
+            if answers and self.text_target_feat is not None:
+                answer = self.generator.choice(answers)
+            else:
+                answer = None
+
+            input_text = self.question_prefix + question + self.question_suffix + self.answer_prefix
+            if answer is not None:
+                input_text += answer + self.answer_suffix
+
+            output = {
+                self.text_input_name: input_text,
+                '_parse': {
+                    'question_id': question_id,
+                    'question': question,
+                    'answers': answers,  # list, all answers included in parse
+                }
+            }
+
+            if self.text_process_fn is not None:
+                processed = self.text_process_fn(input_text)
+                assert self.text_input_name in processed, \
+                    f"Text input name '{self.text_input_name}' not found in processed sample."
+                if self.text_target_feat is not None:
+                    assert self.text_target_name in processed, f"Expected a text target named '{self.text_target_name}' in processed sample."
+                output.update(processed)
+            else:
+                output[self.text_input_name] = input_text
+
+            return output, [0], 1
+
+    def _expand_anno(self, anno, count: int):
+        # FIXME implement expansion for multi-qa (and eventually multi-page option)
+        pass

chugr/doc/synth_text.py

@@ -0,0 +1,212 @@
+""" Synthetic text rendering for reading warmup.
+
+Renders (real) document text lines w/ varied fonts, sizes, and layouts. The training target is exactly the rendered
+text, so rendered pages need no layout annotations (bboxes), only the text.
+"""
+import glob
+import io
+import logging
+import os
+import random
+import re
+from dataclasses import dataclass, field
+from functools import lru_cache
+from typing import List, Optional, Sequence, Tuple
+
+from PIL import Image, ImageDraw, ImageFilter, ImageFont
+
+_logger = logging.getLogger(__name__)
+
+_FONT_DIRS = ('/usr/share/fonts', '/usr/local/share/fonts', '~/.fonts', '~/.local/share/fonts')
+
+# Latin text fonts (excludes symbol, dingbat, script, math, CJK fallback fonts)
+_DEFAULT_FONT_RE = (
+    r'^(DejaVu(Sans|Serif)(Condensed|Mono)?(-(Bold|Oblique|Italic|BoldOblique|BoldItalic))?'
+    r'|Liberation(Sans|Serif|Mono|SansNarrow)-.*'
+    r'|Free(Sans|Serif|Mono).*'
+    r'|(C059|P052|NimbusSans|NimbusSansNarrow|NimbusRoman|NimbusMonoPS|URWBookman|URWGothic)-.*'
+    r')\.(ttf|otf)$'
+)
+
+# printable latin-1 (IDL OCR text is ~all latin-1), other chars are dropped from rendered text + target
+_CHARSET = set(chr(c) for c in range(0x20, 0x7f)) | set(chr(c) for c in range(0xa1, 0x100))
+
+
+@dataclass
+class SynthTextCfg:
+    font_paths: Optional[List[str]] = None  # explicit font files, default: discover system fonts (font_pattern)
+    font_pattern: str = _DEFAULT_FONT_RE  # regex on font file basename for discovery
+    lines_range: Tuple[int, int] = (1, 48)  # rendered lines per sample (after wrapping), log-uniform
+    max_chars: int = 3000  # cap on rendered characters (keep the target within the text length budget)
+    font_size_range: Tuple[int, int] = (16, 44)  # px
+    width_range: Tuple[int, int] = (500, 1600)  # px, text block width
+    line_spacing_range: Tuple[float, float] = (1.0, 1.6)  # x font size
+    margin_range: Tuple[float, float] = (0.02, 0.12)  # fraction of text block width
+    page_prob: float = 0.3  # place the text block on a letter aspect page (vs a tight crop around the text)
+    paragraph_gap_prob: float = 0.1  # per line chance of an extra blank gap before it
+    center_prob: float = 0.08  # per line chance of centered (vs left aligned) line
+    blur_prob: float = 0.2
+    noise_prob: float = 0.2
+    jpeg_prob: float = 0.2
+
+
+def discover_fonts(pattern: str = _DEFAULT_FONT_RE, dirs: Sequence[str] = _FONT_DIRS) -> List[str]:
+    regex = re.compile(pattern)
+    found = {}
+    for d in dirs:
+        for path in glob.glob(os.path.join(os.path.expanduser(d), '**', '*.*'), recursive=True):
+            name = os.path.basename(path)
+            if regex.match(name) and name not in found:
+                found[name] = path
+    return sorted(found.values())
+
+
+def _check_font(path: str) -> bool:
+    """ Font loads and (if fontTools is available) covers the charset. """
+    try:
+        ImageFont.truetype(path, 20)
+    except Exception:
+        return False
+    try:
+        from fontTools.ttLib import TTFont
+    except ImportError:
+        return True
+    try:
+        cmap = TTFont(path, lazy=True).getBestCmap() or {}
+    except Exception:
+        return False
+    return all(ord(c) in cmap for c in _CHARSET if c != ' ')
+
+
+@lru_cache(maxsize=512)
+def _load_font(path: str, size: int):
+    return ImageFont.truetype(path, size)
+
+
+def _wrap_line(text: str, font, width: int) -> List[str]:
+    """ Greedy word wrap to `width` px, words longer than the width are split by character. """
+    if font.getlength(text) <= width:
+        return [text]
+    out = []
+    current = ''
+    for word in text.split(' '):
+        candidate = word if not current else current + ' ' + word
+        if font.getlength(candidate) <= width:
+            current = candidate
+            continue
+        if current:
+            out.append(current)
+        while font.getlength(word) > width and len(word) > 1:
+            n = len(word)
+            while n > 1 and font.getlength(word[:n]) > width:
+                n -= 1
+            out.append(word[:n])
+            word = word[n:]
+        current = word
+    if current:
+        out.append(current)
+    return out
+
+
+def clean_text_lines(lines: Sequence[str]) -> List[str]:
+    out = []
+    for line in lines:
+        line = ''.join(c for c in line.replace('\t', ' ') if c in _CHARSET)
+        line = ' '.join(line.split())
+        if line:
+            out.append(line)
+    return out
+
+
+class SynthTextRenderer:
+    """ Render text lines to a grayscale document image, returns (image, rendered lines). """
+
+    def __init__(self, cfg: Optional[SynthTextCfg] = None):
+        self.cfg = cfg = cfg or SynthTextCfg()
+        paths = cfg.font_paths or discover_fonts(cfg.font_pattern)
+        self.fonts = [p for p in paths if _check_font(p)]
+        assert self.fonts, f'No usable fonts found (pattern {cfg.font_pattern}).'
+        _logger.info(f'SynthTextRenderer using {len(self.fonts)} fonts.')
+
+    def __call__(self, lines: Sequence[str], rng: Optional[random.Random] = None) -> Tuple[Image.Image, List[str]]:
+        cfg = self.cfg
+        rng = rng or random
+        lines = clean_text_lines(lines)
+        if not lines:
+            raise RuntimeError('No renderable text.')
+
+        font_path = rng.choice(self.fonts)
+        font_size = rng.randint(*cfg.font_size_range)
+        font = _load_font(font_path, font_size)
+        block_w = rng.randint(*cfg.width_range)
+        lo, hi = cfg.lines_range
+        num_lines = int(round(lo * (hi / lo) ** rng.random())) if hi > lo else lo
+
+        # contiguous span of source lines, wrapped to the block width
+        start = rng.randrange(len(lines))
+        rendered = []
+        line_center = []
+        line_gap = []
+        num_chars = 0
+        for src in lines[start:]:
+            center = rng.random() < cfg.center_prob
+            gap = rng.random() < cfg.paragraph_gap_prob
+            for i, wrapped in enumerate(_wrap_line(src, font, block_w)):
+                if len(rendered) >= num_lines or (rendered and num_chars + len(wrapped) > cfg.max_chars):
+                    break
+                rendered.append(wrapped)
+                line_center.append(center)
+                line_gap.append(gap and i == 0 and len(rendered) > 1)
+                num_chars += len(wrapped) + 1
+            else:
+                continue
+            break
+
+        line_h = int(font_size * rng.uniform(*cfg.line_spacing_range))
+        text_w = int(max(font.getlength(l) for l in rendered))
+        text_h = line_h * (len(rendered) + sum(line_gap))
+        if not any(line_center):
+            block_w = max(text_w, min(block_w, text_w + font_size * 4))  # tight-ish block for left aligned text
+        margin_x = int(block_w * rng.uniform(*cfg.margin_range)) + 2
+        margin_y = int(block_w * rng.uniform(*cfg.margin_range)) + 2
+        canvas_w = block_w + 2 * margin_x
+        canvas_h = text_h + 2 * margin_y
+        offset_x, offset_y = margin_x, margin_y
+        if rng.random() < cfg.page_prob:
+            # letter-ish page, text block placed at a random position
+            page_w = max(canvas_w, int(canvas_w * rng.uniform(1.0, 1.4)))
+            page_h = max(canvas_h, int(page_w * rng.uniform(1.25, 1.45)))
+            offset_x += rng.randint(0, page_w - canvas_w)
+            offset_y += rng.randint(0, page_h - canvas_h)
+            canvas_w, canvas_h = page_w, page_h
+
+        bg = rng.randint(215, 255)
+        ink = rng.randint(0, min(80, bg - 120))
+        img = Image.new('L', (canvas_w, canvas_h), color=bg)
+        draw = ImageDraw.Draw(img)
+        y = offset_y
+        for text, center, gap in zip(rendered, line_center, line_gap):
+            if gap:
+                y += line_h
+            x = offset_x
+            if center:
+                x += max(0, (block_w - int(font.getlength(text))) // 2)
+            draw.text((x, y), text, font=font, fill=ink)
+            y += line_h
+
+        if rng.random() < cfg.blur_prob:
+            img = img.filter(ImageFilter.GaussianBlur(radius=rng.uniform(0.3, 1.0)))
+        if rng.random() < cfg.noise_prob:
+            img = _add_noise(img, rng.uniform(3, 12), rng)
+        if rng.random() < cfg.jpeg_prob:
+            buf = io.BytesIO()
+            img.save(buf, format='JPEG', quality=rng.randint(30, 90))
+            img = Image.open(io.BytesIO(buf.getvalue())).convert('L')
+        return img, rendered
+
+
+def _add_noise(img: Image.Image, sigma: float, rng: random.Random) -> Image.Image:
+    import numpy as np
+    arr = np.asarray(img, dtype=np.float32)
+    arr = arr + np.random.default_rng(rng.randrange(2 ** 31)).normal(0, sigma, arr.shape).astype(np.float32)
+    return Image.fromarray(arr.clip(0, 255).astype(np.uint8), mode='L')

chugr/hfds/__init__.py

@@ -0,0 +1,2 @@
+from .loader import create_loader_hf
+from .wrappers import SafeDataset

chugr/hfds/collate.py

@@ -0,0 +1,83 @@
+from typing import List, Callable, Optional
+
+import torch.utils
+from torch.utils.data import IterableDataset, DataLoader
+
+from chugr.common import collate
+
+
+def invoke(f, *args, **kwargs):
+    if isinstance(f, (IterableDataset, DataLoader)) and len(args) == 0:
+        return iter(f)
+    if isinstance(f, list):
+        return iter(f)
+    if callable(f):
+        result = f(*args, **kwargs)
+        return result
+    raise ValueError(f"{f}: not a valid pipeline stage")
+
+
+"""
+pipe = [wds.rename(image='jpg'), wds.map_dict(image=tf), wds.to_tuple('image', 'cls')]
+hfc = HfCollate(pipe)
+dl = DataLoader(ds, batch_size=32, num_workers=4, persistent_workers=True, collate_fn=hfc)
+"""
+
+def flatten_bytes(data):
+    for sample in data:
+        to_replace = {k for k, v in sample.items() if isinstance(v, dict) and 'bytes' in v}
+        if to_replace:
+            result = {k: v for k, v in sample.items() if k not in to_replace}
+            result.update({k: sample[k]['bytes'] for k in to_replace})
+            yield result
+        else:
+            yield sample
+
+
+class HfCollate:
+    """ Collation wrapper that applies processing pipeline for HF datasets use
+    """
+    def __init__(
+        self,
+            pipeline: List[Callable],
+            collate_fn: Optional[Callable] = None,
+            apply_collate: bool = True,
+    ):
+        """
+        Args:
+            pipeline: list of pipeline functions
+            collate_fn: use a custom collation function, otherwise defaults to torch default_collate
+        """
+        self.pipeline = pipeline
+        self.collate_fn = collate_fn or collate
+        self.apply_collate = apply_collate
+        self._debug = False
+
+    def __call__(self, batch):
+        item = False
+        if not self.apply_collate and isinstance(batch, dict):
+            batch = [batch]
+            item = True
+
+        if self._debug:
+            for b in batch:
+                for k, v in b.items():
+                    print(k, type(v))
+                    if isinstance(v, torch.Tensor):
+                        print(v.shape)
+
+        if self.pipeline:
+            for pipe_fn in self.pipeline:
+                batch = invoke(pipe_fn, batch)
+
+        batch = list(batch)
+
+        if self._debug:
+            for b in batch:
+                for k, v in b.items():
+                    print(k, v)
+
+        if self.apply_collate:
+            return self.collate_fn(batch)
+        else:
+            return batch[0] if item else batch