Metadata-Version: 2.5
Name: grtools
Version: 2026.9.1
Summary: Research tools and utilities for data processing, benchmarking, and analysis
Project-URL: Repository, https://github.com/artefactory/argimi
Author-email: Hicham Randrianarivo <hicham.randrianarivo@artefact.com>
License-Expression: MIT
License-File: LICENSE
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Requires-Python: <3.13,>=3.10
Requires-Dist: absl-py>=1.4.0
Requires-Dist: beartype>=0.21.0
Requires-Dist: chex>=0.1.0
Requires-Dist: clu>=0.0.12
Requires-Dist: datasets>=4.5.0
Requires-Dist: einops>=0.8.0
Requires-Dist: etils[eapp,edc,enp,epath,etqdm]>=1.10.0
Requires-Dist: flax>=0.12.7
Requires-Dist: google-metrax>=0.0.1
Requires-Dist: grain>=0.2.16
Requires-Dist: ir-measures>=0.4.3
Requires-Dist: jax>=0.10; sys_platform == 'darwin'
Requires-Dist: jax[cuda12]>=0.10; sys_platform == 'linux'
Requires-Dist: jaxtyping>=0.3.7
Requires-Dist: keras-hub<0.31,>=0.25.1
Requires-Dist: keras<3.15,>=3.12.1
Requires-Dist: nvidia-ml-py>=13.590.48
Requires-Dist: optax>=0.2.8
Requires-Dist: orbax-checkpoint>=0.11.19
Requires-Dist: orjson>=3.11.7
Requires-Dist: safetensors>=0.4.0
Requires-Dist: tokenizers>=0.15.0
Requires-Dist: toolz>=1.0.0
Requires-Dist: trackio>=0.2.0
Provides-Extra: preprocessing
Requires-Dist: datatrove[cli,io,processing]>=0.3.0; extra == 'preprocessing'
Description-Content-Type: text/markdown

# 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

```sh
pip install grtools
```
