DOCSEARCHbyAlgolia
Submit Docs

IndicTrans2 Documentation

repository·main·Indexed 19 days ago

https://github.com/ai4bharat/indictrans2

IndicTrans2 is an open-source transformer-based multilingual Neural Machine Translation (NMT) model supporting 22 scheduled Indic languages across scripts including Devanagari, Perso-Arabic, Ol Chiki, Meitei, and Latin. The repository provides tools for HuggingFace integration, fairseq checkpoint conversion, LoRA fine-tuning, and deployment via Triton server on Azure Machine Learning. It includes pre-trained Base and Distilled models for En-Indic, Indic-En, and Indic-Indic translation directions.

Tokens
8.4K
Snippets
36
Records
41
Agent score
66%

What's inside IndicTrans2

  1. Train your own SPM models and Fairseq dictionary

    main

    If you wish to use custom data instead of the provided BPCC datasets, follow these steps to train your own SentencePiece (SPM) models and learn a Fairseq dictionary:

    1. Collect Data: Gather balanced English and Indic monolingual data (recommended ~3M sentences per language-script combination).
    2. Unify Scripts: Perform script unification for Indic languages using scripts/preprocess_translate.py and concatenate all Indic data into a single file.
    3. Train SPM: Train two separate SPM models (one for English, one for Indic) using spm_train.
    4. Learn Dictionary: Copy the trained SPM models to your experiment directory and run prepare_data_joint_training.sh to learn the Fairseq dictionary.

    Important: Use the same Fairseq dictionary for any subsequent fine-tuning experiments.

    # Train SPM model
spm_train --input=train.indic --model_prefix=<model_name> --vocab_size=<vocab_size> --character_coverage=1.0 --model_type=BPE

# Learn Fairseq dictionary
bash prepare_data_joint_training.sh <exp_dir>

  2. Download IndicTrans2 Multilingual Translation Models

    main

    IndicTrans2 provides several model variants for different translation directions (English-to-Indic, Indic-to-English, and Indic-to-Indic) in both Base and Distilled versions. You can download them via Fairseq (tar.gz) or Hugging Face (HF) formats.

    Model Variants

    • Base Models: Used for benchmarking.
    • Distilled Models: Smaller, more efficient versions.

    Available Directions

    • En-Indic: English to Indic languages.
    • Indic-En: Indic languages to English.
    • Indic-Indic: Indic language to Indic language.

    Refer to the official repository for specific Hugging Face links for each language pair.

  3. Run inference using Fairseq (Bash or Python)

    main

    You can perform translation using the Fairseq backend via a Bash script or a Python interface.

    Bash Interface

    Use the joint_translate.sh script to run inference on a file.

    Python Interface

    Use the Model class from inference.engine. Initialize the model by specifying the checkpoint directory and setting model_type="fairseq".

    # Bash interface
bash joint_translate.sh <infname> <outfname> <src_lang> <tgt_lang> <ckpt_dir>
    from inference.engine import Model

model = Model(ckpt_dir, model_type="fairseq")

sents = [sent1, sent2, ...]

# For a batch of sentences
model.batch_translate(sents, src_lang, tgt_lang)

# For a paragraph
model.translate_paragraph(text, src_lang, tgt_lang)

  4. Build the IndicTrans2 Triton server Docker image

    main

    To build the Docker image for the Triton server, navigate to the indicTrans2/inference/ directory and use the Dockerfile located in the triton_server/ subdirectory. Tag the image as indictrans2_triton.

    cd indicTrans2/inference/
docker build -f triton_server/Dockerfile -t indictrans2_triton .

  5. Run inference using CT2 (Python)

    main

    To use the CTranslate2 (CT2) ported models, use the Model class from inference.engine and set the model_type to "ctranslate2".

    from inference.engine import Model

model = Model(ckpt_dir, model_type="ctranslate2")

sents = [sent1, sent2, ...]

