naklitechie/indictrans2-en-indic-dist-200M-ONNX

IndicTrans2 200M (en→indic) — ONNX bundle for in-browser inference

ONNX-exported version of naklitechie/indictrans2-en-indic-dist-200M (itself a verbatim mirror of ai4bharat/indictrans2-en-indic-dist-200M, AI4Bharat's distilled 200M en→indic model).

Built for use with Transformers.js and onnxruntime-web in the browser, with fast BPE tokenizer.json files that don't require the SentencePiece WASM runtime.

Benchmarks vs PyTorch original

Token-level AND text-level bit-exact equivalent to the original PyTorch model at naklitechie/indictrans2-en-indic-dist-200M on greedy decoding.

Test	Pass rate
528 fixtures × greedy decode, token-exact match vs PyTorch	528/528 (100%)
528 fixtures × greedy decode, text-exact match vs PyTorch	528/528 (100%)

Fixtures span 4 categories (generic, politics, numerals, lexicon) × 11 main Indic languages × 12 sentences each. The PyTorch baseline was captured by running the original model with num_beams=1, do_sample=False, max_new_tokens=128.

The verification chain:

fp32 ONNX bundle  →  encoder.run + decoder.run loop (onnxruntime CPU)  →  output token IDs
                                                                             ‖ exact-match
PyTorch original  →  model.generate(num_beams=1)                       →  output token IDs

Reproducer: scripts/04_parity_test.py in the source repo. Full report: parity_report.json in this repo.

See the int8-quantized variant at naklitechie/indictrans2-en-indic-dist-200M-ONNX-int8 for a 4× smaller bundle (~360 MB) at the cost of ~20% sentences drifting in word choice. Use this fp32 repo for any production / research use.

Files

File	Purpose	Size
encoder_model.onnx + .data	encoder graph (fp32, weights externalised)	280 MB
decoder_model.onnx + .data	decoder first-step (no past KV in, full KV out)	530 MB
decoder_with_past_model.onnx + .data	decoder steps 2..N (past KV in/out)	495 MB
tokenizer_src.json	encoder-side fast tokenizer (BPE + lang-tag AddedTokens + NFKC)	3.2 MB
tokenizer_tgt.json	decoder-side fast tokenizer	17 MB
tokenizer_meta.json	dict-size cutoffs for the extended-vocab → <unk> remap	<1 KB
config.json, generation_config.json	model config	—
dict.SRC.json, dict.TGT.json	raw vocabs (~3.7 MB combined)	—
model.SRC, model.TGT	raw SPM models for slow-tokenizer fallback	~3.8 MB
tokenization_indictrans.py + tokenizer_config.json + special_tokens_map.json	HF AutoTokenizer.from_pretrained compat	—

Usage (Python, onnxruntime)

import json, numpy as np, onnxruntime as ort
from tokenizers import Tokenizer
from huggingface_hub import snapshot_download

snap = snapshot_download(repo_id="naklitechie/indictrans2-en-indic-dist-200M-ONNX")

src = Tokenizer.from_file(f"{snap}/tokenizer_src.json")
tgt = Tokenizer.from_file(f"{snap}/tokenizer_tgt.json")
meta = json.load(open(f"{snap}/tokenizer_meta.json"))

enc = ort.InferenceSession(f"{snap}/encoder_model.onnx")
dec = ort.InferenceSession(f"{snap}/decoder_model.onnx")
decp = ort.InferenceSession(f"{snap}/decoder_with_past_model.onnx")

# Tokenize "Who will win the election?" en→hi
text = f"eng_Latn hin_Deva Who will win the election?"
e = src.encode(text)
input_ids = np.array([[i if i < meta["src_dict_size"] else meta["unk_id"] for i in e.ids]], dtype=np.int64)
attn_mask = np.array([e.attention_mask], dtype=np.int64)

# Encoder
enc_h = enc.run(["last_hidden_state"], {"input_ids": input_ids, "attention_mask": attn_mask})[0]

# Greedy decode (see scripts/04_parity_test.py in the source repo for the full loop)
# ...

For the complete end-to-end greedy loop, see 04_parity_test.py.

Production caveats

1. Extended-vocab remap on encoder input. tokenizer_src.json has 33,888 tokens (32,322 canonical + 1,566 SPM-only chars BPE needs as merge halves). Any output ID >= tokenizer_meta.src_dict_size (= 32322) must be replaced with the <unk> ID (3) before feeding to ONNX — the model's encoder embedding only goes up to 32,322. Same applies to tokenizer_tgt.json with tgt_dict_size = 122,672. The slow HF tokenizer does this implicitly via encoder.get(token, unk_id).

2. Devanagari-normalized output. For non-Devanagari Indic languages (Tamil, Bengali, Telugu, etc.), the model emits its output in Devanagari script. AI4Bharat's IndicProcessor.postprocess_batch handles the script conversion. A JS port lives in the source repo at browser-prep/js/indic_processor.js.

3. fp32, not yet quantized. The bundle is ~1.3 GB total — not browser-ready for many users without quantization. q4f16 weights bring this to roughly 200 MB; tooling for that is in progress.

Source

Built via the 01_export_onnx_manual.py → 02_build_tokenizer.py → 03_capture_truth.py → 04_parity_test.py pipeline at prashnam/prashnam-voice (browser-prep/scripts/).

License

MIT (preserved from upstream AI4Bharat). See LICENSE and NOTICE.md.