indictrans2-indic-indic-dist-320M-onnx

indictrans2-indic-indic-dist-320M-onnx

ONNX export of ai4bharat/indictrans2-indic-indic-dist-320M, the IndicTrans2 translation model from AI4Bharat / IIT Madras.

Direction: Indic to Indic directly, with no English pivot. Both src_lang and tgt_lang are Indic tags.

Files

Sizes: fp32 2.05 GB, int8 519 MB.

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 in the wrong language or with corrupted content.

postprocess_batch reverses the entity protection. 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-indic-indic-dist-320M-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 = ["आज मौसम अच्छा है और बच्चे बाहर खेल रहे हैं।", "मुझे थोड़ी चीनी के साथ एक कप चाय चाहिए।"]
src_lang, tgt_lang = "hin_Deva", "tam_Taml"

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=64, use_cache=True)

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

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

Sample output

Parity against the PyTorch original

10 sentences over 5 Indic languages (Hindi, Tamil, Bengali, Marathi, Malayalam), num_beams=4, max_new_tokens=64, greedy string comparison against AutoModelForSeq2SeqLM.from_pretrained(..., trust_remote_code=True).

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 instead of vocab_size). Opset 17. Optimum validated encoder and both decoder graphs against PyTorch at atol=1e-3 during export.

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.