# For a batch of sentences
model.batch_translate(sents, src_lang, tgt_lang)

# For a paragraph
model.translate_paragraph(text, src_lang, tgt_lang)

  7. Publish an online inference endpoint

    main

    Create an online endpoint for inference using the configuration in azure_ml/endpoint.yml.

    Important: After creation, you must go to the Azure Portal, open your Container Registry, and grant ACR_PULL permissions to the endpoint to allow it to download the Docker image.

    az ml online-endpoint create -f azure_ml/endpoint.yml -g $RESOURCE_GROUP -w $WORKSPACE_NAME

  8. Push the Triton Docker image to Azure Container Registry

    main

    Log in to your Azure Container Registry (ACR) and push the indictrans2_triton image using the specified tag format.

    az acr login --name $DOCKER_REGISTRY
docker tag indictrans2_triton $DOCKER_REGISTRY.azurecr.io/nmt/triton-indictrans-v2:latest
docker push $DOCKER_REGISTRY.azurecr.io/nmt/triton-indictrans-v2:latest

  9. Fine-tune IndicTrans2 models with LoRA

    main

    To fine-tune models using LoRA, you must first structure your training and development data in a specific directory format.

    Data Directory Structure:

    en-indic-exp
├── train
│   ├── {src_lang}-{tgt_lang}
│   │   ├── train.{src_lang}
│   │   └── train.{tgt_lang}
│   └── ...
└── dev
    ├── {src_lang}-{tgt_lang}
    │   ├── dev.{src_lang}
    │   └── dev.{tgt_lang}
    └── ...

    Training Command: Use the train_lora.sh script to start fine-tuning. It is recommended to check the script for default hyperparameter arguments.

    bash train_lora.sh <data_dir> <model_name> <output_dir> <direction> <src_lang_list> <tgt_lang_list>

  10. Train or Fine-tune IndicTrans2 models

    main

    Once your data is binarized, you can train a new model or fine-tune an existing one using the provided shell scripts.

    Training a new model

    Use train.sh to start training. You can provide a custom transformer architecture by defining it in model_configs/custom_transformer.py.

    Arguments:

    • <exp_dir>: path to the binarized data.
    • <model_arch>: custom transformer architecture used for training.

    Fine-tuning a model

    Use finetune.sh to perform fine-tuning. This requires a pretrained checkpoint.

    Arguments:

    • <exp_dir>: path to the binarized data.
    • <model_arch>: architecture type. Use transformer_18_18 for IT2 Base models or transformer_base18L for IT2 Distilled models.
    • <pretrained_ckpt>: path to the fairseq model checkpoint to be loaded.

    Note on Checkpoints: When downloading artifacts, use the fairseq_model directory for fine-tuning. The CT2 directories are intended for efficient inference and should be generated via the fairseq-ct2-converter after fine-tuning is complete.

    # Training
bash train.sh <exp_dir> <model_arch>

# Fine-tuning
bash finetune.sh <exp_dir> <model_arch> <pretrained_ckpt>

  11. Install IndicTrans2

    main

    To set up IndicTrans2, clone the repository and run the provided installation script. It is recommended to use a virtual environment with Python version 3.7 or higher.

    # Clone the repository and enter the directory
git clone https://github.com/AI4Bharat/IndicTrans2
cd IndicTrans2

# Install dependencies and requirements
source install.sh
    git clone https://github.com/AI4Bharat/IndicTrans2
cd IndicTrans2

source install.sh

  12. Set up environment for Azure Machine Learning deployment

    main

    Before deploying IndicTrans2 on Azure Machine Learning (AML), navigate to the Triton server directory and configure the required environment variables. You must also ensure that the provided .yml files in the azure_ml/ directory are edited to match your specific Azure configuration.

    cd inference/triton_server

export RESOURCE_GROUP=Dhruva-prod
export WORKSPACE_NAME=dhruva--central-india
export DOCKER_REGISTRY=dhruvaprod