pixr 0.1.0.dev0__tar.gz

pixr-0.1.0.dev0/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright 2024 Hugging Face Inc.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

pixr-0.1.0.dev0/PKG-INFO

@@ -0,0 +1,349 @@
+Metadata-Version: 2.4
+Name: pixr
+Version: 0.1.0.dev0
+Summary: Pixel Reading: OCR-free document understanding models (image encoder + text decoder), training and evaluation
+Keywords: document-ai,ocr-free,document-understanding,vision-language,docvqa,pytorch
+Author-Email: Ross Wightman <rwightman@gmail.com>
+License-Expression: Apache-2.0
+License-File: LICENSE
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Education
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Topic :: Scientific/Engineering
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Scientific/Engineering :: Image Recognition
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Project-URL: homepage, https://github.com/rwightman/pixr
+Project-URL: repository, https://github.com/rwightman/pixr
+Requires-Python: >=3.10
+Requires-Dist: torch>=2.8
+Requires-Dist: timm>=1.0.31
+Requires-Dist: transformers
+Requires-Dist: huggingface_hub
+Requires-Dist: safetensors
+Requires-Dist: fsspec
+Requires-Dist: numpy
+Requires-Dist: chugr>=0.2.0.dev0
+Requires-Dist: Levenshtein
+Requires-Dist: pyyaml
+Requires-Dist: jsonargparse
+Requires-Dist: tiktoken
+Provides-Extra: wandb
+Requires-Dist: wandb; extra == "wandb"
+Provides-Extra: tensorboard
+Requires-Dist: tensorboard; extra == "tensorboard"
+Description-Content-Type: text/markdown
+
+# Pixel Reading (`pixr`)
+
+OCR-free, end-to-end document understanding models trained on open data: pixels in, text out. `pixr` covers the
+models (image encoder + text decoder families), the training tasks (page reading pretraining, DocVQA fine-tuning),
+generation (KV cached greedy / beam search) and a lightweight PyTorch training framework (DDP / FSDP2, compile, EMA).
+Data pipelines are provided by [`chugr`](https://github.com/rwightman/chugr).
+
+`pixr` is the continuation of the original `pixparse` project (huggingface/pixparse, 2023), renamed and revived.
+
+## Status
+
+Alpha, expect breaking changes. The current focus is reproducing decent OCR / DocVQA results w/ open data and modern
+(NaFlex, native aspect ratio) image encoders.
+
+### Changes since `pixparse`
+
+* Ported to `chugr` 0.2, NaFlexVit (native aspect / variable resolution) image encoders, and self-contained text
+  decoders (BART, Ettin, Qwen3) w/ three model families (see below)
+* FSDP2, torch.compile, EMA, NaFlex batching, chunked LM loss, KV cached greedy / beam search generation
+* Fixes for issues that hampered the original DocVQA results:
+  * DocVQA fine-tune loss was applied to question + answer tokens (prompt masking used the task start token)
+  * train / eval image preprocessing mismatched (aspect preserving resize + pad vs stretch, long axis rotation)
+  * DocVQA eval decoded by re-tokenizing the generated string each step (BPE round trip is not idempotent)
+  * `scale_embedding` was forced on for BART decoders (copied from Donut's mBART, BART is trained w/o)
+  * `ModelArgs` command line overrides (image size, etc) were silently dropped
+  * pretrained BART cross-attention weights significantly slow learning w/ an image encoder, cross-attention is now
+    re-initialized w/ zero init output projections by default (decoder starts as the pretrained LM)
+
+## Installation
+
+```bash
+pip install --pre pixr   # pulls in chugr
+pip install --pre "pixr[wandb]"   # optional logging extras: wandb, tensorboard
+```
+
+NOTE the NaFlexVit attention backends (`attn_backend`) and pretrained patch embedding channel adaptation used by the
+NaFlex models need `timm` changes that are not in a release yet, until they are, install `timm` from source w/ them.
+
+From a source checkout (the training / eval scripts and the test launcher live in the repository):
+
+```bash
+pip install -e ./chugr -e ./pixr
+```
+
+## Models
+
+Three model families share the same tasks, data pipelines, loss masking and (greedy / beam search, KV cached) generation:
+
+* **Cruller** -- image encoder + autoregressive text decoder w/ cross-attention (Donut style)
+* **Fritter** -- image encoder -> (optional m x m token merge) -> projector -> image token prefix for a decoder-only
+  LM w/ prefix-LM attention (PaLI / LLaVA style)
+* **Cronut** -- native decoder-only transformer over image patches and text (GenLIP style single trunk w/ 2D/1D
+  interleaved MRoPE, qk-norm, gated attention, SwiGLU), trained from scratch w/ a tiktoken tokenizer
+
+Model configs live in `src/pixr/models/configs`:
+
+| config | image encoder | text decoder |
+|---|---|---|
+| `cruller_naflex_clip_base` | NaFlexVit B/16, CLIP style (pre-norm, cls token), OpenAI CLIP init, 4096 patches | BART base (6L) |
+| `cruller_naflex_clip_base_8k` | as above, 8192 patches | BART base (6L) |
+| `cruller_naflex_clip_base_ettin150m` | as above, 4096 patches | Ettin (ModernBERT) decoder 150M + cross-attn |
+| `cruller_naflex_siglip2_base` | NaFlexVit B/16 SigLIP2 | BART base (6L) |
+| `cruller_naflex_clip_large` | NaFlexVit L/14 CLIP style, OpenAI CLIP init, 6144 patches | BART large (8L) |
+| `cruller_naflex_clip_large_ettin400m` | as above | Ettin decoder 400M + cross-attn |
+| `cruller_naflex_siglip2_base_ettin150m` | NaFlexVit B/16 SigLIP2 | Ettin decoder 150M + cross-attn |
+| `cruller_naflex_siglip2_so400m` | NaFlexVit SO400M/16 SigLIP2 | BART large (8L) |
+| `fritter_naflex_siglip2_base_ettin150m` | NaFlexVit B/16 SigLIP2, 2x2 token merge | Ettin decoder 150M (prefix-LM) |
+| `fritter_naflex_siglip2_base_ettin400m` | as above | Ettin decoder 400M (prefix-LM) |
+| `fritter_naflex_siglip2_so400m_ettin1b` | NaFlexVit SO400M/16 SigLIP2, 2x2 token merge | Ettin decoder 1B (prefix-LM) |
+| `cruller_naflex_siglip2_base_qwen3_0.6b` | NaFlexVit B/16 SigLIP2 | Qwen3 0.6B + zero init cross-attn every N layers |
+| `fritter_naflex_siglip2_base_qwen3_0.6b` | NaFlexVit B/16 SigLIP2, 2x2 token merge | Qwen3 0.6B (prefix-LM) |
+| `cronut_s16_tiktoken` / `cronut_b16_tiktoken` | -- (patch embed in trunk) | -- (single trunk), `cl100k_base` |
+| `cruller_vit_clip_base_1280` | fixed size ViT B/16 CLIP @ 1280x960 | BART base (6L) |
+| `cruller_hybrid_r50_base_1280` | fixed size hybrid R50 + ViT B/16 @ 1280x960 | BART base (6L) |
+
+NaFlex models consume pages at their native aspect ratio, resized to fit a patch budget (`--task.model.max-seq-len`),
+the budget can be changed between pretraining and fine-tuning.
+
+Ettin (ModernBERT) decoders use 64 token sliding window attention in 2 of every 3 layers, a poor fit for a long
+image prefix (Fritter), Qwen3 (full attention) is the recommended prefix-LM decoder.
+
+### Attention backends
+
+Padded NaFlex sequences can use a dense mask or a variable length kernel, selected per model config field:
+
+| field | options |
+|---|---|
+| `image_encoder.attn_backend` | `sdpa` (dense mask), `varlen` (packed, CUDA), `flex` (FlexAttention block mask), batches w/o padding skip masking for all -- requires the timm `naflex-attn-backends` branch |
+| `text_decoder.cross_attn_backend` | `sdpa`, `varlen`, `flex` |
+| `text_decoder.self_attn_backend` | `sdpa` (dense prefix-LM mask), `flex` (prefix-LM block mask), Ettin / Qwen3 |
+| `cronut.attn_backend` | `sdpa`, `flex` |
+
+Override from the command line w/ `--task.model.image-attn-backend`, `--task.model.cross-attn-backend`, and
+`--task.model.self-attn-backend` (decoder self-attention, or the Cronut trunk), e.g. `--task.model.image-attn-backend flex`.
+
+### Tokenizers
+
+HF tokenizers by name (`facebook/bart-base`, `Qwen/Qwen3-0.6B`, ...) or tiktoken encodings via `tiktoken:<encoding>`
+(`tiktoken:cl100k_base`, `tiktoken:o200k_base`, ...), w/ eos / pad / bos and task tokens appended above the base vocab.
+
+## Usage
+
+### Training / evaluation apps
+
+The apps are installed as `pixr-train` / `pixr-eval` (or `python -m pixr.app.train` / `python -m pixr.app.eval`).
+Distributed runs use `torchrun` (or `python -m torch.distributed.run`):
+
+```bash
+torchrun --nproc_per_node 2 -m pixr.app.train --task read --task.model.name cruller_naflex_siglip2_base \
+  --train-data.source "<idl-wds shards>" --train-data.num-samples 200000 --trainer.num-intervals 40 ...
+pixr-eval --task docvqa --checkpoint-path <checkpoint.pt> \
+  --eval-data.source pixparse/docvqa-single-page-questions --eval-data.split validation --eval-data.format hfds
+```
+
+Scripts in `scripts/` (source checkout) cover the main workflow w/ tuned defaults, see each script for env var options:
+
+```bash
+# 1. pretrain (read all text on a page) on IDL, streamed from the HF hub
+NUM_GPUS=2 MODEL=cruller_naflex_siglip2_base ./scripts/pretrain_idl.sh
+
+# held out reading CER / WER
+CHECKPOINT=output/pretrain_cruller_naflex_siglip2_base_idl/checkpoints/checkpoint-latest.pt ./scripts/eval_ocr.sh
+
+# 2. fine-tune on DocVQA (ANLS on validation each eval interval)
+CHECKPOINT=output/pretrain_cruller_naflex_siglip2_base_idl/checkpoints/checkpoint-latest.pt ./scripts/finetune_docvqa.sh
+
+# 3. evaluate (validation ANLS, or SPLIT=test for an eval server submission file)
+CHECKPOINT=output/finetune_docvqa/checkpoints/checkpoint-latest.pt ./scripts/eval_docvqa.sh
+```
+
+### Command line / config
+
+Apps are configured by dataclasses parsed w/ [jsonargparse](https://github.com/omni-us/jsonargparse):
+
+* `--task read|docvqa` selects the task (and the type of the `task` config section), default `read` for train
+* nested fields map to dotted args, written w/ dashes: `--trainer.opt.learning-rate 1e-4` (underscores also accepted)
+  * `--task.*` -- task: model, tokenizer, image transforms, dtype / amp, max text length, eval generation
+  * `--trainer.*` -- optimization: intervals, warmup, optimizer / schedule (`--trainer.opt.*`), EMA, compile, grad ckpt,
+    FSDP
+  * `--train-data.*` / `--eval-data.*` -- chugr data args (source, format, split, batch size, workers, NaFlex batching)
+* `--config file.yaml` loads a config (CLI args override), `--print-config` dumps the resolved config, `--help` lists all
+* lists are passed as `--train-data.naflex-seq-lens [1024,2048,4096]`
+
+The full resolved config is written to `<output_dir>/<experiment>/config.yaml` and can be passed back via `--config`.
+
+### Distributed / compile
+
+Distributed runs (`torch.distributed.run`) use DDP by default, `--trainer.fsdp true` shards params, grads and
+optimizer state w/ FSDP2 (`fully_shard` per encoder block / decoder layer + root, models list their blocks in
+`fsdp_shard_modules()`). The handling follows the OpenCLIP `TrainingTask` (`prepare_distributed` / `prepare_fsdp` /
+`compile` on the task):
+
+* FSDP mixed precision is a `MixedPrecisionPolicy` (bf16 params, fp32 gradient reduction), autocast is suppressed
+* `--trainer.grad-checkpointing` w/ FSDP applies composable activation checkpointing per block (inline otherwise)
+* `--trainer.compile true --trainer.compile-strategy task|model|blocks`: `task` compiles the task training forward
+  (model + loss), `model` the model before DDP, `blocks` the blocks in place before DDP (checkpoint recompute stays
+  in eager autograd, bounding compiled backward memory). W/ FSDP, blocks are always compiled inside `prepare_fsdp()`,
+  ordered compile -> composable checkpointing -> `fully_shard` (checkpoint hooks on the compiled wrapper so recompute
+  re-enters the FSDP all-gather hooks), `--compile-mode` / `--compile-dynamic` / `--compile-backend` pass through
+* eval and train metrics run w/ the params gathered (`task.full_params()`), so ranks decode independently
+* checkpoints are gathered full state dicts (model, EMA, optimizer), loadable w/o FSDP, model / EMA weights are
+  interchangeable between FSDP and DDP runs (optimizer state is not)
+
+`scripts/check_distributed.py` checks single / DDP / FSDP parity (loss trajectories, final params), grad
+checkpointing, accumulation, compile, EMA, resume and generation through the Trainer on synthetic batches.
+
+### Generation
+
+`pixr.generation` provides batched, KV cached greedy (`generate_greedy`) and beam search (`generate_beam`) decoding,
+`--task.eval-num-beams` selects beam search for eval. Cruller, Fritter and Cronut share one incremental decoding
+interface:
+
+```python
+memory = model.encode_image(image_input)                         # per image, computed once
+logits, state = model.prefill(memory, prompt_ids, prompt_mask)   # left padded prompts
+logits, state = model.step(token_ids, state)                     # one token per hypothesis
+state.reorder(parent_idx)                                        # select / duplicate hypotheses (beam search)
+```
+
+The per image memory and the per hypothesis sequence state are separate (`pixr/models/kv_cache.py`). Cruller's
+cross-attention K / V are projected once inside each decoder layer and shared by all hypotheses of an image (several
+prompts or beams per image, image-major order), only the self-attention caches are reordered. Under FSDP, generation
+must run inside `task.full_params()` (ranks then decode independently).
+
+### LR scaling
+
+Params are grouped by component via the model's `param_layer_map()`: `encoder` (image encoder), `bridge`
+(cross-attention / prefix projector), `decoder` (text decoder / LM trunk). Each has an LR multiplier
+(`--trainer.opt.encoder-lr-scale`, `--trainer.opt.bridge-lr-scale`, `--trainer.opt.decoder-lr-scale`) and the
+encoder / decoder support layer-wise LR decay (`--trainer.opt.encoder-layer-decay 0.75`,
+`--trainer.opt.decoder-layer-decay`). `--trainer.opt.lr-scale-overrides '{"embed_tokens": 2.0}'` adds multipliers
+by param name regex. Training pretrained encoders at full LR from the start wrecks their features before the decoder
+learns to use them (reading never takes off), `scripts/pretrain_idl.sh` uses `ENC_LR_SCALE=0.1`.
+
+### LM loss
+
+The LM head + cross-entropy is fused and chunked over tokens by default (`--task.loss-chunk-size 4096`, 0 to
+materialize full logits), padding is dropped before the head GEMM and each chunk's logits are recomputed in backward,
+so the [tokens, vocab] logits (~150k vocab for Qwen3) are never materialized. `--task.z-loss-weight` (e.g. `1e-4`) adds
+a z-loss (mean squared logsumexp), `--task.label-smoothing` is supported by both paths.
+
+### NaFlex batching
+
+For NaFlex models the train loader can draw a sequence length per batch and size the batch to a token budget, w/
+samples resized to the batch's patch budget (rank-identical schedule, from OpenCLIP's streaming NaFlex loader):
+
+```bash
+--train-data.naflex-seq-lens [1024,2048,4096] --train-data.naflex-seq-len-probs [0.25,0.25,0.5] \
+--train-data.naflex-max-tokens 65536 --train-data.naflex-text-tokens 512 --trainer.opt.loss-batch-scale sqrt
+```
+
+`--train-data.length-bucketing true` groups samples of similar target text length. Neither requires resampled shards.
+
+Checkpoints include the resolved model config and the tokens added to the tokenizer, so fine-tuning and evaluation
+only need `--checkpoint-path`, task tokens are appended and their embeddings initialized from the mean embedding.
+
+## Code Organization
+
+Within `src/pixr`:
+* `app/` - CLI applications for training and evaluation
+  * `app/train.py` - main training CLI entrypoint, will attempt to keep useable across tasks
+  * `app/eval.py` - main evaluation CLI entrypoint
+* `data/` - data loaders, image and text preprocessing
+* `framework/` - lightweight train & evaluation scaffolding on top of canonical PyTorch
+  * `framework/task.py` - `TrainingTask` base (model + objective + eval)
+  * `framework/trainer.py` - `Trainer` (optimizer, LR schedule, AMP, grad accumulation, DDP / FSDP2 / compile setup,
+    EMA, train state)
+* `models/` - modelling code with associated factory methods and helpers
+* `task/` - tasks (`read` pretraining, `docvqa` fine-tune), registered in `task/task_registry.py`
+* `tokenizers/` - tokenizer creation and task token handling
+* `generation.py` - batched, KV cached greedy and beam search decoding (`--task.eval-num-beams`)
+* `utils/` - misc utils that don't have a home (`utils/cli.py` -- jsonargparse config / CLI parsing)
+
+## Concepts & Terminology
+
+Some terms and concepts used in this project that may be a bit unfamiliar.
+
+### Task
+A key organization concept in this project. Package the model with its loss, pre/post-processing, and optimization setup together for a given objective.
+
+Examples of tasks conceptually:
+  * Pretraining a Donut style (image enc + text dec) model on supervised (OCR annotation) doc-text dataset
+  * Pretraining a Pix2Struct style (image enc + text dec) model w/ a dataset of webpage/screenshots and structured, simplified HTML
+  * Pretraining a PaLI style (image enc + text enc-dec) model w/ prefix & masked-token completion on datasets as above
+  * Fine-tuning any of the above pretrained models on a possibly large variety of downstream tasks
+    * Semi-structured doc parsing - receipt, invoice, business cards, etc.
+    * VQA
+    * Captioning
+    * ... and more
+
+With the Task concept, the data pipeline exists outside the task. The data pipeline is coupled to the task by passing the pre-processing functions created within the task to the data pipeline on creation (`task.train_data_cfg()` / `task.eval_data_cfg()`).
+
+A task is a `TrainingTask` (an `nn.Module` wrapping the model) implementing `training_forward(batch) -> (outputs, losses)`,
+optional `train_metrics()`, and `evaluate(loader)`. The generic `Trainer` owns everything else (optimizer, schedule,
+AMP, accumulation, setup ordering for DDP / FSDP2 / compile, EMA, checkpointing state). Tasks are agnostic to the model family.
+
+### Interval
+
+You'll see the term 'interval' in the code, sometimes next to epoch. It's related, but an epoch means 'one complete pass of the dataset', an interval may be an epoch, but it may not. Interval is a span of training between checkpoints, ideally meaningful enough in duration to warrant evaluating and archiving each interval checkpoint.
+
+In OpenCLIP development the term arose when using shard sampling with replacement, where the intervals between checkpoints were determined by limitations on job durations or likelihood of crashes.
+
+## Development
+
+Tests: `pytest tests` (`tests/test_distributed.py` runs FSDP2 / DDP on a single rank NCCL process group, skipped w/o CUDA).
+
+For short test / train / eval runs on two GPUs (source checkout), use the scenario launcher from an environment with editable
+`pixr`, `chugr`, and the local `timm` installed:
+
+```bash
+python scripts/run_review.py --list
+python scripts/run_review.py all --validate                 # parse app args, no GPU jobs
+python scripts/run_review.py smoke --dry-run                # inspect commands
+python scripts/run_review.py smoke --devices 0,1
+python scripts/run_review.py fsdp_compile naflex_repeat prefix_flex
+python scripts/run_review.py data_ranks                     # check rank streams are distinct
+python scripts/run_review.py tests parity
+python scripts/run_review.py all
+```
+
+The launcher runs scenarios sequentially through `pixr.app.train`, `pixr.app.eval`, pytest and
+`scripts/check_distributed.py`. Scenarios cover DDP / FSDP2, gradient accumulation, compile, checkpointing, EMA,
+NaFlex batching, all three model families, attention backends, checkpoint resume and DocVQA fine-tuning.
+`data_ranks` checks actual resampled WDS sample identities across ranks, without loading a model or decoding PDFs.
+It fails if any two ranks receive an identical stream. App runs also fail if no updates or non-finite losses are logged.
+`parity` compares FP32 DDP / FSDP2 model and EMA parameters and fails above a relative L2 difference of `1e-5`;
+its resumed losses and generated outputs are also recorded for inspection.
+
+Defaults are short execution checks: random initialization, 64 global samples per interval, batch 2 per GPU,
+accumulation 2, 256 image patches and 128 text tokens. `--pretrained` enables hub weight initialization for the
+app runs; tokenizers and Qwen3 / Ettin configs can still require downloads without it. `pretrained_tests` downloads
+HF reference weights; `docvqa` uses the HF DocVQA dataset. `--samples`, `--batch-size`, `--accum`, `--patches`,
+`--text-length`, and `--workers` adjust the workload. These short random-weight runs do not measure convergence.
+
+The default data paths are sibling `../data/idl/idl-train-000{01..07}.tar` for training and shard `00000` for eval;
+override with `--train-source` / `--eval-source`. `naflex_repeat` needs at least `GPUs * workers` training shards.
+Logs, exact commands, resolved configs, checkpoints and `summary.json` go into a new `../runs/review-<timestamp>`
+directory (`--output` overrides it). Each command has a 20-minute timeout (`--timeout` in seconds). Failed scenarios
+are logged and the remaining scenarios continue; `--fail-fast` stops at the first failure. The launcher exits
+nonzero on failure and prints a `tail -f` command for the active log.
+
+## Acknowledgements
+
+Model and training design draws on [Donut](https://github.com/clovaai/donut), Pix2Struct, PaLI, LLaVA and GenLIP.
+Image encoders come from [`timm`](https://github.com/huggingface/pytorch-image-models), the training task / FSDP2
+handling and NaFlex batching follow [OpenCLIP](https://github.com/mlfoundations/open_clip), datasets are on the
+Hugging Face Hub under the [`pixparse`](https://huggingface.co/pixparse) organization.
+
+## License
+
+Apache 2.0, see [LICENSE](LICENSE).