indictrans2-en-indic-1B-onnx

indictrans2-en-indic-1B-onnx

ONNX export of ai4bharat/indictrans2-en-indic-1B, the full-size IndicTrans2 translation model from AI4Bharat / IIT Madras (the distilled 200M variant is published separately at indictrans2-en-indic-dist-200M-onnx).

Direction: English into 22 Indic languages. src_lang is always eng_Latn.

Files

Sizes: fp32 6.7 GB, int8 1.7 GB.

If you load these graphs with your own code (not ORTModelForSeq2SeqLM.from_pretrained against this repo directly), fetch decoder_model.onnx_data and decoder_with_past_model.onnx_data alongside their .onnx files — the encoder has no external data at this size, but both decoder graphs do.

Preprocessing is mandatory

IndicTrans2 does not take raw text. You must run IndicProcessor from IndicTransToolkit first:

pip install IndicTransToolkit optimum[onnxruntime] transformers sentencepiece

IndicProcessor.preprocess_batch normalises the script, applies Indic-specific punctuation and numeral handling, protects entities, and prefixes the two language tags. Skipping it gives silently wrong output — the model still produces fluent-looking text, but for several target scripts (anything transliterated through Devanagari, e.g. Tamil, Bengali) it comes back correctly-worded but in the wrong script, and nothing raises.

postprocess_batch reverses the entity protection and the script transliteration. It is also required.

Language tags

The tags are plain text prefixed to the source sentence, in the order <src_lang> <tgt_lang> <sentence>. IndicProcessor adds them for you; you only pass src_lang= and tgt_lang=.

For example, ip.preprocess_batch(["This is a test."], src_lang="eng_Latn", tgt_lang="hin_Deva") produces "eng_Latn hin_Deva This is a test .".

The tags are <iso639-3>_<ISO 15924 script>. Languages with two scripts have two tags. Supported tags:

Usage

import torch
from transformers import AutoTokenizer
from optimum.onnxruntime import ORTModelForSeq2SeqLM
from IndicTransToolkit.processor import IndicProcessor

REPO = "TigreGotico/indictrans2-en-indic-1B-onnx"

tokenizer = AutoTokenizer.from_pretrained(REPO, trust_remote_code=True)
model = ORTModelForSeq2SeqLM.from_pretrained(REPO, trust_remote_code=True, use_cache=True)
ip = IndicProcessor(inference=True)          # REQUIRED - see "Preprocessing"

sentences = ["The weather is nice today and the children are playing outside.", "I would like a cup of tea with a little sugar."]
src_lang, tgt_lang = "eng_Latn", "hin_Deva"

batch = ip.preprocess_batch(sentences, src_lang=src_lang, tgt_lang=tgt_lang)
enc = tokenizer(batch, return_tensors="pt", padding=True, truncation=True, max_length=256)

with torch.inference_mode():
    out = model.generate(**enc, num_beams=4, max_new_tokens=128, length_penalty=1.0, use_cache=True)

decoded = tokenizer.batch_decode(out, skip_special_tokens=True)
print(ip.postprocess_batch(decoded, lang=tgt_lang))

To use the quantized graphs, pass subfolder="int8".

Sample output

Parity against the PyTorch original

100 sentences from ai4bharat/IN22-Gen (English source, 3 Indic targets: Hindi, Bengali, Tamil), both greedy (num_beams=1) and beam search (num_beams=4), max_new_tokens=128, length_penalty=1.0, against AutoModelForSeq2SeqLM.from_pretrained(..., trust_remote_code=True, torch_dtype=torch.float32). Preprocessing/postprocessing via the vendored IndicProcessor on both sides.

Quality against human reference (IN22-Gen)

chrF of the model's own output against the dataset's human reference translation (this measures translation quality, not fp32/ONNX parity, which is the table above):

Differential / language-identification check

The same English sentence translated into Hindi, Bengali and Tamil produces three distinct outputs, each independently confirmed to be in the requested language by linguonnx's GlotLID detector (linguonnx.load_detector()): hi → detected hi, bn → detected bn, ta → detected ta. No <unk>, ⁇, �, or repeated-phrase loops found across any of the greedy/beam4/int8 outputs checked for this release.

Limits

Maximum sequence length is 256 tokens on both sides. The sinusoidal position table is frozen into the graph at export time, so longer inputs are not supported. Truncate with max_length=256.

Export route

Exported with optimum.exporters.onnx.onnx_export_from_model and a custom OnnxConfig registered for the IndicTrans model type. The config subclasses M2M100OnnxConfig and remaps the field names IndicTrans2 uses (encoder_embed_dim instead of d_model, encoder_vocab_size for the dummy-input vocab bound since IndicTrans2 has separate source/target vocabularies). Recipe and gotchas (registration casing, the vocab split, the .onnx.data→.onnx_data external-data rename, the decoder-merge protobuf limit at this scale) are documented in scripts/export/README.md in the linguonnx repo.

optimum-cli export onnx alone does not work: it has no config for the IndicTrans architecture.

Attribution

The model is the work of AI4Bharat, IIT Madras.

@article{gala2023indictrans2,
  title   = {IndicTrans2: Towards High-Quality and Accessible Machine Translation Models for all 22 Scheduled Indian Languages},
  author  = {Jay Gala and Pranjal A. Chitale and Raghavan AK and Varun Gumma and Sumanth Doddapaneni and Aswanth Kumar and Janki Nawale and Anupama Sujatha and Ratish Puduppully and Vivek Raghavan and Pratyush Kumar and Mitesh M. Khapra and Raj Dabre and Anoop Kunchukuttan},
  journal = {Transactions on Machine Learning Research},
  year    = {2023}
}

Licence: MIT, same as the original.