Python API

The public pragmatiq.api surface — generated from the live docstrings so signatures never drift.

pragmatiq is an independent implementation inspired by the PRAGMA paper (arXiv 2604.08649) and is not affiliated with or endorsed by Revolut.

pragmatiq/api.py is the single library surface behind the CLI, the notebooks, and production callers. Notebooks also get PragmaModel.from_pretrained(run) and model.embed_records(list_of_dicts) for interactive use without the shard pipeline.

The reference below is generated from the live function signatures and docstrings by scripts/docs_facts.py (and checked in CI), so it always matches the installed code.

api.synthesize

Generate a synthetic dataset.

synthesize(config: 'str | Path | dict[str, Any] | None' = None, out: 'str | Path' = 'data/synth', n_users: 'int | None' = None, seed: 'int | None' = None, n_workers: 'int' = 0, write_report: 'bool | None' = None, **overrides: 'Any') -> 'dict[str, Any]'

api.tokenize

Fit (or load) the tokenizer and write tokenized parquet shards + index.

tokenize(data_dir: 'str | Path', out: 'str | Path', config: 'str | Path | dict[str, Any] | None' = None, tokenizer_dir: 'str | Path | None' = None, max_users: 'int | None' = None, rows_per_shard: 'int' = 4096, n_workers: 'int' = 0) -> 'dict[str, Any]'

api.pretrain

Pretrain a pragmatiq model on tokenized shards.

pretrain(shard_dir: 'str | Path', run_name: 'str', model_size: 'str' = 'small', config: 'str | Path | dict[str, Any] | None' = None, runs_root: 'str | Path' = 'runs', resume: 'str | None' = None, **overrides: 'Any') -> 'dict[str, Any]'

api.pretrain_plan

Resolve a pretrain configuration without training (``pretrain --show-config``).

pretrain_plan(shard_dir: 'str | Path', model_size: 'str' = 'small', config: 'str | Path | dict[str, Any] | None' = None, run_name: 'str' = 'plan', runs_root: 'str | Path' = 'runs', resume: 'str | None' = None, **overrides: 'Any') -> 'dict[str, Any]'

api.finetune

LoRA fine-tune a trained model's adapters + head on a label table.

finetune(shard_dir: 'str | Path', run: 'str | Path', label_path: 'str | Path', config: 'str | Path | dict[str, Any] | None' = None, device: 'str' = 'auto', **overrides: 'Any') -> 'dict[str, Any]'

api.embed

Embed every user in ``shard_dir`` with a trained model.

embed(shard_dir: 'str | Path', run: 'str | Path', out: 'str | Path | None' = None, token_budget: 'int' = 16384, device: 'str' = 'auto') -> 'dict[str, Any]'

api.probe

Probe a trained model on a label table; compares to a raw-count baseline.

probe(shard_dir: 'str | Path', run: 'str | Path', label_path: 'str | Path', device: 'str' = 'auto', token_budget: 'int' = 16384, seed: 'int' = 0, with_baseline: 'bool' = True, probe_model: 'str' = 'gbdt', staleness_window: 'str | int | None' = None) -> 'dict[str, Any]'

api.uplift

Evaluate communication-campaign uplift on a trained model.

uplift(shard_dir: 'str | Path', run: 'str | Path', label_path: 'str | Path', device: 'str' = 'auto', token_budget: 'int' = 16384, seed: 'int' = 0, learner: 'str' = 't') -> 'dict[str, Any]'

api.export

Export the dense ONNX reformulation of the model from one example user.

export(run: 'str | Path', shard_dir: 'str | Path', out: 'str | Path' = 'pragmatiq_embedder.onnx', device: 'str' = 'auto') -> 'dict[str, Any]'

api.benchmark

Benchmark batch-embedding throughput and write a results markdown file.

benchmark(run: 'str | Path', shard_dir: 'str | Path', device: 'str' = 'auto', out: 'str | Path' = 'benchmark_results.md', max_users: 'int | None' = None, precision: 'str' = 'auto') -> 'dict[str, Any]'

api.gnn

Run the four-arm AML GNN ablation.

gnn(shard_dir: 'str | Path', run: 'str | Path', transfers_path: 'str | Path', aml_label_path: 'str | Path', seeds: 'tuple[int, ...]' = (0, 1, 2), device: 'str' = 'auto', epochs: 'int' = 150) -> 'dict[str, Any]'

api.validate

Validate a raw dataset against the data contract.

validate(data_dir: 'str | Path') -> 'dict[str, Any]'

api.quickstart

End-to-end smoke: synth → tokenize → nano pretrain → probe.

quickstart(out: 'str | Path' = 'runs/quickstart', n_users: 'int' = 50000, seed: 'int' = 0, model_size: 'str' = 'nano', max_steps: 'int' = 400, n_workers: 'int' = 0) -> 'dict[str, Any]'

api.runs_list

List runs under ``runs_root`` with their last logged step/loss/metrics.

runs_list(runs_root: 'str | Path' = 'runs') -> 'list[dict[str, Any]]'

api.runs_compare

Compare several runs' last metrics side by side (missing runs flagged).

runs_compare(names: 'list[str]', runs_root: 'str | Path' = 'runs') -> 'list[dict[str, Any]]'

api.calibrate

Fit generator priors to bank-shareable aggregate statistics.

calibrate(stats: 'str | Path', config: 'str | Path | dict[str, Any] | None' = None, out: 'str | Path | None' = None) -> 'dict[str, Any]'

api.info

Describe this installation: versions, devices, kernels, extras and env vars.

info() -> 'dict[str, Any]'

Interactive entry points

Beyond api, the model exposes two notebook-friendly methods:

from pragmatiq.models.pragmatiq import PragmaModel

model = PragmaModel.from_pretrained("runs/demo")     # loads model + tokenizer on cuda if visible, else cpu
emb = model.embed_records([                           # plain dicts, no shards needed
    {"user_id": "u1", "events": [
        {"ts": 1_700_000_000_000_000, "source": "transaction",
         "fields": {"amount": "42.10", "merchant": "TESCO"}}],
     "attributes": {"country": "GB"}, "lifelong": []},
])  # -> np.ndarray [n_users, dim]

from_pretrained verifies the checkpoint's tokenizer hash against the run's tokenizer and refuses to run on a mismatch. embed_records(records, precision="auto", token_budget=None) runs under torch.inference_mode — bf16 on CUDA, fp32 on CPU — and token_budget splits a large request into bounded forward passes.