ShrutamMT-v1

ShrutamMT-v1 🎧

"श्रुतम्" (Shrutam) = "that which is heard" — a multilingual 11 Indic languages → English neural machine translation model, built as a Transformer trained completely from scratch (no pretrained weights, no fine-tuning shortcuts) on the AI4Bharat Samanantar parallel corpus.

Supports: Assamese, Bengali, Gujarati, Hindi, Kannada, Malayalam, Marathi, Odia, Punjabi, Tamil, Telugu → English, all in a single shared model.

This is a from-first-principles implementation of the original "Attention Is All You Need" (Vaswani et al., 2017) transformer architecture — encoder, decoder, multi-head attention, positional encodings, label smoothing, Noam learning-rate schedule, beam search — all hand-written in PyTorch, scaled up and extended for multilingual training.

Model Details

Design choices beyond the base Transformer paper

• 3-way tied embeddings (source embedding = target embedding = output projection) — fewer params, better generalization across languages sharing a joint vocabulary
• Source language-tag tokens (<hi>, <bn>, <ta>, ...) prepended to input, so the encoder knows which language it's translating from — same idea used in Google's multilingual NMT and IndicTrans
• Length-ratio filtering during data cleaning (0.3x–3.0x) to remove noisy/misaligned pairs
• Frequent checkpointing (every 500 steps) to survive GPU driver instability during the long training run
• Checkpoint averaging (last 5 checkpoints) as the final released weights

Results

Best result (final model, step 400,000)

Per-language BLEU / chrF

Evaluated on 30 held-out Samanantar test sentences per language (330 total) with beam search (beam size 4):

Training curve (validation loss / perplexity)

Full interactive curves (train loss, val loss, val ppl, learning rate) are on the live Weights & Biases dashboard.

The curve shows the classic Noam-schedule shape: fast initial drop during warmup, a long plateau through the middle of training while LR is still relatively high, then steady improvement as LR decays in the final third — val PPL nearly halved from step 20,000 (70.0) to step 400,000 (16.9).

Lower-resource languages (Assamese: 137k training pairs) score lower than higher-resource ones (Bengali: 8.5M pairs) — expected behavior, directly correlated with per-language data volume in Samanantar.

Sample translations

TA: உள்ளே செல்ல அனுமதி கிடையாது.
REF: Nobody is allowed to go inside.
HYP: Entry is not allowed.

TE: ఏం పని వుండి చేశారో ఏమో!
REF: What have you just done !
HYP: What have you done!

Repository Structure

.
├── checkpoint/
│   └── model_step400000.pt   # final trained weights (step 400,000)
├── tokenizer/
│   ├── spm_joint.model       # SentencePiece BPE model (joint, 64k vocab, 11 scripts + English)
│   └── spm_joint.vocab
└── code/
    ├── config.py              # all hyperparameters
    ├── models/                # encoder/decoder/attention/embedding — full from-scratch architecture
    ├── dataset/                # data loading + token-bucket batching + masks
    ├── prepare_data.py         # builds the multilingual train/val split from Samanantar
    ├── train_tokenizer.py      # SentencePiece training script
    ├── train.py                # training loop (AMP, checkpointing, Noam LR, wandb)
    ├── translate.py             # beam-search inference (single sentence, language-tagged)
    ├── infer_and_compare.py     # per-language BLEU/chrF evaluation
    ├── average_checkpoints.py   # checkpoint averaging utility
    └── run_train_loop.sh        # auto-restart wrapper for long unattended training

Usage

import torch
import sentencepiece as spm
import sys
sys.path.insert(0, "code")

import config
from models.transformer import Transformer
from dataset.masks import make_src_mask, make_tgt_mask

sp = spm.SentencePieceProcessor()
sp.load("tokenizer/spm_joint.model")

device = "cuda" if torch.cuda.is_available() else "cpu"
model = Transformer(
    src_vocab_size=sp.get_piece_size(),
    tgt_vocab_size=sp.get_piece_size(),
    d_model=config.D_MODEL,
    n_heads=config.N_HEADS,
    n_encoder_layers=config.N_ENCODER_LAYERS,
    n_decoder_layers=config.N_DECODER_LAYERS,
    d_ff=config.D_FF,
    dropout=config.DROPOUT,
    max_len=config.MAX_LEN,
    tie_embeddings=config.TIE_EMBEDDINGS,
).to(device)
ckpt = torch.load("checkpoint/model_step400000.pt", map_location=device)
model.load_state_dict(ckpt["model"])
model.eval()

# translate — see code/translate.py for the full beam_search() implementation
# input must be prefixed with a language tag, e.g. "<hi> your text here"

Or run the included script directly:

python3 code/translate.py --text "यह एक परीक्षण है" --lang hi --checkpoint checkpoint/model_step400000.pt --beam_size 4

Training Data

Trained on ShrutamMT-v1-data, a cleaned and multilingual-tagged derivative of AI4Bharat Samanantar — the largest publicly available parallel corpora collection for 11 Indic languages, created by the AI4Bharat lab. All credit for the underlying parallel data collection and curation goes to AI4Bharat; this repository only adds cleaning, multilingual formatting, and language tagging on top of their original release.

Limitations

• Beam search inference in translate.py is not KV-cached / not batched across beams — correctness-oriented, not throughput-optimized.
• Trained on news/formal-register text (Samanantar is largely mined from news and web sources); informal or code-mixed text may translate less reliably.
• Direction is X→English only (not English→X, and not Indic-to-Indic directly).
• Lower-resource languages (Assamese, Odia, Punjabi) have noticeably lower quality than higher-resource ones (Bengali, Telugu, Gujarati), proportional to training data volume.

Citation

This model is entirely dependent on the Samanantar dataset. Please cite AI4Bharat's work if you use this model:

@article{ramesh2021samanantar,
  title={Samanantar: The Largest Publicly Available Parallel Corpora Collection for 11 Indic Languages},
  author={Ramesh, Gowtham and others},
  journal={Transactions of the Association for Computational Linguistics},
  year={2022}
}