A framework to streamline the research of neural sequence models, from machine translation to LLMs.

Documentation · Quickstart · Examples · Report a bug

AutoNMT is a modular research toolkit that takes the repetitive half of sequence-model experimentation - tokenization, training, scoring, logging, plotting, file management - off your hands so you can focus on the model. Declare a grid of datasets × language pairs × subword models × vocab sizes, and AutoNMT runs the cross-product, persists every intermediate artifact on disk, and produces a single comparable report at the end. It started as a neural machine translation toolkit; the same pipeline now also trains decoder-only LLMs and encoder-only masked LMs.

Every layer - datasets, vocabularies, models, decoding, metrics, reports - is designed to be subclassed, replaced, or extended via callable hooks, so researchers can plug in custom components without forking the core. The same script can train AutoNMT's own PyTorch Lightning models, fine-tune HuggingFace seq2seq checkpoints, or shell out to Fairseq - backends are swapped by changing one class.

• Grid-first API - describe an experiment as a cross-product, not a for-loop.
• Three model families - encoder-decoder (translation), decoder-only LLMs (GPT), and encoder-only masked LMs (MLMTransformer), on one pipeline - trained in-house or fine-tuned from HuggingFace.
• Pluggable backends - AutonmtTranslator (Lightning), HuggingFaceTranslator, FairseqTranslator (deprecated); plus LMTrainer / MLMTrainer for language models.
• Reproducible by construction - every stage writes to a numbered folder; re-runs skip completed steps.
• Subword variants out of the box - word, char, bytes, bpe, unigram, with optional byte fallback.
• Built-in evaluation - sacreBLEU, BERTScore, COMET, HuggingFace metrics, wired into the report.
• Extension-friendly core - subclass LitSeq2Seq or pass callable hooks instead of patching internals.

Requires Python 3.12+.

pip install -e .                       # core
pip install -e '.[hf]'                 # HuggingFace dataset loader
pip install -e '.[hf-models]'          # HuggingFace model backends (seq2seq + LLM/MLM)
pip install -e '.[wandb]'              # W&B logger
pip install -e '.[all]'                # everything above

See the installation guide for optional extras and GPU notes.

Fetch a dataset from HuggingFace, train a small Transformer, and score it - in one script:

from autonmt.datasets import ParallelCorpusBuilder
from autonmt.datasets.sources.hf_loader import download_hf_dataset
from autonmt.backends import AutonmtTranslator
from autonmt.backends._base.config import FitConfig, PredictConfig
from autonmt.core.nn.models import Transformer

# 1. Pull a parallel corpus from the Hub into AutoNMT's on-disk layout.
download_hf_dataset(
    hf_id="bentrevett/multi30k", base_path="datasets/quickstart",
    dataset_name="multi30k", lang_pair="de-en", src_field="de", tgt_field="en",
)

# 2. Declare the grid (one cell here) and materialize it: clean, tokenize, encode.
builder = ParallelCorpusBuilder(
    base_path="datasets/quickstart",
    datasets=[{"name": "multi30k", "languages": ["de-en"], "sizes": [("original", None)]}],
    encoding=[{"subword_models": ["bpe"], "vocab_sizes": [4000]}],
).build()

# 3. Take the single variant and build its source/target vocabularies.
train_ds = builder.get_train_ds()[0]
src_vocab, tgt_vocab = train_ds.build_vocabs(max_tokens=150)

# 4. Bind a Transformer to the native backend, then train and score.
trainer = AutonmtTranslator.from_dataset(
    train_ds,
    model=Transformer.from_vocabs(src_vocab, tgt_vocab),
    src_vocab=src_vocab, tgt_vocab=tgt_vocab,
    run_prefix="quickstart",
)
trainer.fit(train_ds, config=FitConfig(max_epochs=3, batch_size=128))
scores = trainer.predict(builder.get_test_ds(), config=PredictConfig(metrics={"bleu"}))

Full walkthroughs live in examples/ - a step-by-step tutorial that builds from this snippet up to a multi-axis grid and a HuggingFace backend swap.

Full docs are published at salvacarrion.github.io/autonmt:

• Get started - install, first experiment, understanding the output.
• User guide - data, models, training, generation, evaluation, backends.
• How-to guides - task-oriented recipes for common workflows.
• Concepts - design philosophy, mental model, architecture, on-disk layout, reproducibility.
• API reference - autodoc from docstrings.

Contributions are welcome - bug reports, feature requests, docs fixes, or new backends/metrics. See docs/contributing.md for the dev setup, test commands, and PR conventions.

pytest tests/
flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics

If you use AutoNMT in academic work, please cite:

@misc{carrión2023autonmtframeworkstreamlineresearch,
      title={AutoNMT: A Framework to Streamline the Research of Seq2Seq Models},
      author={Salvador Carrión and Francisco Casacuberta},
      year={2023},
      eprint={2302.04981},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2302.04981},
}

MIT © Salva Carrión.