grtools

Core JAX/Flax training infrastructure for generative retrieval research: models (T5 seq2seq, diffusion denoisers), training loops, optimizers, callbacks (evaluation, checkpointing, monitoring), data pipeline transforms, and configs.

Contents

• grtools.models — T5 encoder-decoder and diffusion (MDLM/BD3LM-style) backbones.
• grtools.training — a jax.lax.scan-based training loop, sharding, EMA, samplers.
• grtools.callbacks — periodic evaluation, checkpointing, and monitoring hooks.
• grtools.dataloading / grtools.pipeline — tokenization, docid codecs, and document/query transform steps for building generative-retrieval datasets.
• grtools.retrieval — grounding strategies (trie-constrained decode, Hamming-nearest snap) and retrieval evaluation metrics.
• grtools.configs — dataclass configs (simple_parsing.SerializableMixin) composing training, optimizer, schedule, and callback settings.

Requirements

Python 3.10–3.12. JAX with CUDA extras on Linux, CPU/Metal on macOS.

Installation

pip install grtools