Upload 79 files

IndicTrans2/.gitignore

@@ -0,0 +1,148 @@
+# ignore libs and data folder we use
+indic_nlp_library
+indic_nlp_resources
+fairseq
+devtest
+checkpoints
+eval_benchmarks
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+.DS_Store

IndicTrans2/LICENSE

@@ -0,0 +1,21 @@
+    MIT License
+
+    Copyright (c) AI4Bharat.
+
+    Permission is hereby granted, free of charge, to any person obtaining a copy
+    of this software and associated documentation files (the "Software"), to deal
+    in the Software without restriction, including without limitation the rights
+    to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+    copies of the Software, and to permit persons to whom the Software is
+    furnished to do so, subject to the following conditions:
+
+    The above copyright notice and this permission notice shall be included in all
+    copies or substantial portions of the Software.
+
+    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+    FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+    AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+    LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+    OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+    SOFTWARE

IndicTrans2/README.md

@@ -0,0 +1,523 @@
+# IndicTrans2
+
+[📜 Paper](https://arxiv.org/abs/2305.16307) | [🌐 Website](https://ai4bharat.iitm.ac.in/indic-trans2) | [▶️ Demo](https://models.ai4bharat.org/#/nmt/v2) | [🤗 HF Interface](https://github.com/AI4Bharat/IndicTrans2/tree/main/huggingface_interface) | [![colab link](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/AI4Bharat/IndicTrans2/blob/main/huggingface_interface/colab_inference.ipynb)
+
+IndicTrans2 is the first open-source transformer-based multilingual NMT model that supports high-quality translations across all the 22 scheduled Indic languages — including multiple scripts for low-resouce languages like Kashmiri, Manipuri and Sindhi. It adopts script unification wherever feasible to leverage transfer learning by lexical sharing between languages. Overall, the model supports five scripts Perso-Arabic (Kashmiri, Sindhi, Urdu), Ol Chiki (Santali), Meitei (Manipuri), Latin (English), and Devanagari (used for all the remaining languages).
+
+We open-souce all our training dataset (BPCC), back-translation data (BPCC-BT), final IndicTrans2 models, evaluation benchmarks (IN22, which includes IN22-Gen and IN22-Conv) and training and inference scripts for easier use and adoption within the research community. We hope that this will foster even more research in low-resource Indic languages, leading to further improvements in the quality of low-resource translation through contributions from the research community.
+
+This code repository contains instructions for downloading the artifacts associated with IndicTrans2, as well as the code for training/fine-tuning the multilingual NMT models.
+
+Here is the list of languages supported by the IndicTrans2 models:
+
+<table>
+<tbody>
+  <tr>
+    <td>Assamese (asm_Beng)</td>
+    <td>Kashmiri (Arabic) (kas_Arab)</td>
+    <td>Punjabi (pan_Guru)</td>
+  </tr>
+  <tr>
+    <td>Bengali (ben_Beng)</td>
+    <td>Kashmiri (Devanagari) (kas_Deva)</td>
+    <td>Sanskrit (san_Deva)</td>
+  </tr>
+  <tr>
+    <td>Bodo (brx_Deva)</td>
+    <td>Maithili (mai_Deva)</td>
+    <td>Santali (sat_Olck)</td>
+  </tr>
+  <tr>
+    <td>Dogri (doi_Deva)</td>
+    <td>Malayalam (mal_Mlym)</td>
+    <td>Sindhi (Arabic) (snd_Arab)</td>
+  </tr>
+  <tr>
+    <td>English (eng_Latn)</td>
+    <td>Marathi (mar_Deva)</td>
+    <td>Sindhi (Devanagari) (snd_Deva)</td>
+  </tr>
+  <tr>
+    <td>Konkani (gom_Deva)</td>
+    <td>Manipuri (Bengali) (mni_Beng)</td>
+    <td>Tamil (tam_Taml)</td>
+  </tr>
+  <tr>
+    <td>Gujarati (guj_Gujr)</td>
+    <td>Manipuri (Meitei) (mni_Mtei)</td>
+    <td>Telugu (tel_Telu)</td>
+  </tr>
+  <tr>
+    <td>Hindi (hin_Deva)</td>
+    <td>Nepali (npi_Deva)</td>
+    <td>Urdu (urd_Arab)</td>
+  </tr>
+  <tr>
+    <td>Kannada (kan_Knda)</td>
+    <td>Odia (ory_Orya)</td>
+    <td></td>
+  </tr>
+</tbody>
+</table>
+
+## Updates
+
+- 🚨 Dec 30, 2023 - Migrated IndicTrans2 tokenizer for HF compatible IndicTrans2 models to [IndicTransTokenizer](https://github.com/VarunGumma/IndicTransTokenizer) and will be maintained separately there from now onwards. Add LoRA fine-tuning scripts for our IndicTrans2 models in [huggingface_interface](https://github.com/AI4Bharat/IndicTrans2/tree/main/huggingface_interface).
+- 🚨 Dec 1, 2023 - Release of Indic-Indic model and corresponding distilled variants for each base model. Please refer to the [Download section](https://github.com/AI4Bharat/IndicTrans2#multilingual-translation-models) for the checkpoints.
+- 🚨 Sep 9, 2023 - Added HF compatible IndicTrans2 models. Please refer to the [README](https://github.com/AI4Bharat/IndicTrans2/tree/main/huggingface_interface) for detailed example usage.
+
+## Tables of Contents
+
+- [Download Models and Other Artifacts](#download-models-and-other-artifacts)
+  - [Multilingual Translation Models](#multilingual-translation-models)
+  - [Training Data](#training-data)
+  - [Evaluation Data](#evaluation-data)
+- [Installation](#installation)
+- [Data](#data)
+  - [Training](#training)
+  - [Evaluation](#evaluation)
+- [Preparing Data for Training](#preparing-data-for-training)
+  - [Using our SPM model and Fairseq dictionary](#using-our-spm-model-and-fairseq-dictionary)
+  - [Training your own SPM models and learning Fairseq dictionary](#training-your-own-spm-models-and-learning-fairseq-dictionary)
+- [Training / Fine-tuning](#training--fine-tuning)
+- [Inference](#inference)
+  - [Fairseq Inference](#fairseq-inference)
+  - [CT2 Inference](#ct2-inference)
+- [Evaluations](#evaluations)
+  - [Baseline Evaluation](#baseline-evaluation)
+- [LICENSE](#license)
+- [Citation](#citation)
+
+## Download Models and Other Artifacts
+
+### Multilingual Translation Models
+
+| Model                        | En-Indic                                                                                                    | Indic-En                                                                                                    | Indic-Indic                                                                                            | Evaluations                                                                                                                                                                                                          |
+| ---------------------------- | ----------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------ | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| Base (used for benchmarking) | [download](https://indictrans2-public.objectstore.e2enetworks.net/it2_preprint_ckpts/en-indic-preprint.zip) | [download](https://indictrans2-public.objectstore.e2enetworks.net/it2_preprint_ckpts/indic-en-preprint.zip) | [download](https://indictrans2-public.objectstore.e2enetworks.net/it2_preprint_ckpts/indic-indic.zip)  | [translations](https://indictrans2-public.objectstore.e2enetworks.net/translation_outputs.zip) (as of May 10, 2023), [metrics](https://drive.google.com/drive/folders/1lOOdaU0VdRSBgJEsNav5zC7wwLBis9NI?usp=sharing) |
+| Distilled                    | [download](https://indictrans2-public.objectstore.e2enetworks.net/it2_distilled_ckpts/en-indic.zip)         | [download](https://indictrans2-public.objectstore.e2enetworks.net/it2_distilled_ckpts/indic-en.zip)         | [download](https://indictrans2-public.objectstore.e2enetworks.net/it2_distilled_ckpts/indic-indic.zip) |
+
+### Training Data
+
+| Data                                     | URL                                                                            |
+| ---------------------------------------- | ------------------------------------------------------------------------------ |
+| Bharat Parallel Corpus Collection (BPCC) | [download](https://indictrans2-public.objectstore.e2enetworks.net/BPCC.zip)    |
+| Back-translation (BPCC-BT)               | [download](https://indictrans2-public.objectstore.e2enetworks.net/BT_data.zip) |
+
+### Evaluation Data
+
+| Data                    | URL                                                                                  |
+| ----------------------- | ------------------------------------------------------------------------------------ |
+| IN22 test set           | [download](https://indictrans2-public.objectstore.e2enetworks.net/IN22_testset.zip)  |
+| FLORES-22 Indic dev set | [download](https://indictrans2-public.objectstore.e2enetworks.net/flores-22_dev.zip) |
+
+## Installation
+
+Instructions to setup and install everything before running the code.
+
+```bash
+# Clone the github repository and navigate to the project directory.
+git clone https://github.com/AI4Bharat/IndicTrans2
+cd IndicTrans2
+
+# Install all the dependencies and requirements associated with the project.
+source install.sh
+```
+
+Note: We recommend creating a virtual environment with python>=3.7.
+
+### Additional notes about Installation
+The ``prepare_data_joint_finetuning.sh`` and ``prepare_data_joint_training.sh`` scripts expect that the sentencepiece commandline utility and GNU parallel are installed.
+1. To install the sentencepiece command line utility, please follow the instructions [here](https://github.com/google/sentencepiece?tab=readme-ov-file#build-and-install-sentencepiece-command-line-tools-from-c-source).
+2. Please check if GNU parallel is installed, if not please install the same or alternatively in case of installation issues, remove ``parallel --pipe --keep-order`` from the respective training / finetuning script as well as ``apply_sentence_piece.sh``.
+
+
+## Data
+
+### Training
+
+Bharat Parallel Corpus Collection (BPCC) is a comprehensive and publicly available parallel corpus that includes both existing and new data for all 22 scheduled Indic languages. It is comprised of two parts: BPCC-Mined and BPCC-Human, totaling approximately 230 million bitext pairs. BPCC-Mined contains about 228 million pairs, with nearly 126 million pairs newly added as a part of this work. On the other hand, BPCC-Human consists of 2.2 million gold standard English-Indic pairs, with an additional 644K bitext pairs from English Wikipedia sentences (forming the BPCC-H-Wiki subset) and 139K sentences covering everyday use cases (forming the BPCC-H-Daily subset). It is worth highlighting that BPCC provides the first available datasets for 7 languages and significantly increases the available data for all languages covered.
+
+You can find the contribution from different sources in the following table:
+
+<table>
+<tbody>
+  <tr>
+    <td rowspan="4">BPCC-Mined</th>
+    <td rowspan="2">Existing</th>
+    <td>Samanantar</th>
+    <td>19.4M</th>
+  </tr>
+  <tr>
+    <td>NLLB</th>
+    <td>85M</th>
+  </tr>
+  <tr>
+    <td rowspan="2">Newly Added</th>
+    <td>Samanantar++</th>
+    <td>121.6M</th>
+  </tr>
+  <tr>
+    <td>Comparable</th>
+    <td>4.3M</th>
+  </tr>
+  <tr>
+    <td rowspan="5">BPCC-Human</td>
+    <td rowspan="3">Existing</td>
+    <td>NLLB</td>
+    <td>18.5K</td>
+  </tr>
+  <tr>
+    <td>ILCI</td>
+    <td>1.3M</td>
+  </tr>
+  <tr>
+    <td>Massive</td>
+    <td>115K</td>
+  </tr>
+  <tr>
+    <td rowspan="2">Newly Added</td>
+    <td>Wiki</td>
+    <td>644K</td>
+  </tr>
+  <tr>
+    <td>Daily</td>
+    <td>139K</td>
+  </tr>
+</tbody>
+</table>
+
+Additionally, we provide augmented back-translation data generated by our intermediate IndicTrans2 models for training purposes. Please refer our paper for more details on the selection of sample proportions and sources.
+
+<table>
+<tbody>
+  <tr>
+    <td>English BT data (English Original)</td>
+    <td>401.9M</td>
+  </tr>
+  <tr>
+    <td>Indic BT data (Indic Original)</td>
+    <td>400.9M</td>
+  </tr>
+</tbody>
+</table>
+
+<br>
+
+### Evaluation
+
+IN22 test set is a newly created comprehensive benchmark for evaluating machine translation performance in multi-domain, n-way parallel contexts across 22 Indic languages. It has been created from three distinct subsets, namely IN22-Wiki, IN22-Web and IN22-Conv. The Wikipedia and Web sources subsets offer diverse content spanning news, entertainment, culture, legal, and India-centric topics. IN22-Wiki and IN22-Web have been combined and considered for evaluation purposes and released as IN22-Gen. Meanwhile, IN22-Conv the conversation domain subset is designed to assess translation quality in typical day-to-day conversational-style applications.
+
+<table>
+<tbody>
+  <tr>
+    <td>IN22-Gen (IN22-Wiki + IN22-Web)</td>
+    <td>1024 sentences</td>
+    <td>🤗 <a href="https://huggingface.co/datasets/ai4bharat/IN22-Gen">ai4bharat/IN22-Gen</td>
+  </tr>
+  <tr>
+    <td>IN22-Conv</td>
+    <td>1503 sentences</td>
+    <td>🤗 <a href="https://huggingface.co/datasets/ai4bharat/IN22-Conv">ai4bharat/IN22-Conv</td>
+  </tr>
+</tbody>
+</table>
+
+You can download the data artifacts released as a part of this work from the [following section](#download-models-and-other-artifacts).
+
+## Preparing Data for Training
+
+BPCC data is organized under different subsets as described above, where each subset contains language pair subdirectories with the sentences pairs. We also provide LaBSE and LASER for the mined subsets of BPCC. In order to replicate our training setup, you will need to combine the data for corresponding language pairs from different subsets and remove overlapping bitext pairs if any.
+
+Here is the expected directory structure of the data:
+
+```bash
+BPCC
+├── eng_Latn-asm_Beng
+│   ├── train.eng_Latn
+│   └── train.asm_Beng
+├── eng_Latn-ben_Beng
+└── ...
+```
+
+While we provide deduplicated subsets with the current available benchmarks, we highly recommend performing deduplication using the combined monolingual side of all the benchmarks. You can use the following command for deduplication once you combine the monolingual side of all the benchmarks in the directory.
+
+```python3
+python3 scripts/dedup_benchmark.py <in_data_dir> <out_data_dir> <benchmark_dir>
+```
+
+- `<in_data_dir>`: path to the directory containing train data for each language pair in the format `{src_lang}-{tgt_lang}`
+- `<out_data_dir>`: path to the directory where the deduplicated train data will be written for each language pair in the format `{src_lang}-{tgt_lang}`
+- `<benchmark_dir>`: path to the directory containing the language-wise monolingual side of dev/test set, with monolingual files named as `test.{lang}`
+
+### Using our SPM model and Fairseq dictionary
+
+Once you complete the deduplication of the training data with the available benchmarks, you can preprocess and binarize the data for training models. Please download our trained SPM model and learned Fairseq dictionary using the following links for your experiments.
+
+|                    | En-Indic                                                                                     | Indic-En                                                                                     | Indic-Indic                                                                                     |
+| ------------------ | -------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------- |
+| SPM model          | [download](https://indictrans2-public.objectstore.e2enetworks.net/en-indic-spm.zip)          | [download](https://indictrans2-public.objectstore.e2enetworks.net/indic-en-spm.zip)          | [download](https://indictrans2-public.objectstore.e2enetworks.net/indic-indic-spm.zip)          |
+| Fairseq dictionary | [download](https://indictrans2-public.objectstore.e2enetworks.net/en-indic-fairseq-dict.zip) | [download](https://indictrans2-public.objectstore.e2enetworks.net/indic-en-fairseq-dict.zip) | [download](https://indictrans2-public.objectstore.e2enetworks.net/indic-indic-fairseq-dict.zip) |
+
+To prepare the data for training En-Indic model, please do the following:
+
+1. Download the SPM model in the experiment directory and rename it as `vocab`.
+2. Download the Fairseq dictionary in the experiment directory and rename it as `final_bin`.
+
+Here is the expected directory for training En-Indic model:
+
+```bash
+en-indic-exp
+├── train
+│   ├── eng_Latn-asm_Beng
+│   │   ├── train.eng_Latn
+│   │   └── train.asm_Beng
+│   ├── eng_Latn-ben_Beng
+│   └── ...
+├── devtest
+│   └── all
+│       ├── eng_Latn-asm_Beng
+│       │   ├── dev.eng_Latn
+│       │   └── dev.asm_Beng
+│       ├── eng_Latn-ben_Beng
+│       └── ...
+├── vocab
+│   ├── model.SRC
+│   ├── model.TGT
+│   ├── vocab.SRC
+│   └── vocab.TGT
+└── final_bin
+    ├── dict.SRC.txt
+    └── dict.TGT.txt
+```
+
+To prepare data for training the Indic-En model, you should reverse the language pair directories within the train and devtest directories. Additionally, make sure to download the corresponding SPM model and Fairseq dictionary and put them in the experiment directory, similar to the procedure mentioned above for En-Indic model training.
+
+You can binarize the data for model training using the following:
+
+```bash
+bash prepare_data_joint_finetuning.sh <exp_dir>
+```
+
+- `<exp_dir>`: path to the directory containing the raw data for binarization
+
+You will need to follow the same steps for data preparation in case of fine-tuning models.
+
+### Training your own SPM models and learning Fairseq dictionary
+
+If you want to train your own SPM model and learn Fairseq dictionary, then please do the following:
+
+1. Collect a balanced amount of English and Indic monolingual data (we use around 3 million sentences per language-script combination). If some languages have limited data available, increase their representation to achieve a fair distribution of tokens across languages.
+2. Perform script unification for Indic languages wherever possible using `scripts/preprocess_translate.py` and concatenate all Indic data into a single file.
+3. Train two SPM models, one for English and other for Indic side using the following:
+
+```bash
+spm_train --input=train.indic --model_prefix=<model_name> --vocab_size=<vocab_size> --character_coverage=1.0 --model_type=BPE
+```
+
+4. Copy the trained SPM models in the experiment directory mentioned earlier and learn the Fairseq dictionary using the following:
+
+```bash
+bash prepare_data_joint_training.sh <exp_dir>
+```
+
+5. You will need to use the same Fairseq dictionary for any subsequent fine-tuning experiments and refer to the steps described above ([link](#using-our-spm-model-and-fairseq-dictionary)).
+
+## Training / Fine-tuning
+
+After binarizing the data, you can use train.sh to train the models. We provide the default hyperparameters used in this work. You can modify the hyperparameters as per your requirement if needed. If you want to train the model on a customized architecture, then please define the architecture in `model_configs/custom_transformer.py`. You can start the model training with the following command:
+
+```bash
+bash train.sh <exp_dir> <model_arch>
+```
+
+- `<exp_dir>`: path to the directory containing the binarized data
+- `<model_arch>`: custom transformer architecture used for model training
+
+For fine-tuning, the initial steps remain the same. However, the `finetune.sh` script includes an additional argument, `pretrained_ckpt`, which specifies the model checkpoint to be loaded for further fine-tuning. You can perform fine-tuning using the following command:
+
+```bash
+bash finetune.sh <exp_dir> <model_arch> <pretrained_ckpt>
+```
+
+- `<exp_dir>`: path to the directory containing the binarized data
+- `<model_arch>`: custom transformer architecture used for model training
+  - `transformer_18_18` - For IT2 Base models
+  - `transformer_base18L` - For IT2 Distilled models
+- `<pretrained_ckpt>`: path to the fairseq model checkpoint to be loaded for further fine-tuning
+
+You can download the model artifacts released as a part of this work from the [following section](#download-models-and-other-artifacts).
+
+The pretrained checkpoints have 3 directories, a fairseq model directory and 2 CT-ported model directories. Please note that the CT2 models are provided only for efficient inference. For fine-tuning purposes you should use the `fairseq_model`. Post that you can use the [fairseq-ct2-converter](https://opennmt.net/CTranslate2/guides/fairseq.html) to port your fine-tuned checkpoints to CT2 for faster inference.
+
+## Inference
+
+### Fairseq Inference
+
+In order to run inference on our pretrained models using bash interface, please use the following:
+
+```bash
+bash joint_translate.sh <infname> <outfname> <src_lang> <tgt_lang> <ckpt_dir>
+```
+
+- `infname`: path to the input file containing sentences
+- `outfname`: path to the output file where the translations should be stored
+- `src_lang`: source language
+- `tgt_lang`: target language
+- `ckpt_dir`: path to the fairseq model checkpoint directory
+
+If you want to run the inference using python interface then please execute the following block of code from the root directory:
+
+```python3
+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)
+```
+
+### CT2 Inference
+
+In order to run inference on CT2-ported model using python inference then please execute the following block of code from the root directory:
+
+```python3
+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)
+```
+
+## Evaluations
+
+We consider the chrF++ as our primary metric. Additionally, we also report the BLEU and Comet scores.
+We also perform statistical significance tests for each metric to ascertain whether the differences are statistically significant.
+
+In order to run our evaluation scripts, you will need to organize the evaluation test sets into the following directory structure:
+
+```bash
+eval_benchmarks
+├── flores
+│   └── eng_Latn-asm_Beng
+│       ├── test.eng_Latn
+│       └── test.asm_Beng
+├── in22-gen
+├── in22-conv
+├── ntrex
+└── ...
+```
+
+To compute the BLEU and chrF++ scores for prediction file, you can use the following command:
+
+```bash
+bash compute_metrics.sh <pred_fname> <ref_fname> <tgt_lang>
+```
+
+- `pred_fname`: path to the model translations
+- `ref_fname`: path to the reference translations
+- `tgt_lang`: target language
+
+In order to automate the inference over the individual test sets for En-Indic, you can use the following command:
+
+```bash
+bash eval.sh <devtest_data_dir> <ckpt_dir> <system>
+```
+
+- `<devtest_data_dir>`: path to the evaluation set with language pair subdirectories (for example, flores directory in the above tree structure)
+- `<ckpt_dir>`: path to the fairseq model checkpoint directory
+- `<system>`: system name suffix to store the predictions in the format `test.{lang}.pred.{system}`
+
+In case of Indic-En evaluation, please use the following command:
+
+```bash
+bash eval_rev.sh  <devtest_data_dir> <ckpt_dir> <system>
+```
+
+- `<devtest_data_dir>`: path to the evaluation set with language pair subdirectories (for example, flores directory in the above tree structure)
+- `<ckpt_dir>`: path to the fairseq model checkpoint directory
+- `<system>`: system name suffix to store the predictions in the format `test.{lang}.pred.{system}`
+
+**_Note: You don’t need to reverse the test set directions for each language pair._**
+
+In case of Indic-Indic evaluation, please use the following command:
+
+```bash
+bash pivot_eval.sh <devtest_data_dir> <pivot_lang> <src2pivot_ckpt_dir> <pivot2tgt_ckpt_dir> <system>
+```
+
+- `<devtest_data_dir>`: path to the evaluation set with language pair subdirectories (for example, flores directory in the above tree structure)
+- `<pivot_lang>`: pivot language (default should be `eng_Latn`)
+- `<src2pivot_ckpt_dir>`: path to the fairseq Indic-En model checkpoint directory
+- `<pivot2tgt_ckpt_dir>`: path to the fairseq En-Indic model checkpoint directory
+- `<system>`: system name suffix to store the predictions in the format test.{lang}.pred.{system}
+
+In order to perform significance testing for BLEU and chrF++ metrics after you have the predictions for different systems, you can use the following command:
+
+```bash
+bash compute_comet_metrics_significance.sh <devtest_data_dir>
+```
+
+- `<devtest_data_dir>`: path to the evaluation set with language pair subdirectories (for example, flores directory in the above tree structure)
+
+Similarly, to compute the COMET scores and perform significance testing on predictions of different systems, you can use the following command.
+
+```bash
+bash compute_comet_score.sh <devtest_data_dir>
+```
+
+- `<devtest_data_dir>`: path to the evaluation set with language pair subdirectories (for example, flores directory in the above tree structure)
+
+Please note that as we compute significance tests with the same script and automate everything, it is best to have all the predictions for all the systems in place to avoid repeating anything.
+Also, we define the systems in the script itself, if you want to try out other systems, make sure to edit it there itself.
+
+### Baseline Evaluation
+
+To generate the translation results for baseline models such as M2M-100, MBART, Azure, Google, and NLLB MoE, you can check the scripts provided in the "baseline_eval" directory of this repository. For NLLB distilled, you can either modify NLLB_MoE eval or use this [repository](https://github.com/pluiez/NLLB-inference). Similarly, for IndicTrans inference, please refer to this [repository](https://github.com/ai4bharat/IndicTrans).
+
+You can download the translation outputs released as a part of this work from the [following section](#download-models-and-other-artifacts).
+
+## LICENSE
+
+The following table lists the licenses associated with the different artifacts released as a part of this work:
+
+| Artifact                                              | LICENSE                                                               |
+| ----------------------------------------------------- | --------------------------------------------------------------------- |
+| Existing Mined Corpora (NLLB & Samanantar)            | [CC0](https://creativecommons.org/share-your-work/public-domain/cc0/) |
+| Existing Seed Corpora (NLLB-Seed, ILCI, MASSIVE)      | [CC0](https://creativecommons.org/share-your-work/public-domain/cc0/) |
+| Newly Added Mined Corpora (Samanantar++ & Comparable) | [CC0](https://creativecommons.org/share-your-work/public-domain/cc0/) |
+| Newly Added Seed Corpora (BPCC-H-Wiki & BPCC-H-Daily) | [CC-BY-4.0](https://creativecommons.org/licenses/by/4.0/)             |
+| Newly Created IN-22 test set (IN22-Gen & IN22-Conv)   | [CC-BY-4.0](https://creativecommons.org/licenses/by/4.0/)             |
+| Back-translation data (BPCC-BT)                       | [CC0](https://creativecommons.org/share-your-work/public-domain/cc0/) |
+| Model checkpoints                                     | [MIT](https://github.com/ai4bharat/IndicTrans2/blob/main/LICENSE)     |
+
+The mined corpora collection (BPCC-Mined), existing seed corpora (NLLB-Seed, ILCI, MASSIVE), Backtranslation data (BPCC-BT), are released under the following licensing scheme:
+
+- We do not own any of the text from which this data has been extracted.
+- We license the actual packaging of this data under the Creative Commons [CC0 license (“no rights reserved”)](https://creativecommons.org/share-your-work/public-domain/cc0/).
+- To the extent possible under law, [AI4Bharat](https://ai4bharat.iitm.ac.in/) has waived all copyright and related or neighboring rights to BPCC-Mined, existing seed corpora (NLLB-Seed, ILCI, MASSIVE) and BPCC-BT.
+
+## Citation
+
+```bibtex
+@article{gala2023indictrans,
+title={IndicTrans2: Towards High-Quality and Accessible Machine Translation Models for all 22 Scheduled Indian Languages},
+author={Jay Gala and Pranjal A Chitale and A K Raghavan and Varun Gumma and Sumanth Doddapaneni and Aswanth Kumar M and Janki Atul 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},
+issn={2835-8856},
+year={2023},
+url={https://openreview.net/forum?id=vfT4YuzAYA},
+note={}
+}
+```

IndicTrans2/apply_sentence_piece.sh

@@ -0,0 +1,48 @@
+#!/bin/bash
+
+# This script tokenizes the preprocessed train and dev set using the trained spm models.
+
+
+echo `date`
+exp_dir=$1                      # path to the experiment directory
+data_dir=$2                     # path to the data directory where all lang pairs are concatenated
+bpe_dir=$3                      # path to the tokenized data directory
+src_lang=$4                     # source language
+tgt_lang=$5                     # target language
+split=$6                        # name of the split
+parallel_installed=${7:-false}  # If GNU Parallel is installed or not
+
+in_split_dir=$data_dir/$split
+out_split_dir=$bpe_dir/$split
+
+echo "Apply Sentence Piece tokenization to SRC corpus"
+# for very large datasets, it is recommended to use gnu-parallel to speed up applying bpe
+
+if $parallel_installed; then
+    parallel --pipe --keep-order \
+    spm_encode --model=$exp_dir/vocab/model.SRC \
+        --output_format=piece \
+        < $in_split_dir.$src_lang \
+        > $out_split_dir.$src_lang
+else
+    spm_encode --model=$exp_dir/vocab/model.SRC \
+        --output_format=piece \
+        < $in_split_dir.$src_lang \
+        > $out_split_dir.$src_lang
+fi
+
+echo "Apply Sentence Piece tokenization to TGT corpus"
+# for very large datasets, it is recommended to use gnu-parallel to speed up applying bpe
+
+if $parallel_installed; then
+    parallel --pipe --keep-order \
+    spm_encode --model=$exp_dir/vocab/model.TGT \
+        --output_format=piece \
+        < $in_split_dir.$tgt_lang \
+        > $out_split_dir.$tgt_lang
+else
+    spm_encode --model=$exp_dir/vocab/model.TGT \
+        --output_format=piece \
+        < $in_split_dir.$tgt_lang \
+        > $out_split_dir.$tgt_lang
+fi

IndicTrans2/baseline_eval/azure_translate.py

@@ -0,0 +1,183 @@
+import os
+import sys
+import glob
+import requests
+from urllib.parse import urlencode
+from dotenv import dotenv_values
+import traceback
+import time
+
+flores_to_iso = {
+    "asm_Beng": "as",
+    "ben_Beng": "bn",
+    "brx_Deva": "brx",
+    "doi_Deva": "doi",
+    "eng_Latn": "en",
+    "gom_Deva": "gom",
+    "guj_Gujr": "gu",
+    "hin_Deva": "hi",
+    "kan_Knda": "kn",
+    "kas_Arab": "ks",
+    "kas_Deva": "ks_Deva",
+    "mai_Deva": "mai",
+    "mal_Mlym": "ml",
+    "mar_Deva": "mr",
+    "mni_Beng": "mni_Beng",
+    "mni_Mtei": "mni",
+    "npi_Deva": "ne",
+    "ory_Orya": "or",
+    "pan_Guru": "pa",
+    "san_Deva": "sa",
+    "sat_Olck": "sat",
+    "snd_Arab": "sd",
+    "snd_Deva": "sd_Deva",
+    "tam_Taml": "ta",
+    "tel_Telu": "te",
+    "urd_Arab": "ur",
+}
+
+
+class AzureTranslator:
+    def __init__(
+        self,
+        subscription_key: str,
+        region: str,
+        endpoint: str = "https://api.cognitive.microsofttranslator.com",
+    ) -> None:
+        self.http_headers = {
+            "Ocp-Apim-Subscription-Key": subscription_key,
+            "Ocp-Apim-Subscription-Region": region,
+        }
+        self.translate_endpoint = endpoint + "/translate?api-version=3.0&"
+        self.languages_endpoint = endpoint + "/languages?api-version=3.0"
+
+        self.supported_languages = self.get_supported_languages()
+
+    def get_supported_languages(self) -> dict:
+        return requests.get(self.languages_endpoint).json()["translation"]
+
+    def batch_translate(self, texts: list, src_lang: str, tgt_lang: str) -> list:
+        if not texts:
+            return texts
+
+        src_lang = flores_to_iso[src_lang]
+        tgt_lang = flores_to_iso[tgt_lang]
+
+        if src_lang not in self.supported_languages:
+            raise NotImplementedError(
+                f"Source language code: `{src_lang}` not supported!"
+            )
+
+        if tgt_lang not in self.supported_languages:
+            raise NotImplementedError(
+                f"Target language code: `{tgt_lang}` not supported!"
+            )
+
+        body = [{"text": text} for text in texts]
+        query_string = urlencode(
+            {
+                "from": src_lang,
+                "to": tgt_lang,
+            }
+        )
+
+        try:
+            response = requests.post(
+                self.translate_endpoint + query_string,
+                headers=self.http_headers,
+                json=body,
+            )
+        except:
+            traceback.print_exc()
+            return None
+
+        try:
+            response = response.json()
+        except:
+            traceback.print_exc()
+            print("Response:", response.text)
+            return None
+
+        return [payload["translations"][0]["text"] for payload in response]
+
+    def text_translate(self, text: str, src_lang: str, tgt_lang: str) -> str:
+        return self.batch_translate([text], src_lang, tgt_lang)[0]
+
+
+if __name__ == "__main__":
+    root_dir = sys.argv[1]
+    
+    # Expects a .env file containing the API credentials.
+    config = dotenv_values(os.path.join(os.path.dirname(__file__), ".env"))
+
+    t = AzureTranslator(
+        config["AZURE_TRANSLATOR_TEXT_SUBSCRIPTION_KEY"],
+        config["AZURE_TRANSLATOR_TEXT_REGION"],
+        config["AZURE_TRANSLATOR_TEXT_ENDPOINT"],
+    )
+
+    pairs = sorted(glob.glob(os.path.join(root_dir, "*")))
+
+    for i, pair in enumerate(pairs):
+        basename = os.path.basename(pair)
+        
+        print(pair)
+    
+        src_lang, tgt_lang = basename.split("-")
+
+        print(f"{src_lang} - {tgt_lang}")
+
+        # source to target translations
+        src_infname = os.path.join(pair, f"test.{src_lang}")
+        tgt_outfname = os.path.join(pair, f"test.{tgt_lang}.pred.azure")
+        if not os.path.exists(src_infname):
+            continue
+
+        src_sents = [
+            sent.replace("\n", "").strip()
+            for sent in open(src_infname, "r").read().split("\n")
+            if sent
+        ]
+
+        if not os.path.exists(tgt_outfname):
+            try:
+                translations = []
+                for i in range(0, len(src_sents), 128):
+                    start, end = i, int(min(i + 128, len(src_sents)))
+                    translations.extend(
+                        t.batch_translate(src_sents[start:end], src_lang, tgt_lang)
+                    )
+                with open(tgt_outfname, "w") as f:
+                    f.write("\n".join(translations))
+                
+                time.sleep(10)
+            except Exception as e:
+                print(e)
+                continue
+
+        # target to source translations
+        tgt_infname = os.path.join(pair, f"test.{tgt_lang}")
+        src_outfname = os.path.join(pair, f"test.{src_lang}.pred.azure")
+        if not os.path.exists(tgt_infname):
+            continue
+
+        tgt_sents = [
+            sent.replace("\n", "").strip()
+            for sent in open(tgt_infname, "r").read().split("\n")
+            if sent
+        ]
+
+        if not os.path.exists(src_outfname):
+            try:
+                translations = []
+                for i in range(0, len(tgt_sents), 128):
+                    start, end = i, int(min(i + 128, len(tgt_sents)))
+                    translations.extend(
+                        t.batch_translate(tgt_sents[start:end], tgt_lang, src_lang)
+                    )
+                with open(src_outfname, "w") as f:
+                    f.write("\n".join(translations))
+            except Exception as e:
+                continue
+
+            time.sleep(10)

IndicTrans2/baseline_eval/google_translate.py

@@ -0,0 +1,129 @@
+import os
+import sys
+import glob
+from tqdm import tqdm
+from google.cloud import translate
+
+# Expects a json file containing the API credentials.
+os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = os.path.join(
+    os.path.dirname(__file__), r"api_key.json"
+)
+
+flores_to_iso = {
+    "asm_Beng": "as",
+    "ben_Beng": "bn",
+    "doi_Deva": "doi",
+    "eng_Latn": "en",
+    "gom_Deva": "gom",
+    "guj_Gujr": "gu",
+    "hin_Deva": "hi",
+    "kan_Knda": "kn",
+    "mai_Deva": "mai",
+    "mal_Mlym": "ml",
+    "mar_Deva": "mr",
+    "mni_Mtei": "mni_Mtei",
+    "npi_Deva": "ne",
+    "ory_Orya": "or",
+    "pan_Guru": "pa",
+    "san_Deva": "sa",
+    "sat_Olck": "sat",
+    "snd_Arab": "sd",
+    "tam_Taml": "ta",
+    "tel_Telu": "te",
+    "urd_Arab": "ur",
+}
+
+
+# Copy the project id from the json file containing API credentials
+def translate_text(text, src_lang, tgt_lang, project_id="project_id"):
+
+    src_lang = flores_to_iso[src_lang]
+    tgt_lang = flores_to_iso[tgt_lang]
+
+    if src_lang == "mni_Mtei":
+        src_lang = "mni-Mtei"
+
+    if tgt_lang == "mni_Mtei":
+        tgt_lang = "mni-Mtei"
+
+    client = translate.TranslationServiceClient()
+
+    location = "global"
+
+    parent = f"projects/{project_id}/locations/{location}"
+
+    response = client.translate_text(
+        request={
+            "parent": parent,
+            "contents": [text],
+            "mime_type": "text/plain",  # mime types: text/plain, text/html
+            "source_language_code": src_lang,
+            "target_language_code": tgt_lang,
+        }
+    )
+
+    translated_text = ""
+    for translation in response.translations:
+        translated_text += translation.translated_text
+
+    return translated_text
+
+
+if __name__ == "__main__":
+    root_dir = sys.argv[1]
+
+    pairs = sorted(glob.glob(os.path.join(root_dir, "*")))
+
+    for pair in pairs:
+
+        print(pair)
+
+        basename = os.path.basename(pair)
+
+        src_lang, tgt_lang = basename.split("-")
+        if src_lang not in flores_to_iso.keys() or tgt_lang not in flores_to_iso.keys():
+            continue
+
+        if src_lang == "eng_Latn":
+            lang = tgt_lang
+        else:
+            lang = src_lang
+
+        lang = flores_to_iso[lang]
+
+        if lang not in "as bn doi gom gu hi kn mai ml mni_Mtei mr ne or pa sa sd ta te ur":
+            continue
+
+        print(f"{src_lang} - {tgt_lang}")
+
+        # source to target translations
+
+        src_infname = os.path.join(pair, f"test.{src_lang}")
+        tgt_outfname = os.path.join(pair, f"test.{tgt_lang}.pred.google")
+        if os.path.exists(src_infname) and not os.path.exists(tgt_outfname):
+            src_sents = [
+                sent.replace("\n", "").strip()
+                for sent in open(src_infname, "r").read().split("\n")
+                if sent
+            ]
+            translations = [
+                translate_text(text, src_lang, tgt_lang).strip() for text in tqdm(src_sents)
+            ]
+            with open(tgt_outfname, "w") as f:
+                f.write("\n".join(translations))
+
+        # # target to source translations
+        tgt_infname = os.path.join(pair, f"test.{tgt_lang}")
+        src_outfname = os.path.join(pair, f"test.{src_lang}.pred.google")
+        if os.path.exists(tgt_infname) and not os.path.exists(src_outfname):
+            tgt_sents = [
+                sent.replace("\n", "").strip()
+                for sent in open(tgt_infname, "r").read().split("\n")
+                if sent
+            ]
+            translations = [
+                translate_text(text, tgt_lang, src_lang).strip() for text in tqdm(tgt_sents)
+            ]
+
+            with open(src_outfname, "w") as f:
+                f.write("\n".join(translations))

IndicTrans2/baseline_eval/m2m100_inference.py

@@ -0,0 +1,148 @@
+import os
+import re
+import sys
+from tqdm import tqdm
+import torch
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+
+
+# dictionary mapping flores codes to M2M-100 supported codes
+langs_supported = {
+    "eng_Latn": "en",
+    "ben_Beng": "bn",
+    "guj_Gujr": "gu",
+    "hin_Deva": "hi",
+    "kan_Knda": "kn",
+    "mal_Mlym": "ml",
+    "mar_Deva": "mr",
+    "npi_Deva": "ne",
+    "ory_Orya": "or",
+    "pan_Guru": "pa",
+    "snd_Arab": "sd",
+    "tam_Taml": "ta",
+    "urd_Arab": "ur",
+}
+
+
+def predict(batch, tokenizer, model, bos_token_id):
+    encoded_batch = tokenizer(batch, padding=True, return_tensors="pt").to(model.device)
+    generated_tokens = model.generate(
+        **encoded_batch,
+        num_beams=5,
+        max_length=256,
+        min_length=0,
+        forced_bos_token_id=bos_token_id,
+    )
+    hypothesis = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
+    return hypothesis
+
+
+def main(devtest_data_dir, batch_size):
+    # load the pre-trained M2M-100 tokenizer and model
+    model_name = "facebook/m2m100-12B-last-ckpt"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+    model.eval()
+
+    # iterate over a list of language pairs from `devtest_data_dir`
+    for pair in sorted(os.listdir(devtest_data_dir)):
+        if "-" not in pair:
+            continue
+
+        src_lang, tgt_lang = pair.split("-")
+
+        # check if the source and target languages are supported
+        if (
+            src_lang not in langs_supported.keys()
+            or tgt_lang not in langs_supported.keys()
+        ):
+            print(f"Skipping {src_lang}-{tgt_lang} ...")
+            continue
+
+        # -------------------------------------------------------------------
+        #                   source to target evaluation
+        # -------------------------------------------------------------------
+        print(f"Evaluating {src_lang}-{tgt_lang} ...")
+
+        infname = os.path.join(devtest_data_dir, pair, f"test.{src_lang}")
+        outfname = os.path.join(devtest_data_dir, pair, f"test.{tgt_lang}.pred.m2m100")
+
+        with open(infname, "r") as f:
+            src_sents = f.read().split("\n")
+
+        add_new_line = False
+        if src_sents[-1] == "":
+            add_new_line = True
+            src_sents = src_sents[:-1]
+
+        # set the source language for tokenization
+        tokenizer.src_lang = langs_supported[src_lang]
+
+        # process sentences in batches and generate predictions
+        hypothesis = []
+        for i in tqdm(range(0, len(src_sents), batch_size)):
+            start, end = i, int(min(len(src_sents), i + batch_size))
+            batch = src_sents[start:end]
+            bos_token_id = tokenizer.lang_code_to_id[langs_supported[tgt_lang]]
+            hypothesis += predict(batch, tokenizer, model, bos_token_id)
+
+        assert len(hypothesis) == len(src_sents)
+
+        hypothesis = [
+            re.sub("\s+", " ", x.replace("\n", " ").replace("\t", " ")).strip()
+            for x in hypothesis
+        ]
+        if add_new_line:
+            hypothesis = hypothesis
+
+        with open(outfname, "w") as f:
+            f.write("\n".join(hypothesis))
+
+        # -------------------------------------------------------------------
+        #                   target to source evaluation
+        # -------------------------------------------------------------------
+        infname = os.path.join(devtest_data_dir, pair, f"test.{tgt_lang}")
+        outfname = os.path.join(devtest_data_dir, pair, f"test.{src_lang}.pred.m2m100")
+
+        with open(infname, "r") as f:
+            src_sents = f.read().split("\n")
+
+        add_new_line = False
+        if src_sents[-1] == "":
+            add_new_line = True
+            src_sents = src_sents[:-1]
+
+        # set the source language for tokenization
+        tokenizer.src_lang = langs_supported[tgt_lang]
+
+        # process sentences in batches and generate predictions
+        hypothesis = []
+        for i in tqdm(range(0, len(src_sents), batch_size)):
+            start, end = i, int(min(len(src_sents), i + batch_size))
+            batch = src_sents[start:end]
+            bos_token_id = tokenizer.lang_code_to_id[langs_supported[src_lang]]
+            hypothesis += predict(batch, tokenizer, model, bos_token_id)
+
+        assert len(hypothesis) == len(src_sents)
+
+        hypothesis = [
+            re.sub("\s+", " ", x.replace("\n", " ").replace("\t", " ")).strip()
+            for x in hypothesis
+        ]
+        if add_new_line:
+            hypothesis = hypothesis
+
+        with open(outfname, "w") as f:
+            f.write("\n".join(hypothesis))
+
+
+if __name__ == "__main__":
+    # expects En-X subdirectories pairs within the devtest data directory
+    devtest_data_dir = sys.argv[1]
+    batch_size = int(sys.argv[2])
+
+    if not torch.cuda.is_available():
+        print("No GPU available")
+        sys.exit(1)
+
+    main(devtest_data_dir, batch_size)

IndicTrans2/baseline_eval/mbart_inference.py

@@ -0,0 +1,159 @@
+import os
+import re
+import sys
+from tqdm import tqdm
+import torch
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+
+
+# dictionary mapping flores codes to mBART supported codes
+langs_supported = {
+    "eng_Latn": "en_XX",
+    "guj_Gujr": "gu_IN",
+    "hin_Deva": "hi_IN",
+    "npi_Deva": "ne_NP",
+    "ben_Beng": "bn_IN",
+    "mal_Mlym": "ml_IN",
+    "mar_Deva": "mr_IN",
+    "tam_Taml": "ta_IN",
+    "tel_Telu": "te_IN",
+    "urd_Arab": "ur_PK",
+}
+
+
+def predict(batch, tokenizer, model, bos_token_id):
+    encoded_batch = tokenizer(batch, padding=True, return_tensors="pt").to(model.device)
+    generated_tokens = model.generate(
+        **encoded_batch,
+        num_beams=5,
+        max_length=256,
+        min_length=0,
+        forced_bos_token_id=bos_token_id,
+    )
+    hypothesis = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
+    return hypothesis
+
+
+def main(devtest_data_dir, batch_size):
+    # load the pre-trained mBART tokenizers and models for English-XX and XX-English translation
+    enxx_model_name = "facebook/mbart-large-50-one-to-many-mmt"
+    xxen_model_name = "facebook/mbart-large-50-many-to-one-mmt"
+    tokenizers = {
+        "enxx": AutoTokenizer.from_pretrained(enxx_model_name),
+        "xxen": AutoTokenizer.from_pretrained(xxen_model_name),
+    }
+    models = {
+        "enxx": AutoModelForSeq2SeqLM.from_pretrained(enxx_model_name).cuda(),
+        "xxen": AutoModelForSeq2SeqLM.from_pretrained(xxen_model_name).cuda(),
+    }
+
+    # set the models to evaluation mode
+    for model_name in models:
+        models[model_name].eval()
+
+    # iterate over a list of language pairs from `devtest_data_dir`
+    for pair in sorted(os.listdir(devtest_data_dir)):
+        if "-" not in pair:
+            continue
+
+        src_lang, tgt_lang = pair.split("-")
+
+        # check if the source and target languages are supported
+        if (
+            src_lang not in langs_supported.keys()
+            or tgt_lang not in langs_supported.keys()
+        ):
+            print(f"Skipping {src_lang}-{tgt_lang} ...")
+            continue
+
+        # -------------------------------------------------------------------
+        #                   source to target evaluation
+        # -------------------------------------------------------------------
+        print(f"Evaluating {src_lang}-{tgt_lang} ...")
+
+        infname = os.path.join(devtest_data_dir, pair, f"test.{src_lang}")
+        outfname = os.path.join(devtest_data_dir, pair, f"test.{tgt_lang}.pred.mbart50")
+
+        with open(infname, "r") as f:
+            src_sents = f.read().split("\n")
+
+        add_new_line = False
+        if src_sents[-1] == "":
+            add_new_line = True
+            src_sents = src_sents[:-1]
+
+        # set the source language for tokenization
+        tokenizers["enxx"].src_lang = langs_supported[src_lang]
+
+        # process sentences in batches and generate predictions
+        hypothesis = []
+        for i in tqdm(range(0, len(src_sents), batch_size)):
+            start, end = i, int(min(len(src_sents), i + batch_size))
+            batch = src_sents[start:end]
+            bos_token_id = tokenizers["enxx"].lang_code_to_id[langs_supported[tgt_lang]]
+            hypothesis += predict(
+                batch, tokenizers["enxx"], models["enxx"], bos_token_id
+            )
+
+        assert len(hypothesis) == len(src_sents)
+
+        hypothesis = [
+            re.sub("\s+", " ", x.replace("\n", " ").replace("\t", " ")).strip()
+            for x in hypothesis
+        ]
+        if add_new_line:
+            hypothesis = hypothesis
+
+        with open(outfname, "w") as f:
+            f.write("\n".join(hypothesis))
+
+        # -------------------------------------------------------------------
+        #                   target to source evaluation
+        # -------------------------------------------------------------------
+        infname = os.path.join(devtest_data_dir, pair, f"test.{tgt_lang}")
+        outfname = os.path.join(devtest_data_dir, pair, f"test.{src_lang}.pred.mbart50")
+
+        with open(infname, "r") as f:
+            src_sents = f.read().split("\n")
+
+        add_new_line = False
+        if src_sents[-1] == "":
+            add_new_line = True
+            src_sents = src_sents[:-1]
+
+        # set the source language for tokenization
+        tokenizers["xxen"].src_lang = langs_supported[tgt_lang]
+
+        # process sentences in batches and generate predictions
+        hypothesis = []
+        for i in tqdm(range(0, len(src_sents), batch_size)):
+            start, end = i, int(min(len(src_sents), i + batch_size))
+            batch = src_sents[start:end]
+            bos_token_id = tokenizers["xxen"].lang_code_to_id[langs_supported[src_lang]]
+            hypothesis += predict(
+                batch, tokenizers["xxen"], models["xxen"], bos_token_id
+            )
+
+        assert len(hypothesis) == len(src_sents)
+
+        hypothesis = [
+            re.sub("\s+", " ", x.replace("\n", " ").replace("\t", " ")).strip()
+            for x in hypothesis
+        ]
+        if add_new_line:
+            hypothesis = hypothesis
+
+        with open(outfname, "w") as f:
+            f.write("\n".join(hypothesis))
+
+
+if __name__ == "__main__":
+    # expects En-X subdirectories pairs within the devtest data directory
+    devtest_data_dir = sys.argv[1]
+    batch_size = int(sys.argv[2])
+
+    if not torch.cuda.is_available():
+        print("No GPU available")
+        sys.exit(1)
+
+    main(devtest_data_dir, batch_size)

IndicTrans2/baseline_eval/nllb_moe_cpu_inference.py

@@ -0,0 +1,157 @@
+import os
+import re
+import sys
+from tqdm import tqdm
+import torch
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+
+langs_supported = [
+    "asm_Beng",
+    "ben_Beng",
+    "guj_Gujr",
+    "eng_Latn",
+    "hin_Deva",
+    "kas_Deva",
+    "kas_Arab",
+    "kan_Knda",
+    "mal_Mlym",
+    "mai_Deva",
+    "mar_Deva",
+    "mni_Beng",
+    "npi_Deva",
+    "ory_Orya",
+    "pan_Guru",
+    "san_Deva",
+    "snd_Arab",
+    "sat_Olck",
+    "tam_Taml",
+    "tel_Telu",
+    "urd_Arab",
+]
+
+
+def predict(batch, tokenizer, model, bos_token_id):
+    encoded_batch = tokenizer(batch, padding=True, return_tensors="pt").to(model.device)
+    generated_tokens = model.generate(
+        **encoded_batch,
+        num_beams=5,
+        max_length=256,
+        min_length=0,
+        forced_bos_token_id=bos_token_id,
+    )
+    hypothesis = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
+    return hypothesis
+
+
+def main(devtest_data_dir, batch_size):
+    # load the pre-trained NLLB tokenizer and model
+    model_name = "facebook/nllb-moe-54b"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+    model.eval()
+
+    # iterate over a list of language pairs from `devtest_data_dir`
+    for pair in sorted(os.listdir(devtest_data_dir)):
+        if "-" not in pair:
+            continue
+
+        src_lang, tgt_lang = pair.split("-")
+
+        # check if the source and target languages are supported
+        if (
+            src_lang not in langs_supported.keys()
+            or tgt_lang not in langs_supported.keys()
+        ):
+            print(f"Skipping {src_lang}-{tgt_lang} ...")
+            continue
+
+        # -------------------------------------------------------------------
+        #                   source to target evaluation
+        # -------------------------------------------------------------------
+        print(f"Evaluating {src_lang}-{tgt_lang} ...")
+
+        infname = os.path.join(devtest_data_dir, pair, f"test.{src_lang}")
+        outfname = os.path.join(
+            devtest_data_dir, pair, f"test.{tgt_lang}.pred.nllb_moe"
+        )
+
+        with open(infname, "r") as f:
+            src_sents = f.read().split("\n")
+
+        add_new_line = False
+        if src_sents[-1] == "":
+            add_new_line = True
+            src_sents = src_sents[:-1]
+
+        # set the source language for tokenization
+        tokenizer.src_lang = src_lang
+
+        # process sentences in batches and generate predictions
+        hypothesis = []
+        for i in tqdm(range(0, len(src_sents), batch_size)):
+            start, end = i, int(min(len(src_sents), i + batch_size))
+            batch = src_sents[start:end]
+            if tgt_lang == "sat_Olck":
+                bos_token_id = tokenizer.lang_code_to_id["sat_Beng"]
+            else:
+                bos_token_id = tokenizer.lang_code_to_id[tgt_lang]
+            hypothesis += predict(batch, tokenizer, model, bos_token_id)
+
+        assert len(hypothesis) == len(src_sents)
+
+        hypothesis = [
+            re.sub("\s+", " ", x.replace("\n", " ").replace("\t", " ")).strip()
+            for x in hypothesis
+        ]
+        if add_new_line:
+            hypothesis = hypothesis
+
+        with open(outfname, "w") as f:
+            f.write("\n".join(hypothesis))
+
+        # -------------------------------------------------------------------
+        #                   target to source evaluation
+        # -------------------------------------------------------------------
+        infname = os.path.join(devtest_data_dir, pair, f"test.{tgt_lang}")
+        outfname = os.path.join(
+            devtest_data_dir, pair, f"test.{src_lang}.pred.nllb_moe"
+        )
+
+        with open(infname, "r") as f:
+            src_sents = f.read().split("\n")
+
+        add_new_line = False
+        if src_sents[-1] == "":
+            add_new_line = True
+            src_sents = src_sents[:-1]
+
+        # set the source language for tokenization
+        tokenizer.src_lang = "sat_Beng" if tgt_lang == "sat_Olck" else tgt_lang
+
+        # process sentences in batches and generate predictions
+        hypothesis = []
+        for i in tqdm(range(0, len(src_sents), batch_size)):
+            start, end = i, int(min(len(src_sents), i + batch_size))
+            batch = src_sents[start:end]
+            bos_token_id = tokenizer.lang_code_to_id[langs_supported[src_lang]]
+            hypothesis += predict(batch, tokenizer, model, bos_token_id)
+
+        assert len(hypothesis) == len(src_sents)
+
+        hypothesis = [
+            re.sub("\s+", " ", x.replace("\n", " ").replace("\t", " ")).strip()
+            for x in hypothesis
+        ]
+        if add_new_line:
+            hypothesis = hypothesis
+
+        with open(outfname, "w") as f:
+            f.write("\n".join(hypothesis))
+
+
+if __name__ == "__main__":
+    # expects En-X subdirectories pairs within the devtest data directory
+    devtest_data_dir = sys.argv[1]
+    batch_size = int(sys.argv[2])
+
+    main(devtest_data_dir, batch_size)

IndicTrans2/compute_comet_score.sh

@@ -0,0 +1,84 @@
+#!/bin/bash
+
+# This script computes COMET metrics and also performs significance testing on the evaluation set
+# where each subdirectory contains En-X pair
+
+
+echo `date`
+devtest_data_dir=$1                         # path to the evaluation directory
+model_name=${2-"Unbabel/wmt22-comet-da"}    # name of the model checkpoint
+
+# predefined list of languages supported by COMET
+langs=(asm_Beng ben_Beng guj_Gujr hin_Deva kan_Knda mal_Mlym mar_Deva ory_Orya pan_Guru tam_Taml tel_Telu urd_Arab)
+
+# we predefine a set of systems which we consider for evaluation
+# feel free to change the below line in case you want to add or remove any system
+system=(google azure nllb mbart50 m2m100 it1 it2)
+
+
+# iterate over the list of predefined languages
+for lang in "${langs[@]}"; do
+
+    mkdir -p "$devtest_data_dir/eng_Latn-$lang/comet"
+
+    # --------------------------------------------------------------
+    #                   COMET score computation
+    # --------------------------------------------------------------
+
+    # iterate over the list of predefined systems
+    for sys in "${system[@]}"; do
+
+        echo "${sys}"
+
+        # en - indic direction
+        if [ -f "$devtest_data_dir/eng_Latn-$lang/test.$lang.pred.$sys" ]; then
+            echo "eng_Latn-${lang}"
+
+            src_fname=$devtest_data_dir/eng_Latn-$lang/test.eng_Latn
+            pred_fname=$devtest_data_dir/eng_Latn-$lang/test.$lang.pred.$sys
+            ref_fname=$devtest_data_dir/eng_Latn-$lang/test.$lang
+            out_fname=$devtest_data_dir/eng_Latn-$lang/comet/eng_Latn_${lang}_${sys}_comet.txt
+
+            # Compute COMET scores using the `comet-score`
+            comet-score -s $src_fname -t $pred_fname -r $ref_fname --gpus 1 --model $model_name --quiet --only_system > $out_fname
+        fi
+
+        # indic - en direction
+        if [ -f "$devtest_data_dir/eng_Latn-$lang/test.eng_Latn.pred.$sys" ]; then
+            echo "${lang}-eng_Latn"
+
+            src_fname=$devtest_data_dir/eng_Latn-$lang/test.$lang
+            pred_fname=$devtest_data_dir/eng_Latn-$lang/test.eng_Latn.pred.$sys
+            ref_fname=$devtest_data_dir/eng_Latn-$lang/test.eng_Latn
+            out_fname=$devtest_data_dir/eng_Latn-$lang/comet/${lang}_eng_Latn_${sys}_comet.txt
+
+            # Compute COMET scores using the `comet-score`
+            comet-score -s $src_fname -t $pred_fname -r $ref_fname --gpus 1 --model $model_name --quiet --only_system > $out_fname
+        fi
+
+    done
+
+    # --------------------------------------------------------------
+    #                  COMET significance testing
+    # --------------------------------------------------------------
+    
+    # en - indic direction
+    src_fname=$devtest_data_dir/eng_Latn-$lang/test.eng_Latn
+    pred_fname=$devtest_data_dir/eng_Latn-$lang/test.$lang.pred.*
+    ref_fname=$devtest_data_dir/eng_Latn-$lang/test.$lang
+    out_fname=$devtest_data_dir/eng_Latn-$lang/comet/eng_Latn_${lang}_comet_stat.txt
+
+    # Compute COMET significance scores using the `comet-compare`
+    comet-compare -s $src_fname -t $pred_fname -r $ref_fname > $out_fname
+
+
+    # indic-en direction
+    src_fname=$devtest_data_dir/eng_Latn-$lang/test.$lang
+    pred_fname=$devtest_data_dir/eng_Latn-$lang/test.eng_Latn.pred.*
+    ref_fname=$devtest_data_dir/eng_Latn-$lang/test.eng_Latn
+    out_fname=$devtest_data_dir/eng_Latn-$lang/comet/${lang}_eng_Latn_comet_stat.txt
+
+    # Compute COMET significance scores using the `comet-compare`
+    comet-compare -s $src_fname -t $pred_fname -r $ref_fname > $out_fname
+
+done

IndicTrans2/compute_metrics.sh

@@ -0,0 +1,29 @@
+#!/bin/bash
+
+# This script compute the evaluation metrics such as BLEU, chrF, chrF++ using the 
+# detokenized predictions of the translation systems using sacrebleu (version 2.3.1).
+# If the target language is:
+#   English: directly use Moses tokenizer that is internally supported (`mteval-v13a`)
+#   Indic: use IndicNLP tokenizers and skip tokenization step in sacrebleu.
+
+
+echo `date`
+pred_fname=$1       # path to the predction file
+ref_fname=$2        # path to the reference file
+tgt_lang=$3         # target language
+
+
+if [ $tgt_lang == 'eng_Latn' ]; then
+    # directly tokenize the prediction and reference files using sacrebleu and compute the metric
+    sacrebleu $ref_fname < $pred_fname -m bleu chrf
+    sacrebleu $ref_fname < $pred_fname -m chrf --chrf-word-order 2
+else
+
+    # indicnlp tokenize prediction and reference files before evaluation
+    input_size=`python scripts/preprocess_translate.py $ref_fname $ref_fname.tok $tgt_lang false false`
+    input_size=`python scripts/preprocess_translate.py $pred_fname $pred_fname.tok $tgt_lang false false`
+
+    # since we are tokenizing with indicnlp separately, we are setting tokenize to none here
+    sacrebleu --tokenize none $ref_fname.tok < $pred_fname.tok -m bleu chrf
+    sacrebleu --tokenize none $ref_fname.tok < $pred_fname.tok -m chrf --chrf-word-order 2
+fi

IndicTrans2/compute_metrics_significance.sh

@@ -0,0 +1,66 @@
+#!/bin/bash
+
+# This script performs significance testing for metrics such as BLEU, chrF++ using sacrebleu on the evaluation set
+# where each subdirectory contains En-X pair
+
+
+echo `date`
+devtest_data_dir=$1                         # path to the evaluation directory
+
+# we predefine a set of systems which we consider for evaluation
+# feel free to change the below line in case you want to add or remove any system
+system=(google azure nllb mbart50 m2m100 it1 it2)
+
+
+# get a list of language pairs in the `devtest_data_dir`
+pairs=$(ls -d $devtest_data_dir/eng_Latn-* | sort)
+
+
+# iterate over each language pair
+for pair in ${pairs[@]}; do
+    # extract the source and target languages from the pair name
+    pair=$(basename $pair)
+    src_lang=$(echo "$pair" | cut -d "-" -f 1)
+    tgt_lang=$(echo "$pair" | cut -d "-" -f 2)
+
+    if [[ $src_lang == "eng_Latn" ]]; then
+    
+        # ----------------------------------------------------------------------
+        #                           en - indic direction
+        # ----------------------------------------------------------------------
+        echo "${src_lang} - ${tgt_lang}"
+
+        # find all the prediction files for different systems and tokenize it using IndicNLP
+        pred_fnames=$devtest_data_dir/$pair/test.${tgt_lang}.pred.*
+        ref_fname=$devtest_data_dir/$pair/test.${tgt_lang}
+
+        for pred_fname in $(find . -type f -name $pred_fnames); do
+            input_size=`python scripts/preprocess_translate.py $pred_fname $pred_fname.tok $tgt_lang false false`
+        done
+
+        input_size=`python scripts/preprocess_translate.py $ref_fname $ref_fname.tok $tgt_lang false false`
+
+        ref_fname=$devtest_data_dir/$pair/test.${tgt_lang}.tok
+        it2_fname=$devtest_data_dir/$pair/test.${tgt_lang}.pred.it2.tok
+        sys_fnames=$devtest_data_dir/$pair/test.${tgt_lang}.pred.*.tok
+        bleu_out_fname=$devtest_data_dir/$pair/${src_lang}_${tgt_lang}_bleu_significance.txt
+        chrF_out_fname=$devtest_data_dir/$pair/${src_lang}_${tgt_lang}_chrF++_significance.txt
+
+        sacrebleu --tokenize none $ref_fname -i $it2_fname $sys_fnames --paired-bs -m bleu --format text > $bleu_out_fname
+        sacrebleu --tokenize none $it2_fname $sys_fnames --paired-bs -m chrf --chrf-word-order 2 --format text > $chrF_out_fname
+
+        # ----------------------------------------------------------------------
+        #                           indic - en direction
+        # ----------------------------------------------------------------------
+        echo "${tgt_lang} - ${src_lang}"
+
+        ref_fname=$devtest_data_dir/$pair/test.${src_lang}
+        it2_fname=$devtest_data_dir/$pair/test.${src_lang}.pred.it2
+        sys_fnames=$devtest_data_dir/$pair/test.${src_lang}.pred.*
+        bleu_out_fname=$devtest_data_dir/$pair/${tgt_lang}_${src_lang}_bleu_significance.txt
+        chrF_out_fname=$devtest_data_dir/$pair/${tgt_lang}_${src_lang}_chrF++_significance.txt
+
+        sacrebleu --tokenize none $ref_fname -i $it2_fname $sys_fnames --paired-bs -m bleu --format text > $bleu_out_fname
+        sacrebleu --tokenize none $it2_fname $sys_fnames --paired-bs -m chrf --chrf-word-order 2 --format text > $chrF_out_fname
+    
+    fi

IndicTrans2/eval.sh

@@ -0,0 +1,54 @@
+#!/bin/bash
+
+# This script evaluates the performance of a machine translation system 
+# on a evaluation set in forward direction. For example, if the evaluation set 
+# consists of language pairs, such as En-X, where En represents the English language 
+# and X represents the target Indic language then this script accesses the translation
+# system from the English language (En) to the target Indic language (X) direction.
+
+
+echo `date`
+devtest_data_dir=$1         # path to the evaluation directory
+ckpt_dir=$2                 # path to the checkpoint directory
+system=${3:-"it2"}          # name of the machine translation system
+
+
+# get a list of language pairs in the `devtest_data_dir`
+pairs=$(ls -d $devtest_data_dir/* | sort)
+
+
+# iterate over each language pair
+for pair in ${pairs[@]}; do
+    # extract the source and target languages from the pair name
+    pair=$(basename $pair)
+    src_lang=$(echo "$pair" | cut -d "-" -f 1)
+    tgt_lang=$(echo "$pair" | cut -d "-" -f 2)
+
+    src_fname=$devtest_data_dir/$src_lang-$tgt_lang/test.$src_lang
+    tgt_fname=$devtest_data_dir/$src_lang-$tgt_lang/test.$tgt_lang
+
+    # check if the source and target files exists
+    if [ -f "$src_fname" ] && [ -f "$tgt_fname" ]; then
+        echo "Evaluating $src_lang-$tgt_lang ..."
+    else 
+        echo "Skipping $src_lang-$tgt_lang ..."
+        continue
+    fi
+
+    # generate translations if the system name contains "it2"
+    if [[ $system == *"it2"* ]]; then
+        echo "Generating Translations"
+        bash joint_translate.sh $src_fname $tgt_fname.pred.$system $src_lang $tgt_lang $ckpt_dir
+    fi
+
+    # compute automatic string-based metrics if the prediction exists for the system
+    if [[ -f "${tgt_fname}.pred.${system}" ]]; then
+        echo "Computing Metrics"
+        bash compute_metrics.sh $tgt_fname.pred.$system $tgt_fname $tgt_lang > $devtest_data_dir/$src_lang-$tgt_lang/${src_lang}_${tgt_lang}_${system}_scores.txt
+    fi
+
+    # remove the intermediate files
+    rm -rf $tgt_fname.pred.$system.*
+    rm -rf $devtest_data_dir/$src_lang-$tgt_lang/*.tok
+
+done

IndicTrans2/eval_rev.sh

@@ -0,0 +1,55 @@
+#!/bin/bash
+
+# This script evaluates the performance of a machine translation system 
+# on a evaluation set in forward direction. For example, if the evaluation set 
+# consists of language pairs, such as En-X, where En represents the English language 
+# and X represents the target Indic language then this script accesses the translation
+# system from the target Indic language (X) to the English language (En) direction.
+
+
+echo `date`
+devtest_data_dir=$1         # path to the evaluation directory
+ckpt_dir=$2                 # path to the checkpoint directory
+system=${3:-"it2"}          # name of the machine translation system
+
+
+# get a list of language pairs in the `devtest_data_dir`
+pairs=$(ls -d $devtest_data_dir/* | sort)
+
+
+# iterate over each language pair
+for pair in ${pairs[@]}; do
+    # extract the source and target languages from the pair name
+    pair=$(basename $pair)
+    src_lang=$(echo "$pair" | cut -d "-" -f 1)
+    tgt_lang=$(echo "$pair" | cut -d "-" -f 2)
+
+    src_fname=$devtest_data_dir/$src_lang-$tgt_lang/test.$tgt_lang
+    tgt_fname=$devtest_data_dir/$src_lang-$tgt_lang/test.$src_lang
+
+    # check if the source and target files exists
+    # in this case, we flip the actual target file as source and vice-versa
+    if [ -f "$src_fname" ] && [ -f "$tgt_fname" ]; then
+        echo "Evaluating $src_lang-$tgt_lang ..."
+    else 
+        echo "Skipping $src_lang-$tgt_lang ..."
+        continue
+    fi
+
+    # generate translations if the system name contains "it2"
+    if [[ $system == *"it2"* ]]; then
+        echo "Generating Translations"
+        bash joint_translate.sh $src_fname $tgt_fname.pred.$system $tgt_lang $src_lang $ckpt_dir
+    fi
+
+    # compute automatic string-based metrics if the prediction exists for the system
+    if [[ -f "${tgt_fname}.pred.${system}" ]]; then
+        echo "Computing Metrics"
+        bash compute_metrics.sh $tgt_fname.pred.$system $tgt_fname $src_lang > $devtest_data_dir/$src_lang-$tgt_lang/${tgt_lang}_${src_lang}_${system}_scores.txt
+    fi
+
+    # remove the intermediate files
+    rm -rf $tgt_fname.pred.$system.*
+    rm -rf $devtest_data_dir/$src_lang-$tgt_lang/*.tok
+
+done

IndicTrans2/finetune.sh

@@ -0,0 +1,54 @@
+#!/bin/bash
+
+# This script finetunes the pretrained translation model on the binarized data using fairseq.
+
+
+echo `date`
+exp_dir=$1                              # path of the experiment directory
+model_arch=${2:-"transformer_18_18"}    # model architecture (defaults to `transformer_18_18`)
+pretrained_ckpt=$3                      # path to the pretrained checkpoint `.pt` file
+
+
+fairseq-train $exp_dir/final_bin \
+--max-source-positions=256 \
+--max-target-positions=256 \
+--source-lang=SRC \
+--target-lang=TGT \
+--max-update=1000000 \
+--save-interval-updates=1000 \
+--arch=$model_arch \
+--activation-fn gelu \
+--criterion=label_smoothed_cross_entropy \
+--label-smoothing=0.1 \
+--optimizer adam \
+--adam-betas "(0.9, 0.98)" \
+--lr-scheduler=inverse_sqrt \
+--clip-norm 1.0 \
+--warmup-init-lr 1e-07 \
+--lr 3e-5 \
+--warmup-updates 2000 \
+--dropout 0.2 \
+--save-dir $exp_dir/model \
+--keep-last-epochs 5 \
+--keep-interval-updates 3 \
+--patience 10 \
+--skip-invalid-size-inputs-valid-test \
+--fp16 \
+--user-dir model_configs \
+--update-freq=4 \
+--distributed-world-size 8 \
+--num-workers 24 \
+--max-tokens 1024 \
+--eval-bleu \
+--eval-bleu-args "{\"beam\": 1, \"lenpen\": 1.0, \"max_len_a\": 1.2, \"max_len_b\": 10}" \
+--eval-bleu-detok moses \
+--eval-bleu-remove-bpe sentencepiece \
+--eval-bleu-print-samples \
+--best-checkpoint-metric bleu \
+--maximize-best-checkpoint-metric \
+--restore-file $pretrained_ckpt \
+--reset-lr-scheduler \
+--reset-meters \
+--reset-dataloader \
+--reset-optimizer \
+--task translation

IndicTrans2/huggingface_interface/.gitignore

@@ -0,0 +1 @@
+IndicTransTokenizer

IndicTrans2/huggingface_interface/README.md

@@ -0,0 +1,119 @@
+# IndicTrans2 HF Compatible Models
+
+[![colab link](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/AI4Bharat/IndicTrans2/blob/main/huggingface_interface/colab_inference.ipynb)
+
+In this section, we provide details on how to use our [IndicTrans2](https://github.com/AI4Bharat/IndicTrans2) models which were originally trained with the [fairseq](https://github.com/facebookresearch/fairseq) to [HuggingFace transformers](https://huggingface.co/docs/transformers/index) for inference purpose. Our scripts for HuggingFace compatible models are adapted from [M2M100 repository](https://github.com/huggingface/transformers/tree/main/src/transformers/models/m2m_100).
+
+> Note: We have migrated IndicTrans2 tokenizer for HF compatible IndicTrans2 models to [IndicTransToolkit](https://github.com/VarunGumma/IndicTransToolkit) and will be maintained separately there from now onwards. This is automatically installed when you call `install.sh` script in `huggingface_interface`.
+
+### Setup
+
+To get started, follow these steps to set up the environment:
+
+```
+# Clone the github repository and navigate to the project directory.
+git clone https://github.com/AI4Bharat/IndicTrans2
+cd IndicTrans2/huggingface_interface
+
+# Install all the dependencies and requirements associated with the project for running HF compatible models.
+source install.sh
+```
+
+> Note: The `install.sh` script in this directory is specifically for running HF compatible models for inference.
+
+### Converting
+
+In order to convert the fairseq checkpoint to a PyTorch checkpoint that is compatible with HuggingFace Transformers, use the following command:
+
+```bash
+python3 convert_indictrans_checkpoint_to_pytorch.py --fairseq_path <fairseq_checkpoint_best.pt> --pytorch_dump_folder_path <hf_output_dir>
+```
+
+- `<fairseq_checkpoint_best.pt>`: path to the fairseq `checkpoint_best.pt` that needs to be converted to HF compatible models
+- `<hf_output_dir>`: path to the output directory where the HF compatible models will be saved
+
+### Models
+
+| Model                            | 🤗 HuggingFace Checkpoints                                                                                        |
+| -------------------------------- | ----------------------------------------------------------------------------------------------------------------- |
+| En-Indic                         | [ai4bharat/indictrans2-en-indic-1B](https://huggingface.co/ai4bharat/indictrans2-en-indic-1B)                     |
+| Indic-En                         | [ai4bharat/indictrans2-indic-en-1B](https://huggingface.co/ai4bharat/indictrans2-indic-en-1B)                     |
+| Distilled En-Indic               | [ai4bharat/indictrans2-en-indic-dist-200M](https://huggingface.co/ai4bharat/indictrans2-en-indic-dist-200M)       |
+| Distilled Indic-En               | [ai4bharat/indictrans2-indic-en-dist-200M](https://huggingface.co/ai4bharat/indictrans2-indic-en-dist-200M)       |
+| Indic-Indic (Stitched)           | [ai4bharat/indictrans2-indic-indic-1B](https://huggingface.co/ai4bharat/indictrans2-indic-indic-1B)               |
+| Distilled Indic-Indic (Stitched) | [ai4bharat/indictrans2-indic-indic-dist-320M](https://huggingface.co/ai4bharat/indictrans2-indic-indic-dist-320M) |
+
+### Inference
+
+With the conversion complete, you can now perform inference using the HuggingFace Transformers.
+
+You can start with the provided `example.py` script and customize it for your specific translation use case:
+
+```bash
+python3 example.py
+```
+
+Feel free to modify the `example.py` script to suit your translation needs.
+
+### Fine-tuning with LoRA
+
+Before starting with fine-tuning IndicTrans2 models, you will need to restructure the training data in the following format.
+
+```
+en-indic-exp
+├── train
+│   ├── eng_Latn-asm_Beng
+│   │   ├── train.eng_Latn
+│   │   └── train.asm_Beng
+│   ├── eng_Latn-ben_Beng
+│   │   └── ...
+│   └── {src_lang}-{tgt_lang}
+│       ├── train.{src_lang}
+│       └── train.{tgt_lang}
+└── dev
+    ├── eng_Latn-asm_Beng
+    │   ├── dev.eng_Latn
+    │   └── dev.asm_Beng
+    ├── eng_Latn-ben_Beng
+    │   └── ...
+    └── {src_lang}-{tgt_lang}
+        ├── dev.{src_lang}
+        └── dev.{tgt_lang}
+```
+
+Once you have data ready in above specified format, use the following command.
+
+```bash
+bash train_lora.sh <data_dir> <model_name> <output_dir> <direction> <src_lang_list> <tgt_lang_list>
+```
+
+We recommend you to refer to `train_lora.sh` for defaults arguments for fine-tuning. Please note that the specified hyperparameters may not be optimal and might require tuning for your use case.
+
+### Inference with LoRA
+
+You can load the LoRA adapters with the base model for inference by modifying the model initialization in `example.py` script.
+
+```python
+from transformers import AutoModelForSeq2SeqLM
+from peft import PeftConfig, PeftModel
+
+base_ckpt_dir = "ai4bharat/indictrans2-en-indic-1B" # you will need to change as per your use case
+base_model = AutoModelForSeq2SeqLM.from_pretrained(base_ckpt_dir, trust_remote_code=True)
+lora_model = PeftModel.from_pretrained(base_model, lora_ckpt_dir)
+```
+
+> Note: Please feel free to open issues on the GitHub repo in case of any queries/issues.
+
+### Citation
+
+```bibtex
+@article{gala2023indictrans,
+title={IndicTrans2: Towards High-Quality and Accessible Machine Translation Models for all 22 Scheduled Indian Languages},
+author={Jay Gala and Pranjal A Chitale and A K Raghavan and Varun Gumma and Sumanth Doddapaneni and Aswanth Kumar M and Janki Atul 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},
+issn={2835-8856},
+year={2023},
+url={https://openreview.net/forum?id=vfT4YuzAYA},
+note={}
+}
+```

IndicTrans2/huggingface_interface/colab_inference.ipynb

@@ -0,0 +1,458 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "8Aa-nRCzPVdF"
+      },
+      "source": [
+        "# IndicTrans2 HF Inference\n",
+        "\n",
+        "We provide an example notebook on how to use our IndicTrans2 models which were originally trained with the fairseq to HuggingFace transformers for inference purpose.\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "Cfsv02IeP2It"
+      },
+      "source": [
+        "## Setup\n",
+        "\n",
+        "Please run the cells below to install the necessary dependencies.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "qKcYlUZYGLrt"
+      },
+      "outputs": [],
+      "source": [
+        "%%capture\n",
+        "!git clone https://github.com/AI4Bharat/IndicTrans2.git"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "U3vs7FkIGSxK"
+      },
+      "outputs": [],
+      "source": [
+        "%%capture\n",
+        "%cd /content/IndicTrans2/huggingface_interface"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ddkRAXQ2Git0"
+      },
+      "outputs": [],
+      "source": [
+        "%%capture\n",
+        "!python3 -m pip install nltk sacremoses pandas regex mock transformers>=4.33.2 mosestokenizer\n",
+        "!python3 -c \"import nltk; nltk.download('punkt')\"\n",
+        "!python3 -m pip install bitsandbytes scipy accelerate datasets\n",
+        "!python3 -m pip install sentencepiece\n",
+        "\n",
+        "!git clone https://github.com/VarunGumma/IndicTransToolkit.git\n",
+        "%cd IndicTransToolkit\n",
+        "!python3 -m pip install --editable ./\n",
+        "%cd .."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "hjN7ub1tO33H"
+      },
+      "source": [
+        "**IMPORTANT : Restart your run-time first and then run the cells below.**"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "_SLBIw6rQB-0"
+      },
+      "source": [
+        "## Inference\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "fYczM2U6G1Zv"
+      },
+      "outputs": [],
+      "source": [
+        "import torch\n",
+        "from transformers import AutoModelForSeq2SeqLM, BitsAndBytesConfig, AutoTokenizer\n",
+        "from IndicTransToolkit import IndicProcessor\n",
+        "\n",
+        "BATCH_SIZE = 4\n",
+        "DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
+        "quantization = None"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "xj1WCNjuHG-d"
+      },
+      "outputs": [],
+      "source": [
+        "def initialize_model_and_tokenizer(ckpt_dir, quantization):\n",
+        "    if quantization == \"4-bit\":\n",
+        "        qconfig = BitsAndBytesConfig(\n",
+        "            load_in_4bit=True,\n",
+        "            bnb_4bit_use_double_quant=True,\n",
+        "            bnb_4bit_compute_dtype=torch.bfloat16,\n",
+        "        )\n",
+        "    elif quantization == \"8-bit\":\n",
+        "        qconfig = BitsAndBytesConfig(\n",
+        "            load_in_8bit=True,\n",
+        "            bnb_8bit_use_double_quant=True,\n",
+        "            bnb_8bit_compute_dtype=torch.bfloat16,\n",
+        "        )\n",
+        "    else:\n",
+        "        qconfig = None\n",
+        "\n",
+        "    tokenizer = AutoTokenizer.from_pretrained(ckpt_dir, trust_remote_code=True)\n",
+        "    model = AutoModelForSeq2SeqLM.from_pretrained(\n",
+        "        ckpt_dir,\n",
+        "        trust_remote_code=True,\n",
+        "        low_cpu_mem_usage=True,\n",
+        "        quantization_config=qconfig,\n",
+        "    )\n",
+        "\n",
+        "    if qconfig == None:\n",
+        "        model = model.to(DEVICE)\n",
+        "        if DEVICE == \"cuda\":\n",
+        "            model.half()\n",
+        "\n",
+        "    model.eval()\n",
+        "\n",
+        "    return tokenizer, model\n",
+        "\n",
+        "\n",
+        "def batch_translate(input_sentences, src_lang, tgt_lang, model, tokenizer, ip):\n",
+        "    translations = []\n",
+        "    for i in range(0, len(input_sentences), BATCH_SIZE):\n",
+        "        batch = input_sentences[i : i + BATCH_SIZE]\n",
+        "\n",
+        "        # Preprocess the batch and extract entity mappings\n",
+        "        batch = ip.preprocess_batch(batch, src_lang=src_lang, tgt_lang=tgt_lang)\n",
+        "\n",
+        "        # Tokenize the batch and generate input encodings\n",
+        "        inputs = tokenizer(\n",
+        "            batch,\n",
+        "            truncation=True,\n",
+        "            padding=\"longest\",\n",
+        "            return_tensors=\"pt\",\n",
+        "            return_attention_mask=True,\n",
+        "        ).to(DEVICE)\n",
+        "\n",
+        "        # Generate translations using the model\n",
+        "        with torch.no_grad():\n",
+        "            generated_tokens = model.generate(\n",
+        "                **inputs,\n",
+        "                use_cache=True,\n",
+        "                min_length=0,\n",
+        "                max_length=256,\n",
+        "                num_beams=5,\n",
+        "                num_return_sequences=1,\n",
+        "            )\n",
+        "\n",
+        "        # Decode the generated tokens into text\n",
+        "\n",
+        "        with tokenizer.as_target_tokenizer():\n",
+        "            generated_tokens = tokenizer.batch_decode(\n",
+        "                generated_tokens.detach().cpu().tolist(),\n",
+        "                skip_special_tokens=True,\n",
+        "                clean_up_tokenization_spaces=True,\n",
+        "            )\n",
+        "\n",
+        "        # Postprocess the translations, including entity replacement\n",
+        "        translations += ip.postprocess_batch(generated_tokens, lang=tgt_lang)\n",
+        "\n",
+        "        del inputs\n",
+        "        torch.cuda.empty_cache()\n",
+        "\n",
+        "    return translations"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "erNCuZTEMt49"
+      },
+      "source": [
+        "### English to Indic Example\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "6OG3Bw-sHnf3",
+        "outputId": "a204f50e-9456-4fb1-900a-e60680b97b99"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "\n",
+            "eng_Latn - hin_Deva\n",
+            "eng_Latn: When I was young, I used to go to the park every day.\n",
+            "hin_Deva: जब मैं छोटा था, मैं हर दिन पार्क जाता था।\n",
+            "eng_Latn: He has many old books, which he inherited from his ancestors.\n",
+            "hin_Deva: उनके पास कई पुरानी किताबें हैं, जो उन्हें अपने पूर्वजों से विरासत में मिली हैं।\n",
+            "eng_Latn: I can't figure out how to solve my problem.\n",
+            "hin_Deva: मुझे समझ नहीं आ रहा है कि मैं अपनी समस्या का समाधान कैसे करूं।\n",
+            "eng_Latn: She is very hardworking and intelligent, which is why she got all the good marks.\n",
+            "hin_Deva: वह बहुत मेहनती और बुद्धिमान है, यही कारण है कि उसे सभी अच्छे अंक मिले।\n",
+            "eng_Latn: We watched a new movie last week, which was very inspiring.\n",
+            "hin_Deva: हमने पिछले हफ्ते एक नई फिल्म देखी, जो बहुत प्रेरणादायक थी।\n",
+            "eng_Latn: If you had met me at that time, we would have gone out to eat.\n",
+            "hin_Deva: अगर आप उस समय मुझसे मिलते तो हम बाहर खाना खाने जाते।\n",
+            "eng_Latn: She went to the market with her sister to buy a new sari.\n",
+            "hin_Deva: वह अपनी बहन के साथ नई साड़ी खरीदने के लिए बाजार गई थी।\n",
+            "eng_Latn: Raj told me that he is going to his grandmother's house next month.\n",
+            "hin_Deva: राज ने मुझे बताया कि वह अगले महीने अपनी दादी के घर जा रहा है।\n",
+            "eng_Latn: All the kids were having fun at the party and were eating lots of sweets.\n",
+            "hin_Deva: पार्टी में सभी बच्चे खूब मस्ती कर रहे थे और खूब मिठाइयां खा रहे थे।\n",
+            "eng_Latn: My friend has invited me to his birthday party, and I will give him a gift.\n",
+            "hin_Deva: मेरे दोस्त ने मुझे अपने जन्मदिन की पार्टी में आमंत्रित किया है, और मैं उसे एक उपहार दूंगा।\n"
+          ]
+        }
+      ],
+      "source": [
+        "en_indic_ckpt_dir = \"ai4bharat/indictrans2-en-indic-1B\"  # ai4bharat/indictrans2-en-indic-dist-200M\n",
+        "en_indic_tokenizer, en_indic_model = initialize_model_and_tokenizer(en_indic_ckpt_dir, quantization)\n",
+        "\n",
+        "ip = IndicProcessor(inference=True)\n",
+        "\n",
+        "en_sents = [\n",
+        "    \"When I was young, I used to go to the park every day.\",\n",
+        "    \"He has many old books, which he inherited from his ancestors.\",\n",
+        "    \"I can't figure out how to solve my problem.\",\n",
+        "    \"She is very hardworking and intelligent, which is why she got all the good marks.\",\n",
+        "    \"We watched a new movie last week, which was very inspiring.\",\n",
+        "    \"If you had met me at that time, we would have gone out to eat.\",\n",
+        "    \"She went to the market with her sister to buy a new sari.\",\n",
+        "    \"Raj told me that he is going to his grandmother's house next month.\",\n",
+        "    \"All the kids were having fun at the party and were eating lots of sweets.\",\n",
+        "    \"My friend has invited me to his birthday party, and I will give him a gift.\",\n",
+        "]\n",
+        "\n",
+        "src_lang, tgt_lang = \"eng_Latn\", \"hin_Deva\"\n",
+        "hi_translations = batch_translate(en_sents, src_lang, tgt_lang, en_indic_model, en_indic_tokenizer, ip)\n",
+        "\n",
+        "print(f\"\\n{src_lang} - {tgt_lang}\")\n",
+        "for input_sentence, translation in zip(en_sents, hi_translations):\n",
+        "    print(f\"{src_lang}: {input_sentence}\")\n",
+        "    print(f\"{tgt_lang}: {translation}\")\n",
+        "\n",
+        "# flush the models to free the GPU memory\n",
+        "del en_indic_tokenizer, en_indic_model"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "OM_1pbPtMpV9"
+      },
+      "source": [
+        "### Indic to English Example"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "PLCEWJKvGG9I",
+        "outputId": "ab9d8726-67c7-490b-ecb3-208df1c0f741"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "\n",
+            "hin_Deva - eng_Latn\n",
+            "hin_Deva: जब मैं छोटा था, मैं हर रोज़ पार्क जाता था।\n",
+            "eng_Latn: When I was young, I used to go to the park every day.\n",
+            "hin_Deva: उसके पास बहुत सारी पुरानी किताबें हैं, जिन्हें उसने अपने दादा-परदादा से विरासत में पाया।\n",
+            "eng_Latn: She has a lot of old books, which she inherited from her grandparents.\n",
+            "hin_Deva: मुझे समझ में नहीं आ रहा कि मैं अपनी समस्या का समाधान कैसे ढूंढूं।\n",
+            "eng_Latn: I don't know how to find a solution to my problem.\n",
+            "hin_Deva: वह बहुत मेहनती और समझदार है, इसलिए उसे सभी अच्छे मार्क्स मिले।\n",
+            "eng_Latn: He is very hardworking and understanding, so he got all the good marks.\n",
+            "hin_Deva: हमने पिछले सप्ताह एक नई फिल्म देखी जो कि बहुत प्रेरणादायक थी।\n",
+            "eng_Latn: We saw a new movie last week that was very inspiring.\n",
+            "hin_Deva: अगर तुम मुझे उस समय पास मिलते, तो हम बाहर खाना खाने चलते।\n",
+            "eng_Latn: If you'd given me a pass at that time, we'd have gone out to eat.\n",
+            "hin_Deva: वह अपनी दीदी के साथ बाजार गयी थी ताकि वह नई साड़ी खरीद सके।\n",
+            "eng_Latn: She had gone to the market with her sister so that she could buy a new sari.\n",
+            "hin_Deva: राज ने मुझसे कहा कि वह अगले महीने अपनी नानी के घर जा रहा है।\n",
+            "eng_Latn: Raj told me that he was going to his grandmother's house next month.\n",
+            "hin_Deva: सभी बच्चे पार्टी में मज़ा कर रहे थे और खूब सारी मिठाइयाँ खा रहे थे।\n",
+            "eng_Latn: All the children were having fun at the party and eating a lot of sweets.\n",
+            "hin_Deva: मेरे मित्र ने मुझे उसके जन्मदिन की पार्टी में बुलाया है, और मैं उसे एक तोहफा दूंगा।\n",
+            "eng_Latn: My friend has invited me to her birthday party, and I'll give her a present.\n"
+          ]
+        }
+      ],
+      "source": [
+        "indic_en_ckpt_dir = \"ai4bharat/indictrans2-indic-en-1B\"  # ai4bharat/indictrans2-indic-en-dist-200M\n",
+        "indic_en_tokenizer, indic_en_model = initialize_model_and_tokenizer(indic_en_ckpt_dir, quantization)\n",
+        "\n",
+        "ip = IndicProcessor(inference=True)\n",
+        "\n",
+        "hi_sents = [\n",
+        "    \"जब मैं छोटा था, मैं हर रोज़ पार्क जाता था।\",\n",
+        "    \"उसके पास बहुत सारी पुरानी किताबें हैं, जिन्हें उसने अपने दादा-परदादा से विरासत में पाया।\",\n",
+        "    \"मुझे समझ में नहीं आ रहा कि मैं अपनी समस्या का समाधान कैसे ढूंढूं।\",\n",
+        "    \"वह बहुत मेहनती और समझदार है, इसलिए उसे सभी अच्छे मार्क्स मिले।\",\n",
+        "    \"हमने पिछले सप्ताह एक नई फिल्म देखी जो कि बहुत प्रेरणादायक थी।\",\n",
+        "    \"अगर तुम मुझे उस समय पास मिलते, तो हम बाहर खाना खाने चलते।\",\n",
+        "    \"वह अपनी दीदी के साथ बाजार गयी थी ताकि वह नई साड़ी खरीद सके।\",\n",
+        "    \"राज ने मुझसे कहा कि वह अगले महीने अपनी नानी के घर जा रहा है।\",\n",
+        "    \"सभी बच्चे पार्टी में मज़ा कर रहे थे और खूब सारी मिठाइयाँ खा रहे थे।\",\n",
+        "    \"मेरे मित्र ने मुझे उसके जन्मदिन की पार्टी में बुलाया है, और मैं उसे एक तोहफा दूंगा।\",\n",
+        "]\n",
+        "src_lang, tgt_lang = \"hin_Deva\", \"eng_Latn\"\n",
+        "en_translations = batch_translate(hi_sents, src_lang, tgt_lang, indic_en_model, indic_en_tokenizer, ip)\n",
+        "\n",
+        "\n",
+        "print(f\"\\n{src_lang} - {tgt_lang}\")\n",
+        "for input_sentence, translation in zip(hi_sents, en_translations):\n",
+        "    print(f\"{src_lang}: {input_sentence}\")\n",
+        "    print(f\"{tgt_lang}: {translation}\")\n",
+        "\n",
+        "# flush the models to free the GPU memory\n",
+        "del indic_en_tokenizer, indic_en_model"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "7VCAkyKBGtnV"
+      },
+      "source": [
+        "### Indic to Indic Example\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "_7TxTTCoKjti",
+        "outputId": "df1a750b-0f32-478d-cfc9-e445f669f3ee"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "\n",
+            "hin_Deva - mar_Deva\n",
+            "hin_Deva: जब मैं छोटा था, मैं हर रोज़ पार्क जाता था।\n",
+            "mar_Deva: मी लहान होतो तेव्हा मी दररोज उद्यानाला जायचे.\n",
+            "hin_Deva: उसके पास बहुत सारी पुरानी किताबें हैं, जिन्हें उसने अपने दादा-परदादा से विरासत में पाया।\n",
+            "mar_Deva: तिच्याकडे बरीच जुनी पुस्तके आहेत, जी तिला तिच्या आजोबांकडून वारशाने मिळाली आहेत.\n",
+            "hin_Deva: मुझे समझ में नहीं आ रहा कि मैं अपनी समस्या का समाधान कैसे ढूंढूं।\n",
+            "mar_Deva: माझ्या समस्येवर तोडगा कसा काढायचा हे मला समजत नाही.\n",
+            "hin_Deva: वह बहुत मेहनती और समझदार है, इसलिए उसे सभी अच्छे मार्क्स मिले।\n",
+            "mar_Deva: तो खूप मेहनती आणि बुद्धिमान आहे, त्यामुळे त्याला सर्व चांगले गुण मिळाले.\n",
+            "hin_Deva: हमने पिछले सप्ताह एक नई फिल्म देखी जो कि बहुत प्रेरणादायक थी।\n",
+            "mar_Deva: आम्ही गेल्या आठवड्यात एक नवीन चित्रपट पाहिला जो खूप प्रेरणादायी होता.\n",
+            "hin_Deva: अगर तुम मुझे उस समय पास मिलते, तो हम बाहर खाना खाने चलते।\n",
+            "mar_Deva: जर तुम्हाला त्या वेळी मला पास मिळाला तर आम्ही बाहेर जेवायला जाऊ.\n",
+            "hin_Deva: वह अपनी दीदी के साथ बाजार गयी थी ताकि वह नई साड़ी खरीद सके।\n",
+            "mar_Deva: ती तिच्या बहिणीसोबत बाजारात गेली होती जेणेकरून ती नवीन साडी खरेदी करू शकेल.\n",
+            "hin_Deva: राज ने मुझसे कहा कि वह अगले महीने अपनी नानी के घर जा रहा है।\n",
+            "mar_Deva: राजने मला सांगितले की तो पुढच्या महिन्यात त्याच्या आजीच्या घरी जात आहे.\n",
+            "hin_Deva: सभी बच्चे पार्टी में मज़ा कर रहे थे और खूब सारी मिठाइयाँ खा रहे थे।\n",
+            "mar_Deva: सर्व मुले पार्टीचा आनंद घेत होती आणि भरपूर मिठाई खात होती.\n",
+            "hin_Deva: मेरे मित्र ने मुझे उसके जन्मदिन की पार्टी में बुलाया है, और मैं उसे एक तोहफा दूंगा।\n",
+            "mar_Deva: माझ्या मित्राने मला त्याच्या वाढदिवसाच्या मेजवानीसाठी आमंत्रित केले आहे आणि मी त्याला भेटवस्तू देईन.\n"
+          ]
+        }
+      ],
+      "source": [
+        "indic_indic_ckpt_dir = \"ai4bharat/indictrans2-indic-indic-1B\"  # ai4bharat/indictrans2-indic-indic-dist-320M\n",
+        "indic_indic_tokenizer, indic_indic_model = initialize_model_and_tokenizer(indic_indic_ckpt_dir, quantization)\n",
+        "\n",
+        "ip = IndicProcessor(inference=True)\n",
+        "\n",
+        "hi_sents = [\n",
+        "    \"जब मैं छोटा था, मैं हर रोज़ पार्क जाता था।\",\n",
+        "    \"उसके पास बहुत सारी पुरानी किताबें हैं, जिन्हें उसने अपने दादा-परदादा से विरासत में पाया।\",\n",
+        "    \"मुझे समझ में नहीं आ रहा कि मैं अपनी समस्या का समाधान कैसे ढूंढूं।\",\n",
+        "    \"वह बहुत मेहनती और समझदार है, इसलिए उसे सभी अच्छे मार्क्स मिले।\",\n",
+        "    \"हमने पिछले सप्ताह एक नई फिल्म देखी जो कि बहुत प्रेरणादायक थी।\",\n",
+        "    \"अगर तुम मुझे उस समय पास मिलते, तो हम बाहर खाना खाने चलते।\",\n",
+        "    \"वह अपनी दीदी के साथ बाजार गयी थी ताकि वह नई साड़ी खरीद सके।\",\n",
+        "    \"राज ने मुझसे कहा कि वह अगले महीने अपनी नानी के घर जा रहा है।\",\n",
+        "    \"सभी बच्चे पार्टी में मज़ा कर रहे थे और खूब सारी मिठाइयाँ खा रहे थे।\",\n",
+        "    \"मेरे मित्र ने मुझे उसके जन्मदिन की पार्टी में बुलाया है, और मैं उसे एक तोहफा दूंगा।\",\n",
+        "]\n",
+        "src_lang, tgt_lang = \"hin_Deva\", \"mar_Deva\"\n",
+        "mr_translations = batch_translate(hi_sents, src_lang, tgt_lang, indic_indic_model, indic_indic_tokenizer, ip)\n",
+        "\n",
+        "print(f\"\\n{src_lang} - {tgt_lang}\")\n",
+        "for input_sentence, translation in zip(hi_sents, mr_translations):\n",
+        "    print(f\"{src_lang}: {input_sentence}\")\n",
+        "    print(f\"{tgt_lang}: {translation}\")\n",
+        "\n",
+        "# flush the models to free the GPU memory\n",
+        "del indic_indic_tokenizer, indic_indic_model"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "uyxXpt--Ma6n"
+      },
+      "outputs": [],
+      "source": []
+    }
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "gpuType": "T4",
+      "provenance": [],
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

IndicTrans2/huggingface_interface/configuration_indictrans.py

@@ -0,0 +1,309 @@
+# coding=utf-8
+# Copyright 2023 The IndicTrans2 Authors and AI4Bharat team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" PyTorch IndicTrans config."""
+
+
+from collections import OrderedDict
+from typing import Any, Mapping, Optional
+
+from transformers import PreTrainedTokenizer
+from transformers.configuration_utils import PretrainedConfig
+from transformers.onnx import OnnxConfig, OnnxSeq2SeqConfigWithPast
+from transformers.onnx.utils import compute_effective_axis_dimension
+from transformers.utils import TensorType, is_torch_available
+
+
+# Copied from transformers.models.m2m_100.configuration_m2m_100.M2M100Config->IndicTrans
+class IndicTransConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`IT2Model`]. It is used to instantiate an
+    IT2 model according to the specified arguments, defining the model architecture. Instantiating a configuration
+    with the defaults will yield a similar configuration to that of the IT2
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+
+    Args:
+        vocab_size (`int`, *optional*, defaults to 50265):
+            Vocabulary size of the IT2 model. Defines the number of different tokens that can be represented by the
+            `inputs_ids` passed when calling [`IT2Model`] or
+        d_model (`int`, *optional*, defaults to 1024):
+            Dimensionality of the layers and the pooler layer.
+        encoder_layers (`int`, *optional*, defaults to 12):
+            Number of encoder layers.
+        decoder_layers (`int`, *optional*, defaults to 12):
+            Number of decoder layers.
+        encoder_attention_heads (`int`, *optional*, defaults to 16):
+            Number of attention heads for each attention layer in the Transformer encoder.
+        decoder_attention_heads (`int`, *optional*, defaults to 16):
+            Number of attention heads for each attention layer in the Transformer decoder.
+        decoder_ffn_dim (`int`, *optional*, defaults to 4096):
+            Dimensionality of the "intermediate" (often named feed-forward) layer in decoder.
+        encoder_ffn_dim (`int`, *optional*, defaults to 4096):
+            Dimensionality of the "intermediate" (often named feed-forward) layer in decoder.
+        activation_function (`str` or `function`, *optional*, defaults to `"gelu"`):
+            The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
+            `"relu"`, `"silu"` and `"gelu_new"` are supported.
+        dropout (`float`, *optional*, defaults to 0.1):
+            The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+        activation_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for activations inside the fully connected layer.
+        classifier_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for classifier.
+        max_position_embeddings (`int`, *optional*, defaults to 1024):
+            The maximum sequence length that this model might ever be used with. Typically set this to something large
+            just in case (e.g., 512 or 1024 or 2048).
+        init_std (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        encoder_layerdrop (`float`, *optional*, defaults to 0.0):
+            The LayerDrop probability for the encoder. See the [LayerDrop paper](see https://arxiv.org/abs/1909.11556)
+            for more details.
+        decoder_layerdrop (`float`, *optional*, defaults to 0.0):
+            The LayerDrop probability for the decoder. See the [LayerDrop paper](see https://arxiv.org/abs/1909.11556)
+            for more details.
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Whether or not the model should return the last key/values attentions (not used by all models).
+    ```"""
+    model_type = "IndicTrans"
+    keys_to_ignore_at_inference = ["past_key_values"]
+    attribute_map = {
+        "num_attention_heads": "encoder_attention_heads",
+        "hidden_size": "d_model",
+    }
+
+    def __init__(
+        self,
+        encoder_vocab_size=None,
+        decoder_vocab_size=None,
+        encoder_embed_dim=512,
+        decoder_embed_dim=512,
+        max_source_positions=210,
+        max_target_positions=210,
+        encoder_layers=6,
+        encoder_ffn_dim=2048,
+        encoder_attention_heads=8,
+        decoder_layers=6,
+        decoder_ffn_dim=2048,
+        decoder_attention_heads=8,
+        encoder_layerdrop=0.00,
+        decoder_layerdrop=0.00,
+        use_cache=True,
+        is_encoder_decoder=True,
+        activation_function="relu",
+        encoder_normalize_before=False,
+        decoder_normalize_before=False,
+        layernorm_embedding=False,
+        share_decoder_input_output_embed=False,
+        dropout=0.1,
+        attention_dropout=0.0,
+        activation_dropout=0.0,
+        init_std=0.02,
+        scale_embedding=True,
+        decoder_start_token_id=2,
+        pad_token_id=1,
+        bos_token_id=0,
+        eos_token_id=2,
+        attn_implementation="eager",
+        **kwargs,
+    ):
+        self.encoder_vocab_size = encoder_vocab_size
+        self.decoder_vocab_size = decoder_vocab_size
+        self.encoder_normalize_before = encoder_normalize_before
+        self.decoder_normalize_before = decoder_normalize_before
+        self.layernorm_embedding = layernorm_embedding
+        self.max_source_positions = max_source_positions
+        self.max_target_positions = max_target_positions
+        self.encoder_embed_dim = encoder_embed_dim
+        self.decoder_embed_dim = decoder_embed_dim
+        self.encoder_ffn_dim = encoder_ffn_dim
+        self.encoder_layers = encoder_layers
+        self.encoder_attention_heads = encoder_attention_heads
+        self.decoder_ffn_dim = decoder_ffn_dim
+        self.decoder_layers = decoder_layers
+        self.decoder_attention_heads = decoder_attention_heads
+        self.dropout = dropout
+        self.attention_dropout = attention_dropout
+        self.activation_dropout = activation_dropout
+        self.activation_function = activation_function
+        self.init_std = init_std
+        self.encoder_layerdrop = encoder_layerdrop
+        self.decoder_layerdrop = decoder_layerdrop
+        self.use_cache = use_cache
+        self.num_hidden_layers = encoder_layers
+        self.scale_embedding = scale_embedding
+        self.share_decoder_input_output_embed = share_decoder_input_output_embed
+        self.attn_implementation = attn_implementation
+        
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            is_encoder_decoder=is_encoder_decoder,
+            decoder_start_token_id=decoder_start_token_id,
+            **kwargs,
+        )
+
+
+class IndicTransOnnxConfig(OnnxSeq2SeqConfigWithPast):
+    @property
+    def inputs(self) -> Mapping[str, Mapping[int, str]]:
+        common_inputs = OrderedDict(
+            [
+                ("input_ids", {0: "batch", 1: "encoder_sequence"}),
+                ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
+            ]
+        )
+
+        if self.use_past:
+            common_inputs["decoder_input_ids"] = {0: "batch"}
+            common_inputs["decoder_attention_mask"] = {
+                0: "batch",
+                1: "past_decoder_sequence + sequence",
+            }
+        else:
+            common_inputs["decoder_input_ids"] = {0: "batch", 1: "decoder_sequence"}
+            common_inputs["decoder_attention_mask"] = {
+                0: "batch",
+                1: "decoder_sequence",
+            }
+
+        if self.use_past:
+            self.fill_with_past_key_values_(common_inputs, direction="inputs")
+        return common_inputs
+
+    # Copied from BartOnnxConfig._generate_dummy_inputs_for_sequence_classification_and_question_answering
+    # A better name would be _generate_dummy_inputs_for_encoder_and_decoder because sequence classification and question
+    # answering are not supported for IT2, but this name is preserved to be able to check that the copy matches what
+    # was done for BART so that it can be updated if need be.
+    def _generate_dummy_inputs_for_sequence_classification_and_question_answering(
+        self,
+        tokenizer: PreTrainedTokenizer,
+        batch_size: int = -1,
+        seq_length: int = -1,
+        is_pair: bool = False,
+        framework: Optional[TensorType] = None,
+    ) -> Mapping[str, Any]:
+        # Copied from OnnxConfig.generate_dummy_inputs
+        # Did not use super(OnnxConfigWithPast, self).generate_dummy_inputs for code clarity.
+        # If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX
+        batch_size = compute_effective_axis_dimension(
+            batch_size,
+            fixed_dimension=OnnxConfig.default_fixed_batch,
+            num_token_to_add=0,
+        )
+
+        # If dynamic axis (-1) we forward with a fixed dimension of 8 tokens to avoid optimizations made by ONNX
+        token_to_add = tokenizer.num_special_tokens_to_add(is_pair)
+        seq_length = compute_effective_axis_dimension(
+            seq_length,
+            fixed_dimension=OnnxConfig.default_fixed_sequence,
+            num_token_to_add=token_to_add,
+        )
+
+        # Generate dummy inputs according to compute batch and sequence
+        dummy_input = [" ".join([tokenizer.unk_token]) * seq_length] * batch_size
+        common_inputs = dict(tokenizer(dummy_input, return_tensors=framework))
+        return common_inputs
+
+    # Copied from transformers.models.bart.configuration_bart.BartOnnxConfig._generate_dummy_inputs_for_default_and_seq2seq_lm
+    def _generate_dummy_inputs_for_default_and_seq2seq_lm(
+        self,
+        tokenizer: PreTrainedTokenizer,
+        batch_size: int = -1,
+        seq_length: int = -1,
+        is_pair: bool = False,
+        framework: Optional[TensorType] = None,
+    ) -> Mapping[str, Any]:
+        encoder_inputs = self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
+            tokenizer, batch_size, seq_length, is_pair, framework
+        )
+
+        # Generate decoder inputs
+        decoder_seq_length = seq_length if not self.use_past else 1
+        decoder_inputs = self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
+            tokenizer, batch_size, decoder_seq_length, is_pair, framework
+        )
+        decoder_inputs = {
+            f"decoder_{name}": tensor for name, tensor in decoder_inputs.items()
+        }
+        common_inputs = dict(**encoder_inputs, **decoder_inputs)
+
+        if self.use_past:
+            if not is_torch_available():
+                raise ValueError(
+                    "Cannot generate dummy past_keys inputs without PyTorch installed."
+                )
+            else:
+                import torch
+            batch, encoder_seq_length = common_inputs["input_ids"].shape
+            decoder_seq_length = common_inputs["decoder_input_ids"].shape[1]
+            (
+                num_encoder_attention_heads,
+                num_decoder_attention_heads,
+            ) = self.num_attention_heads
+            encoder_shape = (
+                batch,
+                num_encoder_attention_heads,
+                encoder_seq_length,
+                self._config.hidden_size // num_encoder_attention_heads,
+            )
+            decoder_past_length = decoder_seq_length + 3
+            decoder_shape = (
+                batch,
+                num_decoder_attention_heads,
+                decoder_past_length,
+                self._config.hidden_size // num_decoder_attention_heads,
+            )
+
+            common_inputs["decoder_attention_mask"] = torch.cat(
+                [
+                    common_inputs["decoder_attention_mask"],
+                    torch.ones(batch, decoder_past_length),
+                ],
+                dim=1,
+            )
+
+            common_inputs["past_key_values"] = []
+            # If the number of encoder and decoder layers are present in the model configuration, both are considered
+            num_encoder_layers, num_decoder_layers = self.num_layers
+            min_num_layers = min(num_encoder_layers, num_decoder_layers)
+            max_num_layers = (
+                max(num_encoder_layers, num_decoder_layers) - min_num_layers
+            )
+            remaining_side_name = (
+                "encoder" if num_encoder_layers > num_decoder_layers else "decoder"
+            )
+
+            for _ in range(min_num_layers):
+                common_inputs["past_key_values"].append(
+                    (
+                        torch.zeros(decoder_shape),
+                        torch.zeros(decoder_shape),
+                        torch.zeros(encoder_shape),
+                        torch.zeros(encoder_shape),
+                    )
+                )
+            # TODO: test this.
+            shape = encoder_shape if remaining_side_name == "encoder" else decoder_shape
+            for _ in range(min_num_layers, max_num_layers):
+                common_inputs["past_key_values"].append(
+                    (torch.zeros(shape), torch.zeros(shape))
+                )
+        return common_inputs
+
+    generate_dummy_inputs = _generate_dummy_inputs_for_default_and_seq2seq_lm

IndicTrans2/huggingface_interface/convert_indictrans_checkpoint_to_pytorch.py

@@ -0,0 +1,107 @@
+# Copyright 2021 The Fairseq Authors and The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+
+import torch
+import torch.nn as nn
+
+from configuration_indictrans import IndicTransConfig
+from modeling_indictrans import IndicTransForConditionalGeneration
+
+
+def remove_ignore_keys_(state_dict):
+    ignore_keys = [
+        "encoder.version",
+        "decoder.version",
+        "model.encoder.version",
+        "model.decoder.version",
+        "_float_tensor",
+        "encoder.embed_positions._float_tensor",
+        "decoder.embed_positions._float_tensor",
+    ]
+    for k in ignore_keys:
+        state_dict.pop(k, None)
+
+
+def make_linear_from_emb(emb):
+    vocab_size, emb_size = emb.shape
+    lin_layer = nn.Linear(vocab_size, emb_size, bias=False)
+    lin_layer.weight.data = emb.data
+    return lin_layer
+
+
+def convert_fairseq_IT2_checkpoint_from_disk(checkpoint_path):
+    model = torch.load(checkpoint_path, map_location="cpu")
+    args = model["args"] or model["cfg"]["model"]
+    state_dict = model["model"]
+    remove_ignore_keys_(state_dict)
+    encoder_vocab_size = state_dict["encoder.embed_tokens.weight"].shape[0]
+    decoder_vocab_size = state_dict["decoder.embed_tokens.weight"].shape[0]
+
+    config = IndicTransConfig(
+        encoder_vocab_size=encoder_vocab_size,
+        decoder_vocab_size=decoder_vocab_size,
+        max_source_positions=args.max_source_positions,
+        max_target_positions=args.max_target_positions,
+        encoder_layers=args.encoder_layers,
+        decoder_layers=args.decoder_layers,
+        layernorm_embedding=args.layernorm_embedding,
+        encoder_normalize_before=args.encoder_normalize_before,
+        decoder_normalize_before=args.decoder_normalize_before,
+        encoder_attention_heads=args.encoder_attention_heads,
+        decoder_attention_heads=args.decoder_attention_heads,
+        encoder_ffn_dim=args.encoder_ffn_embed_dim,
+        decoder_ffn_dim=args.decoder_ffn_embed_dim,
+        encoder_embed_dim=args.encoder_embed_dim,
+        decoder_embed_dim=args.decoder_embed_dim,
+        encoder_layerdrop=args.encoder_layerdrop,
+        decoder_layerdrop=args.decoder_layerdrop,
+        dropout=args.dropout,
+        attention_dropout=args.attention_dropout,
+        activation_dropout=args.activation_dropout,
+        activation_function=args.activation_fn,
+        share_decoder_input_output_embed=args.share_decoder_input_output_embed,
+        scale_embedding=not args.no_scale_embedding,
+    )
+
+    model = IndicTransForConditionalGeneration(config)
+    model.model.load_state_dict(state_dict, strict=False)
+    if not args.share_decoder_input_output_embed:
+        model.lm_head = make_linear_from_emb(
+            state_dict["decoder.output_projection.weight"]
+        )
+    print(model)
+    return model
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    # Required parameters
+    parser.add_argument(
+        "--fairseq_path",
+        default="indic-en/model/checkpoint_best.pt",
+        type=str,
+        help="path to a model.pt on local filesystem.",
+    )
+    parser.add_argument(
+        "--pytorch_dump_folder_path",
+        default="indic-en/hf_model",
+        type=str,
+        help="Path to the output PyTorch model.",
+    )
+
+    args = parser.parse_args()
+    model = convert_fairseq_IT2_checkpoint_from_disk(args.fairseq_path)
+    model.save_pretrained(args.pytorch_dump_folder_path)

IndicTrans2/huggingface_interface/example.py

@@ -0,0 +1,275 @@
+import sys
+import torch
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, BitsAndBytesConfig
+from transformers.utils import is_flash_attn_2_available, is_flash_attn_greater_or_equal_2_10
+from IndicTransToolkit import IndicProcessor
+from mosestokenizer import MosesSentenceSplitter
+from nltk import sent_tokenize
+from indicnlp.tokenize.sentence_tokenize import sentence_split, DELIM_PAT_NO_DANDA
+
+
+en_indic_ckpt_dir = "ai4bharat/indictrans2-en-indic-1B"  # ai4bharat/indictrans2-en-indic-dist-200M
+indic_en_ckpt_dir = "ai4bharat/indictrans2-indic-en-1B"  # ai4bharat/indictrans2-indic-en-dist-200M
+indic_indic_ckpt_dir = (
+    "ai4bharat/indictrans2-indic-indic-dist-320M"  # ai4bharat/indictrans2-indic-indic-dist-320M
+)
+BATCH_SIZE = 4
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+if len(sys.argv) > 1:
+    quantization = sys.argv[1]
+    attn_implementation = sys.argv[2]
+else:
+    quantization = ""
+    attn_implementation = "eager"
+
+
+# FLORES language code mapping to 2 letter ISO language code for compatibility
+# with Indic NLP Library (https://github.com/anoopkunchukuttan/indic_nlp_library)
+flores_codes = {
+    "asm_Beng": "as",
+    "awa_Deva": "hi",
+    "ben_Beng": "bn",
+    "bho_Deva": "hi",
+    "brx_Deva": "hi",
+    "doi_Deva": "hi",
+    "eng_Latn": "en",
+    "gom_Deva": "kK",
+    "guj_Gujr": "gu",
+    "hin_Deva": "hi",
+    "hne_Deva": "hi",
+    "kan_Knda": "kn",
+    "kas_Arab": "ur",
+    "kas_Deva": "hi",
+    "kha_Latn": "en",
+    "lus_Latn": "en",
+    "mag_Deva": "hi",
+    "mai_Deva": "hi",
+    "mal_Mlym": "ml",
+    "mar_Deva": "mr",
+    "mni_Beng": "bn",
+    "mni_Mtei": "hi",
+    "npi_Deva": "ne",
+    "ory_Orya": "or",
+    "pan_Guru": "pa",
+    "san_Deva": "hi",
+    "sat_Olck": "or",
+    "snd_Arab": "ur",
+    "snd_Deva": "hi",
+    "tam_Taml": "ta",
+    "tel_Telu": "te",
+    "urd_Arab": "ur",
+}
+
+
+def split_sentences(input_text, lang):
+    if lang == "eng_Latn":
+        input_sentences = sent_tokenize(input_text)
+        with MosesSentenceSplitter(flores_codes[lang]) as splitter:
+            sents_moses = splitter([input_text])
+        sents_nltk = sent_tokenize(input_text)
+        if len(sents_nltk) < len(sents_moses):
+            input_sentences = sents_nltk
+        else:
+            input_sentences = sents_moses
+        input_sentences = [sent.replace("\xad", "") for sent in input_sentences]
+    else:
+        input_sentences = sentence_split(
+            input_text, lang=flores_codes[lang], delim_pat=DELIM_PAT_NO_DANDA
+        )
+    return input_sentences
+
+
+def initialize_model_and_tokenizer(ckpt_dir, quantization, attn_implementation):
+    if quantization == "4-bit":
+        qconfig = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_compute_dtype=torch.bfloat16,
+        )
+    elif quantization == "8-bit":
+        qconfig = BitsAndBytesConfig(
+            load_in_8bit=True,
+            bnb_8bit_use_double_quant=True,
+            bnb_8bit_compute_dtype=torch.bfloat16,
+        )
+    else:
+        qconfig = None 
+        
+    if attn_implementation == "flash_attention_2":
+        if is_flash_attn_2_available() and is_flash_attn_greater_or_equal_2_10():
+            attn_implementation = "flash_attention_2"
+        else:
+            attn_implementation = "eager"
+
+    tokenizer = AutoTokenizer.from_pretrained(ckpt_dir, trust_remote_code=True)
+    model = AutoModelForSeq2SeqLM.from_pretrained(
+        ckpt_dir,
+        trust_remote_code=True,
+        attn_implementation=attn_implementation,
+        low_cpu_mem_usage=True,
+        quantization_config=qconfig,
+    )
+
+    if qconfig == None:
+        model = model.to(DEVICE)
+        model.half()
+
+    model.eval()
+
+    return tokenizer, model
+
+
+def batch_translate(input_sentences, src_lang, tgt_lang, model, tokenizer, ip):
+    translations = []
+    for i in range(0, len(input_sentences), BATCH_SIZE):
+        batch = input_sentences[i : i + BATCH_SIZE]
+
+        # Preprocess the batch and extract entity mappings
+        batch = ip.preprocess_batch(batch, src_lang=src_lang, tgt_lang=tgt_lang)
+
+        # Tokenize the batch and generate input encodings
+        inputs = tokenizer(
+            batch,
+            truncation=True,
+            padding="longest",
+            return_tensors="pt",
+            return_attention_mask=True,
+        ).to(DEVICE)
+
+        # Generate translations using the model
+        with torch.no_grad():
+            generated_tokens = model.generate(
+                **inputs,
+                use_cache=True,
+                min_length=0,
+                max_length=256,
+                num_beams=5,
+                num_return_sequences=1,
+            )
+
+        # Decode the generated tokens into text
+        with tokenizer.as_target_tokenizer():
+            generated_tokens = tokenizer.batch_decode(
+                generated_tokens.detach().cpu().tolist(),
+                skip_special_tokens=True,
+                clean_up_tokenization_spaces=True,
+            )
+
+        # Postprocess the translations, including entity replacement
+        translations += ip.postprocess_batch(generated_tokens, lang=tgt_lang)
+
+        del inputs
+        torch.cuda.empty_cache()
+
+    return translations
+
+
+def translate_paragraph(input_text, src_lang, tgt_lang, model, tokenizer, ip):
+    input_sentences = split_sentences(input_text, src_lang)
+    translated_text = batch_translate(input_sentences, src_lang, tgt_lang, model, tokenizer, ip)
+    return " ".join(translated_text)
+
+
+ip = IndicProcessor(inference=True)
+
+en_indic_tokenizer, en_indic_model = initialize_model_and_tokenizer(
+    en_indic_ckpt_dir, quantization, attn_implementation
+)
+
+indic_en_tokenizer, indic_en_model = initialize_model_and_tokenizer(
+    indic_en_ckpt_dir, quantization, attn_implementation
+)
+
+indic_indic_tokenizer, indic_indic_model = initialize_model_and_tokenizer(
+    indic_indic_ckpt_dir, quantization, attn_implementation
+)
+
+# ---------------------------------------------------------------------------
+#                              Hindi to English
+# ---------------------------------------------------------------------------
+hi_sents = [
+    "जब मैं छोटा था, मैं हर रोज़ पार्क जाता था।",
+    "उसके पास बहुत सारी पुरानी किताबें हैं, जिन्हें उसने अपने दादा-परदादा से विरासत में पाया।",
+    "मुझे समझ में नहीं आ रहा कि मैं अपनी समस्या का समाधान कैसे ढूंढूं।",
+    "वह बहुत मेहनती और समझदार है, इसलिए उसे सभी अच्छे मार्क्स मिले।",
+    "हमने पिछले सप्ताह एक नई फिल्म देखी जो कि बहुत प्रेरणादायक थी।",
+    "अगर तुम मुझे उस समय पास मिलते, तो हम बाहर खाना खाने चलते।",
+    "वह अपनी दीदी के साथ बाजार गयी थी ताकि वह नई साड़ी खरीद सके।",
+    "राज ने मुझसे कहा कि वह अगले महीने अपनी नानी के घर जा रहा है।",
+    "सभी बच्चे पार्टी में मज़ा कर रहे थे और खूब सारी मिठाइयाँ खा रहे थे।",
+    "मेरे मित्र ने मुझे उसके जन्मदिन की पार्टी में बुलाया है, और मैं उसे एक तोहफा दूंगा।",
+]
+src_lang, tgt_lang = "hin_Deva", "eng_Latn"
+en_translations = batch_translate(
+    hi_sents, src_lang, tgt_lang, indic_en_model, indic_en_tokenizer, ip
+)
+
+print(f"\n{src_lang} - {tgt_lang}")
+for input_sentence, translation in zip(hi_sents, en_translations):
+    print(f"{src_lang}: {input_sentence}")
+    print(f"{tgt_lang}: {translation}")
+
+
+# ---------------------------------------------------------------------------
+#                              English to Hindi
+# ---------------------------------------------------------------------------
+en_sents = [
+    "When I was young, I used to go to the park every day.",
+    "He has many old books, which he inherited from his ancestors.",
+    "I can't figure out how to solve my problem.",
+    "She is very hardworking and intelligent, which is why she got all the good marks.",
+    "We watched a new movie last week, which was very inspiring.",
+    "If you had met me at that time, we would have gone out to eat.",
+    "She went to the market with her sister to buy a new sari.",
+    "Raj told me that he is going to his grandmother's house next month.",
+    "All the kids were having fun at the party and were eating lots of sweets.",
+    "My friend has invited me to his birthday party, and I will give him a gift.",
+]
+src_lang, tgt_lang = "eng_Latn", "hin_Deva"
+hi_translations = batch_translate(
+    en_sents, src_lang, tgt_lang, en_indic_model, en_indic_tokenizer, ip
+)
+
+print(f"\n{src_lang} - {tgt_lang}")
+for input_sentence, translation in zip(en_sents, hi_translations):
+    print(f"{src_lang}: {input_sentence}")
+    print(f"{tgt_lang}: {translation}")
+
+
+# ---------------------------------------------------------------------------
+#                              Hindi to Marathi
+# ---------------------------------------------------------------------------
+hi_sents = [
+    "��ब मैं छोटा था, मैं हर रोज़ पार्क जाता था।",
+    "उसके पास बहुत सारी पुरानी किताबें हैं, जिन्हें उसने अपने दादा-परदादा से विरासत में पाया।",
+    "मुझे समझ में नहीं आ रहा कि मैं अपनी समस्या का समाधान कैसे ढूंढूं।",
+    "वह बहुत मेहनती और समझदार है, इसलिए उसे सभी अच्छे मार्क्स मिले।",
+    "हमने पिछले सप्ताह एक नई फिल्म देखी जो कि बहुत प्रेरणादायक थी।",
+    "अगर तुम मुझे उस समय पास मिलते, तो हम बाहर खाना खाने चलते।",
+    "वह अपनी दीदी के साथ बाजार गयी थी ताकि वह नई साड़ी खरीद सके।",
+    "राज ने मुझसे कहा कि वह अगले महीने अपनी नानी के घर जा रहा है।",
+    "सभी बच्चे पार्टी में मज़ा कर रहे थे और खूब सारी मिठाइयाँ खा रहे थे।",
+    "मेरे मित्र ने मुझे उसके जन्मदिन की पार्टी में बुलाया है, और मैं उसे एक तोहफा दूंगा।",
+]
+src_lang, tgt_lang = "hin_Deva", "mar_Deva"
+mr_translations = batch_translate(
+    hi_sents, src_lang, tgt_lang, indic_indic_model, indic_indic_tokenizer, ip
+)
+
+print(f"\n{src_lang} - {tgt_lang}")
+for input_sentence, translation in zip(hi_sents, mr_translations):
+    print(f"{src_lang}: {input_sentence}")
+    print(f"{tgt_lang}: {translation}")
+
+
+# ---------------------------------------------------------------------------
+#                             Paragraph translation
+# ---------------------------------------------------------------------------
+src_lang, tgt_lang = "hin_Deva", "eng_Latn"
+hi_text = "यहाँ एक पाराग्राफ है जो हिंदी में लिखा गया है। हिंदी एक सुंदर भाषा है और भारत की राष्ट्रीय भाषा है। इसका विकास विभिन्न कालों में हुआ है और यह विशेषतः भारतीय उपमहाद्वीप में बोली जाती है। हिंदी भाषा का साहित्य, संस्कृति और इतिहास भी बहुत गर्वनीय है।"
+en_translated_text = translate_paragraph(
+    hi_text, src_lang, tgt_lang, indic_en_model, indic_en_tokenizer, ip
+)
+print(f"{src_lang}: {hi_text}")
+print(f"{tgt_lang}: {en_translated_text}")

IndicTrans2/huggingface_interface/install.sh

@@ -0,0 +1,49 @@
+#/bin/bash
+
+root_dir=$(pwd)
+echo "Setting up the environment in the $root_dir"
+
+# --------------------------------------------------------------
+#          create and activate the virtual environment
+# --------------------------------------------------------------
+echo "Creating a virtual environment with python3"
+conda create -n itv2_hf python=3.9 -y
+conda activate itv2_hf
+
+echo "Installing all the dependencies"
+conda install pip
+python3 -m pip install --upgrade pip
+
+
+# --------------------------------------------------------------
+#                   PyTorch Installation
+# --------------------------------------------------------------
+python3 -m pip install torch --extra-index-url https://download.pytorch.org/whl/cu118
+
+
+# --------------------------------------------------------------
+#               Install additional utility packages
+# --------------------------------------------------------------
+python3 -m pip install nltk sacremoses pandas regex mock transformers>=4.33.2 mosestokenizer
+python3 -c "import nltk; nltk.download('punkt')"
+python3 -m pip install bitsandbytes scipy accelerate datasets flash-attn>=2.1
+
+
+# --------------------------------------------------------------
+#               Sentencepiece for tokenization
+# --------------------------------------------------------------
+# build the cpp binaries from the source repo in order to use the command line utility
+# source repo: https://github.com/google/sentencepiece
+python3 -m pip install sentencepiece
+
+
+# -----------------------------------------------------------------
+#       Install IndicTrans2 tokenizer and its dependencies
+# -----------------------------------------------------------------
+git clone https://github.com/VarunGumma/IndicTransToolkit
+cd IndicTransToolkit
+python3 -m pip install --editable ./
+cd $root_dir
+
+
+echo "Setup completed!"

IndicTrans2/huggingface_interface/modeling_indictrans.py

@@ -0,0 +1,1801 @@
+# coding=utf-8
+# Copyright 2023 The IndicTrans2 Authors and AI4Bharat team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" PyTorch IndicTrans model."""
+
+
+import math
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+from transformers.activations import ACT2FN
+
+from transformers.modeling_attn_mask_utils import (
+    _prepare_4d_attention_mask,
+    _prepare_4d_attention_mask_for_sdpa,
+    _prepare_4d_causal_attention_mask,
+    _prepare_4d_causal_attention_mask_for_sdpa,
+)
+
+from transformers.integrations.deepspeed import is_deepspeed_zero3_enabled
+from transformers.modeling_outputs import (
+    BaseModelOutput,
+    BaseModelOutputWithPastAndCrossAttentions,
+    Seq2SeqLMOutput,
+    Seq2SeqModelOutput
+)
+
+from transformers.utils import (
+    logging,
+    is_flash_attn_2_available,
+    is_flash_attn_greater_or_equal_2_10,
+)
+ 
+from transformers.modeling_utils import PreTrainedModel
+
+from .configuration_indictrans import IndicTransConfig
+
+
+logger = logging.get_logger(__name__)
+
+INDICTRANS_PRETRAINED_MODEL_ARCHIVE_LIST = [""]
+
+try:
+    if is_flash_attn_2_available():
+        from flash_attn import flash_attn_func, flash_attn_varlen_func
+        from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input  # noqa    
+except:
+    pass
+
+
+# Copied from transformers.models.llama.modeling_llama._get_unpad_data
+def _get_unpad_data(attention_mask):
+    seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
+    indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
+    max_seqlen_in_batch = seqlens_in_batch.max().item()
+    cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
+    return (
+        indices,
+        cu_seqlens,
+        max_seqlen_in_batch,
+    )
+
+
+# Copied from transformers.models.bart.modeling_bart.shift_tokens_right
+def shift_tokens_right(
+    input_ids: torch.Tensor, pad_token_id: int, decoder_start_token_id: int
+):
+    """
+    Shift input ids one token to the right.
+    """
+    shifted_input_ids = input_ids.new_zeros(input_ids.shape)
+    shifted_input_ids[:, 1:] = input_ids[:, :-1].clone()
+    shifted_input_ids[:, 0] = decoder_start_token_id
+
+    if pad_token_id is None:
+        raise ValueError("self.model.config.pad_token_id has to be defined.")
+    # replace possible -100 values in labels by `pad_token_id`
+    shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id)
+
+    return shifted_input_ids
+
+
+def create_position_ids_from_input_ids(
+    input_ids, padding_idx, past_key_values_length=0
+):
+    """
+    Replace non-padding symbols with their position numbers. Position numbers begin at padding_idx+1. Padding symbols
+    are ignored. This is modified from fairseq's `utils.make_positions`.
+    """
+    # The series of casts and type-conversions here are carefully balanced to both work with ONNX export and XLA.
+    mask = input_ids.ne(padding_idx).int()
+    incremental_indices = (
+        torch.cumsum(mask, dim=1).type_as(mask) + past_key_values_length
+    ) * mask
+    return incremental_indices.long() + padding_idx
+
+
+# Copied from transformers.models.m2m_100.modeling_m2m_100.M2M100SinusoidalPositionalEmbedding->IndicTrans
+class IndicTransSinusoidalPositionalEmbedding(nn.Module):
+    """This module produces sinusoidal positional embeddings of any length."""
+
+    def __init__(
+        self, num_positions: int, embedding_dim: int, padding_idx: Optional[int] = None
+    ):
+        super().__init__()
+        self.offset = 2
+        self.embedding_dim = embedding_dim
+        self.padding_idx = padding_idx
+        self.make_weights(num_positions + self.offset, embedding_dim, padding_idx)
+
+    def make_weights(
+        self, num_embeddings: int, embedding_dim: int, padding_idx: Optional[int] = None
+    ):
+        emb_weights = self.get_embedding(num_embeddings, embedding_dim, padding_idx)
+        if hasattr(self, "weights"):
+            # in forward put the weights on the correct dtype and device of the param
+            emb_weights = emb_weights.to(
+                dtype=self.weights.dtype, device=self.weights.device
+            )
+
+        self.register_buffer("weights", emb_weights, persistent=False)
+
+    @staticmethod
+    def get_embedding(
+        num_embeddings: int, embedding_dim: int, padding_idx: Optional[int] = None
+    ):
+        """
+        Build sinusoidal embeddings.
+
+        This matches the implementation in tensor2tensor, but differs slightly from the description in Section 3.5 of
+        "Attention Is All You Need".
+        """
+        half_dim = embedding_dim // 2
+        emb = math.log(10000) / (half_dim - 1)
+        emb = torch.exp(torch.arange(half_dim, dtype=torch.float) * -emb)
+        emb = torch.arange(num_embeddings, dtype=torch.float).unsqueeze(
+            1
+        ) * emb.unsqueeze(0)
+        emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1).view(
+            num_embeddings, -1
+        )
+        if embedding_dim % 2 == 1:
+            # zero pad
+            emb = torch.cat([emb, torch.zeros(num_embeddings, 1)], dim=1)
+        if padding_idx is not None:
+            emb[padding_idx, :] = 0
+
+        return emb.to(torch.get_default_dtype())
+
+    @torch.no_grad()
+    def forward(
+        self,
+        input_ids: torch.Tensor = None,
+        inputs_embeds: torch.Tensor = None,
+        past_key_values_length: int = 0,
+    ):
+        if input_ids is not None:
+            bsz, seq_len = input_ids.size()
+            # Create the position ids from the input token ids. Any padded tokens remain padded.
+            position_ids = create_position_ids_from_input_ids(
+                input_ids, self.padding_idx, past_key_values_length
+            ).to(input_ids.device)
+        else:
+            bsz, seq_len = inputs_embeds.size()[:-1]
+            position_ids = self.create_position_ids_from_inputs_embeds(
+                inputs_embeds, past_key_values_length
+            )
+
+        # expand embeddings if needed
+        max_pos = self.padding_idx + 1 + seq_len + past_key_values_length
+        if max_pos > self.weights.size(0):
+            self.make_weights(
+                max_pos + self.offset, self.embedding_dim, self.padding_idx
+            )
+
+        return (
+            self.weights.index_select(0, position_ids.view(-1))
+            .view(bsz, seq_len, self.weights.shape[-1])
+            .detach()
+        )
+
+    def create_position_ids_from_inputs_embeds(
+        self, inputs_embeds, past_key_values_length
+    ):
+        """
+        We are provided embeddings directly. We cannot infer which are padded so just generate sequential position ids.
+
+        Args:
+            inputs_embeds: torch.Tensor
+
+        Returns: torch.Tensor
+        """
+        input_shape = inputs_embeds.size()[:-1]
+        sequence_length = input_shape[1]
+
+        position_ids = torch.arange(
+            self.padding_idx + 1,
+            sequence_length + self.padding_idx + 1,
+            dtype=torch.long,
+            device=inputs_embeds.device,
+        )
+        return (
+            position_ids.unsqueeze(0).expand(input_shape).contiguous()
+            + past_key_values_length
+        )
+
+
+# Copied from transformers.models.bart.modeling_bart.BartAttention with Bart->IndicTrans
+class IndicTransAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(
+        self,
+        embed_dim: int,
+        num_heads: int,
+        dropout: float = 0.0,
+        is_decoder: bool = False,
+        bias: bool = True,
+        is_causal: bool = False,
+        config: Optional[IndicTransConfig] = None,
+    ):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.dropout = dropout
+        self.head_dim = embed_dim // num_heads
+        self.config = config
+
+        if (self.head_dim * num_heads) != self.embed_dim:
+            raise ValueError(
+                f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}"
+                f" and `num_heads`: {num_heads})."
+            )
+        self.scaling = self.head_dim**-0.5
+        self.is_decoder = is_decoder
+        self.is_causal = is_causal
+
+        self.k_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+        self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+        self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+
+    def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
+        return (
+            tensor.view(bsz, seq_len, self.num_heads, self.head_dim)
+            .transpose(1, 2)
+            .contiguous()
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        key_value_states: Optional[torch.Tensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        layer_head_mask: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        """Input shape: Batch x Time x Channel"""
+
+        # if key_value_states are provided this layer is used as a cross-attention layer
+        # for the decoder
+        is_cross_attention = key_value_states is not None
+
+        bsz, tgt_len, _ = hidden_states.size()
+
+        # get query proj
+        query_states = self.q_proj(hidden_states) * self.scaling
+        # get key, value proj
+        # `past_key_value[0].shape[2] == key_value_states.shape[1]`
+        # is checking that the `sequence_length` of the `past_key_value` is the same as
+        # the provided `key_value_states` to support prefix tuning
+        if (
+            is_cross_attention
+            and past_key_value is not None
+            and past_key_value[0].shape[2] == key_value_states.shape[1]
+        ):
+            # reuse k,v, cross_attentions
+            key_states = past_key_value[0]
+            value_states = past_key_value[1]
+        elif is_cross_attention:
+            # cross_attentions
+            key_states = self._shape(self.k_proj(key_value_states), -1, bsz)
+            value_states = self._shape(self.v_proj(key_value_states), -1, bsz)
+        elif past_key_value is not None:
+            # reuse k, v, self_attention
+            key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
+            value_states = self._shape(self.v_proj(hidden_states), -1, bsz)
+            key_states = torch.cat([past_key_value[0], key_states], dim=2)
+            value_states = torch.cat([past_key_value[1], value_states], dim=2)
+        else:
+            # self_attention
+            key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
+            value_states = self._shape(self.v_proj(hidden_states), -1, bsz)
+
+        if self.is_decoder:
+            # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states.
+            # Further calls to cross_attention layer can then reuse all cross-attention
+            # key/value_states (first "if" case)
+            # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of
+            # all previous decoder key/value_states. Further calls to uni-directional self-attention
+            # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case)
+            # if encoder bi-directional self-attention `past_key_value` is always `None`
+            past_key_value = (key_states, value_states)
+
+        proj_shape = (bsz * self.num_heads, -1, self.head_dim)
+        query_states = self._shape(query_states, tgt_len, bsz).view(*proj_shape)
+        key_states = key_states.reshape(*proj_shape)
+        value_states = value_states.reshape(*proj_shape)
+
+        src_len = key_states.size(1)
+        attn_weights = torch.bmm(query_states, key_states.transpose(1, 2))
+
+        if attn_weights.size() != (bsz * self.num_heads, tgt_len, src_len):
+            raise ValueError(
+                f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is"
+                f" {attn_weights.size()}"
+            )
+
+        if attention_mask is not None:
+            if attention_mask.size() != (bsz, 1, tgt_len, src_len):
+                raise ValueError(
+                    f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is {attention_mask.size()}"
+                )
+            attn_weights = (
+                attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
+                + attention_mask
+            )
+            attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
+
+        attn_weights = F.softmax(attn_weights, dim=-1)
+
+        if layer_head_mask is not None:
+            if layer_head_mask.size() != (self.num_heads,):
+                raise ValueError(
+                    f"Head mask for a single layer should be of size {(self.num_heads,)}, but is"
+                    f" {layer_head_mask.size()}"
+                )
+            attn_weights = layer_head_mask.view(1, -1, 1, 1) * attn_weights.view(
+                bsz, self.num_heads, tgt_len, src_len
+            )
+            attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
+
+        if output_attentions:
+            # this operation is a bit awkward, but it's required to
+            # make sure that attn_weights keeps its gradient.
+            # In order to do so, attn_weights have to be reshaped
+            # twice and have to be reused in the following
+            attn_weights_reshaped = attn_weights.view(
+                bsz, self.num_heads, tgt_len, src_len
+            )
+            attn_weights = attn_weights_reshaped.view(
+                bsz * self.num_heads, tgt_len, src_len
+            )
+        else:
+            attn_weights_reshaped = None
+
+        attn_probs = F.dropout(attn_weights, p=self.dropout, training=self.training)
+
+        attn_output = torch.bmm(attn_probs, value_states)
+
+        if attn_output.size() != (bsz * self.num_heads, tgt_len, self.head_dim):
+            raise ValueError(
+                f"`attn_output` should be of size {(bsz * self.num_heads, tgt_len, self.head_dim)}, but is"
+                f" {attn_output.size()}"
+            )
+
+        attn_output = attn_output.view(bsz, self.num_heads, tgt_len, self.head_dim)
+        attn_output = attn_output.transpose(1, 2)
+
+        # Use the `embed_dim` from the config (stored in the class) rather than `hidden_state` because `attn_output` can be
+        # partitioned across GPUs when using tensor-parallelism.
+        attn_output = attn_output.reshape(bsz, tgt_len, self.embed_dim)
+
+        attn_output = self.out_proj(attn_output)
+
+        return attn_output, attn_weights_reshaped, past_key_value
+    
+
+class IndicTransFlashAttention2(IndicTransAttention):
+    """
+    IndicTrans flash attention module. This module inherits from `IndicTransAttention` as the weights of the module stays
+    untouched. The only required change would be on the forward pass where it needs to correctly call the public API of
+    flash attention and deal with padding tokens in case the input contains any of them.
+    """
+
+    # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2.__init__
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
+        # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.
+        # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left).
+        self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
+
+    def _reshape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
+        return tensor.view(bsz, seq_len, self.num_heads, self.head_dim)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        key_value_states: Optional[torch.Tensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        layer_head_mask: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        # IndicTransFlashAttention2 attention does not support output_attentions
+        if output_attentions:
+            raise ValueError("IndicTransFlashAttention2 attention does not support output_attentions")
+
+        # if key_value_states are provided this layer is used as a cross-attention layer
+        # for the decoder
+        is_cross_attention = key_value_states is not None
+
+        bsz, q_len, _ = hidden_states.size()
+
+        # get query proj
+        query_states = self._reshape(self.q_proj(hidden_states), -1, bsz)
+        # get key, value proj
+        # `past_key_value[0].shape[2] == key_value_states.shape[1]`
+        # is checking that the `sequence_length` of the `past_key_value` is the same as
+        # the provided `key_value_states` to support prefix tuning
+        if (
+            is_cross_attention
+            and past_key_value is not None
+            and past_key_value[0].shape[2] == key_value_states.shape[1]
+        ):
+            # reuse k,v, cross_attentions
+            key_states = past_key_value[0].transpose(1, 2)
+            value_states = past_key_value[1].transpose(1, 2)
+        elif is_cross_attention:
+            # cross_attentions
+            key_states = self._reshape(self.k_proj(key_value_states), -1, bsz)
+            value_states = self._reshape(self.v_proj(key_value_states), -1, bsz)
+        elif past_key_value is not None:
+            # reuse k, v, self_attention
+            key_states = self._reshape(self.k_proj(hidden_states), -1, bsz)
+            value_states = self._reshape(self.v_proj(hidden_states), -1, bsz)
+            key_states = torch.cat([past_key_value[0].transpose(1, 2), key_states], dim=1)
+            value_states = torch.cat([past_key_value[1].transpose(1, 2), value_states], dim=1)
+        else:
+            # self_attention
+            key_states = self._reshape(self.k_proj(hidden_states), -1, bsz)
+            value_states = self._reshape(self.v_proj(hidden_states), -1, bsz)
+
+        if self.is_decoder:
+            # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states.
+            # Further calls to cross_attention layer can then reuse all cross-attention
+            # key/value_states (first "if" case)
+            # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of
+            # all previous decoder key/value_states. Further calls to uni-directional self-attention
+            # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case)
+            # if encoder bi-directional self-attention `past_key_value` is always `None`
+            past_key_value = (key_states.transpose(1, 2), value_states.transpose(1, 2))
+
+        kv_seq_len = key_states.shape[-2]
+        if past_key_value is not None:
+            kv_seq_len += past_key_value[0].shape[-2]
+
+        # In PEFT, usually we cast the layer norms in float32 for training stability reasons
+        # therefore the input hidden states gets silently casted in float32. Hence, we need
+        # cast them back in the correct dtype just to be sure everything works as expected.
+        # This might slowdown training & inference so it is recommended to not cast the LayerNorms
+        # in fp32. (LlamaRMSNorm handles it correctly)
+
+        input_dtype = query_states.dtype
+        if input_dtype == torch.float32:
+            if torch.is_autocast_enabled():
+                target_dtype = torch.get_autocast_gpu_dtype()
+            # Handle the case where the model is quantized
+            elif hasattr(self.config, "_pre_quantization_dtype"):
+                target_dtype = self.config._pre_quantization_dtype
+            else:
+                target_dtype = self.q_proj.weight.dtype
+
+            logger.warning_once(
+                f"The input hidden states seems to be silently casted in float32, this might be related to"
+                f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
+                f" {target_dtype}."
+            )
+
+            query_states = query_states.to(target_dtype)
+            key_states = key_states.to(target_dtype)
+            value_states = value_states.to(target_dtype)
+
+        attn_output = self._flash_attention_forward(
+            query_states, key_states, value_states, attention_mask, q_len, dropout=self.dropout
+        )
+
+        attn_output = attn_output.reshape(bsz, q_len, -1)
+        attn_output = self.out_proj(attn_output)
+
+        if not output_attentions:
+            attn_weights = None
+
+        return attn_output, attn_weights, past_key_value
+
+    # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2._flash_attention_forward
+    def _flash_attention_forward(
+        self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None
+    ):
+        """
+        Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token
+        first unpad the input, then computes the attention scores and pad the final attention scores.
+
+        Args:
+            query_states (`torch.Tensor`):
+                Input query states to be passed to Flash Attention API
+            key_states (`torch.Tensor`):
+                Input key states to be passed to Flash Attention API
+            value_states (`torch.Tensor`):
+                Input value states to be passed to Flash Attention API
+            attention_mask (`torch.Tensor`):
+                The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
+                position of padding tokens and 1 for the position of non-padding tokens.
+            dropout (`float`):
+                Attention dropout
+            softmax_scale (`float`, *optional*):
+                The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
+        """
+        if not self._flash_attn_uses_top_left_mask:
+            causal = self.is_causal
+        else:
+            # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in LlamaFlashAttention2 __init__.
+            causal = self.is_causal and query_length != 1
+
+        # Contains at least one padding token in the sequence
+        if attention_mask is not None:
+            batch_size = query_states.shape[0]
+            query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
+                query_states, key_states, value_states, attention_mask, query_length
+            )
+
+            cu_seqlens_q, cu_seqlens_k = cu_seq_lens
+            max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
+
+            attn_output_unpad = flash_attn_varlen_func(
+                query_states,
+                key_states,
+                value_states,
+                cu_seqlens_q=cu_seqlens_q,
+                cu_seqlens_k=cu_seqlens_k,
+                max_seqlen_q=max_seqlen_in_batch_q,
+                max_seqlen_k=max_seqlen_in_batch_k,
+                dropout_p=dropout,
+                softmax_scale=softmax_scale,
+                causal=causal,
+            )
+
+            attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)
+        else:
+            attn_output = flash_attn_func(
+                query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=causal
+            )
+
+        return attn_output
+
+    # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2._upad_input
+    def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
+        indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
+        batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
+
+        key_layer = index_first_axis(
+            key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
+        )
+        value_layer = index_first_axis(
+            value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
+        )
+        if query_length == kv_seq_len:
+            query_layer = index_first_axis(
+                query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim), indices_k
+            )
+            cu_seqlens_q = cu_seqlens_k
+            max_seqlen_in_batch_q = max_seqlen_in_batch_k
+            indices_q = indices_k
+        elif query_length == 1:
+            max_seqlen_in_batch_q = 1
+            cu_seqlens_q = torch.arange(
+                batch_size + 1, dtype=torch.int32, device=query_layer.device
+            )  # There is a memcpy here, that is very bad.
+            indices_q = cu_seqlens_q[:-1]
+            query_layer = query_layer.squeeze(1)
+        else:
+            # The -q_len: slice assumes left padding.
+            attention_mask = attention_mask[:, -query_length:]
+            query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)
+
+        return (
+            query_layer,
+            key_layer,
+            value_layer,
+            indices_q,
+            (cu_seqlens_q, cu_seqlens_k),
+            (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
+        )
+
+
+class IndicTransSdpaAttention(IndicTransAttention):
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        key_value_states: Optional[torch.Tensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        layer_head_mask: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        """Input shape: Batch x Time x Channel"""
+        if output_attentions or layer_head_mask is not None:
+            # TODO: Improve this warning with e.g. `model.config._attn_implementation = "manual"` once this is implemented.
+            logger.warning_once(
+                "IndicTransModel is using IndicTransSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True` or `layer_head_mask` not None. Falling back to the manual attention"
+                ' implementation, but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
+            )
+            return super().forward(
+                hidden_states,
+                key_value_states=key_value_states,
+                past_key_value=past_key_value,
+                attention_mask=attention_mask,
+                layer_head_mask=layer_head_mask,
+                output_attentions=output_attentions,
+            )
+
+        # if key_value_states are provided this layer is used as a cross-attention layer
+        # for the decoder
+        is_cross_attention = key_value_states is not None
+
+        bsz, tgt_len, _ = hidden_states.size()
+
+        # get query proj
+        query_states = self.q_proj(hidden_states)
+        # get key, value proj
+        # `past_key_value[0].shape[2] == key_value_states.shape[1]`
+        # is checking that the `sequence_length` of the `past_key_value` is the same as
+        # the provided `key_value_states` to support prefix tuning
+        if (
+            is_cross_attention
+            and past_key_value is not None
+            and past_key_value[0].shape[2] == key_value_states.shape[1]
+        ):
+            # reuse k,v, cross_attentions
+            key_states = past_key_value[0]
+            value_states = past_key_value[1]
+        elif is_cross_attention:
+            # cross_attentions
+            key_states = self._shape(self.k_proj(key_value_states), -1, bsz)
+            value_states = self._shape(self.v_proj(key_value_states), -1, bsz)
+        elif past_key_value is not None:
+            # reuse k, v, self_attention
+            key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
+            value_states = self._shape(self.v_proj(hidden_states), -1, bsz)
+            key_states = torch.cat([past_key_value[0], key_states], dim=2)
+            value_states = torch.cat([past_key_value[1], value_states], dim=2)
+        else:
+            # self_attention
+            key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
+            value_states = self._shape(self.v_proj(hidden_states), -1, bsz)
+
+        if self.is_decoder:
+            # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states.
+            # Further calls to cross_attention layer can then reuse all cross-attention
+            # key/value_states (first "if" case)
+            # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of
+            # all previous decoder key/value_states. Further calls to uni-directional self-attention
+            # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case)
+            # if encoder bi-directional self-attention `past_key_value` is always `None`
+            past_key_value = (key_states, value_states)
+
+        query_states = self._shape(query_states, tgt_len, bsz)
+
+        # NOTE: SDPA with memory-efficient backend is currently (torch==2.1.2) bugged when using non-contiguous inputs and a custom attn_mask,
+        # but we are fine here as `_shape` do call `.contiguous()`. Reference: https://github.com/pytorch/pytorch/issues/112577
+        attn_output = F.scaled_dot_product_attention(
+            query_states,
+            key_states,
+            value_states,
+            attn_mask=attention_mask,
+            dropout_p=self.dropout if self.training else 0.0,
+            # The tgt_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case tgt_len == 1.
+            is_causal=self.is_causal and attention_mask is None and tgt_len > 1,
+        )
+
+        if attn_output.size() != (bsz, self.num_heads, tgt_len, self.head_dim):
+            raise ValueError(
+                f"`attn_output` should be of size {(bsz, self.num_heads, tgt_len, self.head_dim)}, but is"
+                f" {attn_output.size()}"
+            )
+
+        attn_output = attn_output.transpose(1, 2)
+
+        # Use the `embed_dim` from the config (stored in the class) rather than `hidden_state` because `attn_output` can be
+        # partitioned across GPUs when using tensor-parallelism.
+        attn_output = attn_output.reshape(bsz, tgt_len, self.embed_dim)
+
+        attn_output = self.out_proj(attn_output)
+
+        return attn_output, None, past_key_value
+
+
+INDICTRANS_ATTENTION_CLASSES = {
+    "eager": IndicTransAttention,
+    "sdpa": IndicTransSdpaAttention,
+    "flash_attention_2": IndicTransFlashAttention2,
+}
+
+# Copied from transformers.models.mbart.modeling_mbart.MBartEncoderLayer with MBart->IndicTrans
+class IndicTransEncoderLayer(nn.Module):
+    def __init__(self, config: IndicTransConfig):
+        super().__init__()
+        self.embed_dim = config.encoder_embed_dim
+        self.self_attn = INDICTRANS_ATTENTION_CLASSES[config._attn_implementation](
+            embed_dim=self.embed_dim,
+            num_heads=config.encoder_attention_heads,
+            dropout=config.attention_dropout,
+            config=config,
+        )
+        self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
+        self.dropout = config.dropout
+        self.activation_fn = ACT2FN[config.activation_function]
+        self.activation_dropout = config.activation_dropout
+        self.fc1 = nn.Linear(self.embed_dim, config.encoder_ffn_dim)
+        self.fc2 = nn.Linear(config.encoder_ffn_dim, self.embed_dim)
+        self.final_layer_norm = nn.LayerNorm(self.embed_dim)
+        self.normalize_before = config.encoder_normalize_before
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        layer_head_mask: torch.Tensor,
+        output_attentions: bool = False,
+    ) -> torch.Tensor:
+        """
+        Args:
+            hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
+            attention_mask (`torch.FloatTensor`): attention mask of size
+                `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
+            layer_head_mask (`torch.FloatTensor`): mask for attention heads in a given layer of size
+                `(encoder_attention_heads,)`.
+            output_attentions (`bool`, *optional*):
+                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+                returned tensors for more detail.
+        """
+        residual = hidden_states
+        if self.normalize_before:
+            hidden_states = self.self_attn_layer_norm(hidden_states)
+        hidden_states, attn_weights, _ = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            layer_head_mask=layer_head_mask,
+            output_attentions=output_attentions,
+        )
+        hidden_states = F.dropout(hidden_states, p=self.dropout, training=self.training)
+        hidden_states = residual + hidden_states
+        if not self.normalize_before:
+            hidden_states = self.self_attn_layer_norm(hidden_states)
+
+        residual = hidden_states
+        if self.normalize_before:
+            hidden_states = self.final_layer_norm(hidden_states)
+        hidden_states = self.activation_fn(self.fc1(hidden_states))
+        hidden_states = F.dropout(
+            hidden_states, p=self.activation_dropout, training=self.training
+        )
+        hidden_states = self.fc2(hidden_states)
+        hidden_states = F.dropout(hidden_states, p=self.dropout, training=self.training)
+        hidden_states = residual + hidden_states
+        if not self.normalize_before:
+            hidden_states = self.final_layer_norm(hidden_states)
+
+        if hidden_states.dtype == torch.float16 and (
+            torch.isinf(hidden_states).any() or torch.isnan(hidden_states).any()
+        ):
+            clamp_value = torch.finfo(hidden_states.dtype).max - 1000
+            hidden_states = torch.clamp(
+                hidden_states, min=-clamp_value, max=clamp_value
+            )
+
+        outputs = (hidden_states,)
+
+        if output_attentions:
+            outputs += (attn_weights,)
+
+        return outputs
+
+
+# Copied from transformers.models.mbart.modeling_mbart.MBartDecoderLayer with MBart->IndicTrans
+class IndicTransDecoderLayer(nn.Module):
+    def __init__(self, config: IndicTransConfig):
+        super().__init__()
+        self.embed_dim = config.decoder_embed_dim
+
+        self.self_attn = INDICTRANS_ATTENTION_CLASSES[config._attn_implementation](
+            embed_dim=self.embed_dim,
+            num_heads=config.decoder_attention_heads,
+            dropout=config.attention_dropout,
+            is_decoder=True,
+            is_causal=True,
+            config=config,
+        )
+        self.dropout = config.dropout
+        self.activation_fn = ACT2FN[config.activation_function]
+        self.activation_dropout = config.activation_dropout
+
+        self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
+        self.encoder_attn = INDICTRANS_ATTENTION_CLASSES[config._attn_implementation](
+            self.embed_dim,
+            config.decoder_attention_heads,
+            dropout=config.attention_dropout,
+            is_decoder=True,
+            config=config,
+        )
+        self.encoder_attn_layer_norm = nn.LayerNorm(self.embed_dim)
+        self.fc1 = nn.Linear(self.embed_dim, config.decoder_ffn_dim)
+        self.fc2 = nn.Linear(config.decoder_ffn_dim, self.embed_dim)
+        self.final_layer_norm = nn.LayerNorm(self.embed_dim)
+        self.normalize_before = config.decoder_normalize_before
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        encoder_attention_mask: Optional[torch.Tensor] = None,
+        layer_head_mask: Optional[torch.Tensor] = None,
+        cross_attn_layer_head_mask: Optional[torch.Tensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        output_attentions: Optional[bool] = False,
+        use_cache: Optional[bool] = True,
+    ) -> torch.Tensor:
+        """
+        Args:
+            hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
+            attention_mask (`torch.FloatTensor`): attention mask of size
+                `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
+            encoder_hidden_states (`torch.FloatTensor`):
+                cross attention input to the layer of shape `(batch, seq_len, embed_dim)`
+            encoder_attention_mask (`torch.FloatTensor`): encoder attention mask of size
+                `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
+            layer_head_mask (`torch.FloatTensor`): mask for attention heads in a given layer of size
+                `(encoder_attention_heads,)`.
+            cross_attn_layer_head_mask (`torch.FloatTensor`): mask for cross-attention heads in a given layer of
+                size `(decoder_attention_heads,)`.
+            past_key_value (`Tuple(torch.FloatTensor)`): cached past key and value projection states
+            output_attentions (`bool`, *optional*):
+                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+                returned tensors for more detail.
+        """
+        residual = hidden_states
+        if self.normalize_before:
+            hidden_states = self.self_attn_layer_norm(hidden_states)
+
+        # Self Attention
+        # decoder uni-directional self-attention cached key/values tuple is at positions 1,2
+        self_attn_past_key_value = (
+            past_key_value[:2] if past_key_value is not None else None
+        )
+        # add present self-attn cache to positions 1,2 of present_key_value tuple
+        hidden_states, self_attn_weights, present_key_value = self.self_attn(
+            hidden_states=hidden_states,
+            past_key_value=self_attn_past_key_value,
+            attention_mask=attention_mask,
+            layer_head_mask=layer_head_mask,
+            output_attentions=output_attentions,
+        )
+        hidden_states = F.dropout(hidden_states, p=self.dropout, training=self.training)
+        hidden_states = residual + hidden_states
+        if not self.normalize_before:
+            hidden_states = self.self_attn_layer_norm(hidden_states)
+
+        # Cross-Attention Block
+        cross_attn_present_key_value = None
+        cross_attn_weights = None
+        if encoder_hidden_states is not None:
+            residual = hidden_states
+            if self.normalize_before:
+                hidden_states = self.encoder_attn_layer_norm(hidden_states)
+
+            # cross_attn cached key/values tuple is at positions 3,4 of present_key_value tuple
+            cross_attn_past_key_value = (
+                past_key_value[-2:] if past_key_value is not None else None
+            )
+            (
+                hidden_states,
+                cross_attn_weights,
+                cross_attn_present_key_value,
+            ) = self.encoder_attn(
+                hidden_states=hidden_states,
+                key_value_states=encoder_hidden_states,
+                attention_mask=encoder_attention_mask,
+                layer_head_mask=cross_attn_layer_head_mask,
+                past_key_value=cross_attn_past_key_value,
+                output_attentions=output_attentions,
+            )
+            hidden_states = F.dropout(
+                hidden_states, p=self.dropout, training=self.training
+            )
+            hidden_states = residual + hidden_states
+            if not self.normalize_before:
+                hidden_states = self.encoder_attn_layer_norm(hidden_states)
+
+            # add cross-attn to positions 3,4 of present_key_value tuple
+            present_key_value = present_key_value + cross_attn_present_key_value
+
+        # Fully Connected
+        residual = hidden_states
+        if self.normalize_before:
+            hidden_states = self.final_layer_norm(hidden_states)
+        hidden_states = self.activation_fn(self.fc1(hidden_states))
+        hidden_states = F.dropout(
+            hidden_states, p=self.activation_dropout, training=self.training
+        )
+        hidden_states = self.fc2(hidden_states)
+        hidden_states = F.dropout(hidden_states, p=self.dropout, training=self.training)
+        hidden_states = residual + hidden_states
+        if not self.normalize_before:
+            hidden_states = self.final_layer_norm(hidden_states)
+
+        outputs = (hidden_states,)
+
+        if output_attentions:
+            outputs += (self_attn_weights, cross_attn_weights)
+
+        if use_cache:
+            outputs += (present_key_value,)
+
+        return outputs
+
+
+# Copied from transformers.models.m2m_100.modeling_m2m_100.M2M100PretrainedModel->IndicTrans
+class IndicTransPreTrainedModel(PreTrainedModel):
+    config_class = IndicTransConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["IndicTransAttention"]
+
+    def _init_weights(self, module):
+        std = self.config.init_std
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if isinstance(module, (IndicTransDecoder, IndicTransEncoder)):
+            module.gradient_checkpointing = value
+
+
+# Copied from transformers.models.m2m_100.modeling_m2m_100.M2M100EncoderLayer->IndicTrans
+class IndicTransEncoder(IndicTransPreTrainedModel):
+    """
+    Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a
+    [`IndicTransEncoderLayer`].
+
+    Args:
+        config: IndicTransConfig
+        embed_tokens (nn.Embedding): output embedding
+    """
+
+    def __init__(
+        self, config: IndicTransConfig, embed_tokens: Optional[nn.Embedding] = None
+    ):
+        super().__init__(config)
+
+        self.dropout = config.dropout
+        self.layerdrop = config.encoder_layerdrop
+
+        embed_dim = config.encoder_embed_dim
+        self.padding_idx = config.pad_token_id
+        self.max_source_positions = config.max_source_positions
+        self.embed_scale = math.sqrt(embed_dim) if config.scale_embedding else 1.0
+
+        self.embed_tokens = nn.Embedding(
+            config.encoder_vocab_size, embed_dim, self.padding_idx
+        )
+
+        if embed_tokens is not None:
+            self.embed_tokens.weight = embed_tokens.weight
+
+        self.embed_positions = IndicTransSinusoidalPositionalEmbedding(
+            config.max_source_positions,
+            embed_dim,
+            self.padding_idx,
+        )
+        self.layers = nn.ModuleList(
+            [IndicTransEncoderLayer(config) for _ in range(config.encoder_layers)]
+        )
+        self.layer_norm = (
+            nn.LayerNorm(embed_dim) if config.encoder_normalize_before else None
+        )
+        self.layernorm_embedding = (
+            nn.LayerNorm(embed_dim) if config.layernorm_embedding else None
+        )
+
+        self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
+        self._use_sdpa = config._attn_implementation == "sdpa"
+
+        self.gradient_checkpointing = False
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ):
+        r"""
+        Args:
+            input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+                Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you
+                provide it.
+
+                Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+                [`PreTrainedTokenizer.__call__`] for details.
+
+                [What are input IDs?](../glossary#input-ids)
+            attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
+
+                - 1 for tokens that are **not masked**,
+                - 0 for tokens that are **masked**.
+
+                [What are attention masks?](../glossary#attention-mask)
+            head_mask (`torch.Tensor` of shape `(encoder_layers, encoder_attention_heads)`, *optional*):
+                Mask to nullify selected heads of the attention modules. Mask values selected in `[0, 1]`:
+
+                - 1 indicates the head is **not masked**,
+                - 0 indicates the head is **masked**.
+
+            inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
+                Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation.
+                This is useful if you want more control over how to convert `input_ids` indices into associated vectors
+                than the model's internal embedding lookup matrix.
+            output_attentions (`bool`, *optional*):
+                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+                returned tensors for more detail.
+            output_hidden_states (`bool`, *optional*):
+                Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
+                for more detail.
+            return_dict (`bool`, *optional*):
+                Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+        """
+        output_attentions = (
+            output_attentions
+            if output_attentions is not None
+            else self.config.output_attentions
+        )
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+
+        # retrieve input_ids and inputs_embeds
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError(
+                "You cannot specify both input_ids and inputs_embeds at the same time"
+            )
+        elif input_ids is not None:
+            self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
+            input_shape = input_ids.size()
+            input_ids = input_ids.view(-1, input_shape[-1])
+        elif inputs_embeds is not None:
+            input_shape = inputs_embeds.size()[:-1]
+        else:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids) * self.embed_scale
+
+        embed_pos = self.embed_positions(input_ids, inputs_embeds)
+        embed_pos = embed_pos.to(inputs_embeds.device)
+
+        hidden_states = inputs_embeds + embed_pos
+        if self.layernorm_embedding is not None:
+            hidden_states = self.layernorm_embedding(hidden_states)
+        hidden_states = F.dropout(hidden_states, p=self.dropout, training=self.training)
+
+        if attention_mask is not None:
+            if self._use_flash_attention_2:
+                attention_mask = attention_mask if 0 in attention_mask else None
+            elif self._use_sdpa and head_mask is None and not output_attentions:
+                # output_attentions=True & head_mask can not be supported when using SDPA, fall back to
+                # the manual implementation that requires a 4D causal mask in all cases.
+                # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+                attention_mask = _prepare_4d_attention_mask_for_sdpa(attention_mask, inputs_embeds.dtype)
+            else:
+                # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+                attention_mask = _prepare_4d_attention_mask(attention_mask, inputs_embeds.dtype)
+
+
+        encoder_states = () if output_hidden_states else None
+        all_attentions = () if output_attentions else None
+
+        # check if head_mask has a correct number of layers specified if desired
+        if head_mask is not None:
+            if head_mask.size()[0] != len(self.layers):
+                raise ValueError(
+                    f"The head_mask should be specified for {len(self.layers)} layers, but it is for"
+                    f" {head_mask.size()[0]}."
+                )
+        deepspeed_zero3_is_enabled = is_deepspeed_zero3_enabled()
+
+        for idx, encoder_layer in enumerate(self.layers):
+            if output_hidden_states:
+                encoder_states = encoder_states + (hidden_states,)
+
+            # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description)
+            dropout_probability = torch.rand([])
+
+            skip_the_layer = (
+                True
+                if self.training and (dropout_probability < self.layerdrop)
+                else False
+            )
+            if not skip_the_layer or deepspeed_zero3_is_enabled:
+                # under deepspeed zero3 all gpus must run in sync
+
+                if self.gradient_checkpointing and self.training:
+                    # create gradient checkpointing function
+                    def create_custom_forward(module):
+                        def custom_forward(*inputs):
+                            return module(*inputs, output_attentions)
+
+                        return custom_forward
+
+                    layer_outputs = torch.utils.checkpoint.checkpoint(
+                        create_custom_forward(encoder_layer),
+                        hidden_states,
+                        attention_mask,
+                        (head_mask[idx] if head_mask is not None else None),
+                    )
+                else:
+                    layer_outputs = encoder_layer(
+                        hidden_states,
+                        attention_mask,
+                        layer_head_mask=(
+                            head_mask[idx] if head_mask is not None else None
+                        ),
+                        output_attentions=output_attentions,
+                    )
+
+                hidden_states = layer_outputs[0]
+
+            if skip_the_layer:
+                layer_outputs = (None, None)
+
+            if output_attentions:
+                all_attentions = all_attentions + (layer_outputs[1],)
+
+        if self.layer_norm is not None:
+            hidden_states = self.layer_norm(hidden_states)
+
+        if output_hidden_states:
+            encoder_states = encoder_states + (hidden_states,)
+
+        if not return_dict:
+            return tuple(
+                v
+                for v in [hidden_states, encoder_states, all_attentions]
+                if v is not None
+            )
+        return BaseModelOutput(
+            last_hidden_state=hidden_states,
+            hidden_states=encoder_states,
+            attentions=all_attentions,
+        )
+
+
+# Copied from transformers.models.m2m_100.modeling_m2m_100.M2M100DecoderLayer->IndicTrans
+class IndicTransDecoder(IndicTransPreTrainedModel):
+    """
+    Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a [`IndicTransDecoderLayer`]
+
+    Args:
+        config: IndicTransConfig
+        embed_tokens (nn.Embedding): output embedding
+    """
+
+    def __init__(
+        self, config: IndicTransConfig, embed_tokens: Optional[nn.Embedding] = None
+    ):
+        super().__init__(config)
+        self.dropout = config.dropout
+        self.layerdrop = config.decoder_layerdrop
+
+        embed_dim = config.encoder_embed_dim
+        self.padding_idx = config.pad_token_id
+        self.max_target_positions = config.max_target_positions
+        self.embed_scale = math.sqrt(embed_dim) if config.scale_embedding else 1.0
+
+        self.embed_tokens = nn.Embedding(
+            config.decoder_vocab_size, embed_dim, self.padding_idx
+        )
+
+        if embed_tokens is not None:
+            self.embed_tokens.weight = embed_tokens.weight
+
+        self.embed_positions = IndicTransSinusoidalPositionalEmbedding(
+            config.max_target_positions,
+            embed_dim,
+            self.padding_idx,
+        )
+        self.layers = nn.ModuleList(
+            [IndicTransDecoderLayer(config) for _ in range(config.decoder_layers)]
+        )
+        self.layer_norm = (
+            nn.LayerNorm(embed_dim) if config.decoder_normalize_before else None
+        )
+        self.layernorm_embedding = (
+            nn.LayerNorm(embed_dim) if config.layernorm_embedding else None
+        )
+
+        self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
+        self._use_sdpa = config._attn_implementation == "sdpa"
+
+        self.gradient_checkpointing = False
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        encoder_attention_mask: Optional[torch.Tensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        cross_attn_head_mask: Optional[torch.Tensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ):
+        r"""
+        Args:
+            input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+                Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you
+                provide it.
+
+                Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+                [`PreTrainedTokenizer.__call__`] for details.
+
+                [What are input IDs?](../glossary#input-ids)
+            attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
+
+                - 1 for tokens that are **not masked**,
+                - 0 for tokens that are **masked**.
+
+                [What are attention masks?](../glossary#attention-mask)
+            encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, encoder_sequence_length, hidden_size)`, *optional*):
+                Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention
+                of the decoder.
+            encoder_attention_mask (`torch.LongTensor` of shape `(batch_size, encoder_sequence_length)`, *optional*):
+                Mask to avoid performing cross-attention on padding tokens indices of encoder input_ids. Mask values
+                selected in `[0, 1]`:
+
+                - 1 for tokens that are **not masked**,
+                - 0 for tokens that are **masked**.
+
+                [What are attention masks?](../glossary#attention-mask)
+            head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*):
+                Mask to nullify selected heads of the attention modules. Mask values selected in `[0, 1]`:
+
+                - 1 indicates the head is **not masked**,
+                - 0 indicates the head is **masked**.
+
+            cross_attn_head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*):
+                Mask to nullify selected heads of the cross-attention modules in the decoder to avoid performing
+                cross-attention on hidden heads. Mask values selected in `[0, 1]`:
+
+                - 1 indicates the head is **not masked**,
+                - 0 indicates the head is **masked**.
+
+            past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+                Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of
+                shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of
+                shape `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`.
+
+                Contains pre-computed hidden-states (key and values in the self-attention blocks and in the
+                cross-attention blocks) that can be used (see `past_key_values` input) to speed up sequential decoding.
+
+                If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those
+                that don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of
+                all `decoder_input_ids` of shape `(batch_size, sequence_length)`. inputs_embeds (`torch.FloatTensor` of
+                shape `(batch_size, sequence_length, hidden_size)`, *optional*): Optionally, instead of passing
+                `input_ids` you can choose to directly pass an embedded representation. This is useful if you want more
+                control over how to convert `input_ids` indices into associated vectors than the model's internal
+                embedding lookup matrix.
+            output_attentions (`bool`, *optional*):
+                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+                returned tensors for more detail.
+            output_hidden_states (`bool`, *optional*):
+                Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
+                for more detail.
+            return_dict (`bool`, *optional*):
+                Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+        """
+        output_attentions = (
+            output_attentions
+            if output_attentions is not None
+            else self.config.output_attentions
+        )
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+
+        # retrieve input_ids and inputs_embeds
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError(
+                "You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time"
+            )
+        elif input_ids is not None:
+            input_shape = input_ids.size()
+            input_ids = input_ids.view(-1, input_shape[-1])
+        elif inputs_embeds is not None:
+            input_shape = inputs_embeds.size()[:-1]
+        else:
+            raise ValueError(
+                "You have to specify either decoder_input_ids or decoder_inputs_embeds"
+            )
+
+        # past_key_values_length
+        past_key_values_length = (
+            past_key_values[0][0].shape[2] if past_key_values is not None else 0
+        )
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids) * self.embed_scale
+
+
+        if self._use_flash_attention_2:
+            # 2d mask is passed through the layers
+            attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None
+        elif self._use_sdpa and not output_attentions and cross_attn_head_mask is None:
+            # output_attentions=True & cross_attn_head_mask can not be supported when using SDPA, and we fall back on
+            # the manual implementation that requires a 4D causal mask in all cases.
+            attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
+                attention_mask,
+                input_shape,
+                inputs_embeds,
+                past_key_values_length,
+            )
+        else:
+            # 4d mask is passed through the layers
+            attention_mask = _prepare_4d_causal_attention_mask(
+                attention_mask, input_shape, inputs_embeds, past_key_values_length
+            )
+
+        # expand encoder attention mask
+        if encoder_hidden_states is not None and encoder_attention_mask is not None:
+            if self._use_flash_attention_2:
+                encoder_attention_mask = encoder_attention_mask if 0 in encoder_attention_mask else None
+            elif self._use_sdpa and cross_attn_head_mask is None and not output_attentions:
+                # output_attentions=True & cross_attn_head_mask can not be supported when using SDPA, and we fall back on
+                # the manual implementation that requires a 4D causal mask in all cases.
+                # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+                encoder_attention_mask = _prepare_4d_attention_mask_for_sdpa(
+                    encoder_attention_mask,
+                    inputs_embeds.dtype,
+                    tgt_len=input_shape[-1],
+                )
+            else:
+                # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+                encoder_attention_mask = _prepare_4d_attention_mask(
+                    encoder_attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]
+                )
+
+        # embed positions
+        positions = self.embed_positions(
+            input_ids, inputs_embeds, past_key_values_length
+        )
+        positions = positions.to(inputs_embeds.device)
+
+        hidden_states = inputs_embeds + positions
+        if self.layernorm_embedding is not None:
+            hidden_states = self.layernorm_embedding(hidden_states)
+
+        hidden_states = F.dropout(hidden_states, p=self.dropout, training=self.training)
+
+        if self.gradient_checkpointing and self.training:
+            if use_cache:
+                logger.warning_once(
+                    "`use_cache=True` is incompatible with gradient checkpointing. Setting"
+                    " `use_cache=False`..."
+                )
+                use_cache = False
+
+        # decoder layers
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+        all_cross_attentions = () if output_attentions else None
+        next_decoder_cache = () if use_cache else None
+
+        # check if head_mask/cross_attn_head_mask has a correct number of layers specified if desired
+        for attn_mask, mask_name in zip(
+            [head_mask, cross_attn_head_mask], ["head_mask", "cross_attn_head_mask"]
+        ):
+            if attn_mask is not None:
+                if attn_mask.size()[0] != len(self.layers):
+                    raise ValueError(
+                        f"The `{mask_name}` should be specified for {len(self.layers)} layers, but it is for"
+                        f" {head_mask.size()[0]}."
+                    )
+        deepspeed_zero3_is_enabled = is_deepspeed_zero3_enabled()
+
+        for idx, decoder_layer in enumerate(self.layers):
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description)
+            dropout_probability = torch.rand([])
+
+            skip_the_layer = (
+                True
+                if self.training and (dropout_probability < self.layerdrop)
+                else False
+            )
+            if not skip_the_layer or deepspeed_zero3_is_enabled:
+                # under deepspeed zero3 all gpus must run in sync
+
+                past_key_value = (
+                    past_key_values[idx] if past_key_values is not None else None
+                )
+
+                if self.gradient_checkpointing and self.training:
+
+                    def create_custom_forward(module):
+                        def custom_forward(*inputs):
+                            # None for past_key_value
+                            return module(*inputs, output_attentions, use_cache)
+
+                        return custom_forward
+
+                    layer_outputs = torch.utils.checkpoint.checkpoint(
+                        create_custom_forward(decoder_layer),
+                        hidden_states,
+                        attention_mask,
+                        encoder_hidden_states,
+                        encoder_attention_mask,
+                        head_mask[idx] if head_mask is not None else None,
+                        cross_attn_head_mask[idx]
+                        if cross_attn_head_mask is not None
+                        else None,
+                        None,
+                    )
+                else:
+                    layer_outputs = decoder_layer(
+                        hidden_states,
+                        attention_mask=attention_mask,
+                        encoder_hidden_states=encoder_hidden_states,
+                        encoder_attention_mask=encoder_attention_mask,
+                        layer_head_mask=(
+                            head_mask[idx] if head_mask is not None else None
+                        ),
+                        cross_attn_layer_head_mask=(
+                            cross_attn_head_mask[idx]
+                            if cross_attn_head_mask is not None
+                            else None
+                        ),
+                        past_key_value=past_key_value,
+                        output_attentions=output_attentions,
+                        use_cache=use_cache,
+                    )
+
+                hidden_states = layer_outputs[0]
+
+            if skip_the_layer:
+                continue
+
+            if use_cache:
+                next_decoder_cache += (layer_outputs[3 if output_attentions else 1],)
+
+            if output_attentions:
+                all_self_attns += (layer_outputs[1],)
+                all_cross_attentions += (layer_outputs[2],)
+
+        if self.layer_norm is not None:
+            hidden_states = self.layer_norm(hidden_states)
+
+        # add hidden states from the last decoder layer
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        next_cache = next_decoder_cache if use_cache else None
+        if not return_dict:
+            return tuple(
+                v
+                for v in [
+                    hidden_states,
+                    next_cache,
+                    all_hidden_states,
+                    all_self_attns,
+                    all_cross_attentions,
+                ]
+                if v is not None
+            )
+        return BaseModelOutputWithPastAndCrossAttentions(
+            last_hidden_state=hidden_states,
+            past_key_values=next_cache,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+            cross_attentions=all_cross_attentions,
+        )
+
+
+# Copied from transformers.models.m2m_100.modeling_m2m_100.M2M100Model->IndicTrans
+class IndicTransModel(IndicTransPreTrainedModel):
+    _tied_weights_keys = None
+
+    def __init__(self, config: IndicTransConfig):
+        super().__init__(config)
+
+        self.encoder = IndicTransEncoder(config)
+        self.decoder = IndicTransDecoder(config)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_encoder(self):
+        return self.encoder
+
+    def get_decoder(self):
+        return self.decoder
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        decoder_input_ids: Optional[torch.LongTensor] = None,
+        decoder_attention_mask: Optional[torch.LongTensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        decoder_head_mask: Optional[torch.Tensor] = None,
+        cross_attn_head_mask: Optional[torch.Tensor] = None,
+        encoder_outputs: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
+        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        decoder_inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple[torch.Tensor], Seq2SeqModelOutput]:
+        output_attentions = (
+            output_attentions
+            if output_attentions is not None
+            else self.config.output_attentions
+        )
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+
+        if encoder_outputs is None:
+            encoder_outputs = self.encoder(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                head_mask=head_mask,
+                inputs_embeds=inputs_embeds,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+            )
+        # If the user passed a tuple for encoder_outputs, we wrap it in a BaseModelOutput when return_dict=True
+        elif return_dict and not isinstance(encoder_outputs, BaseModelOutput):
+            encoder_outputs = BaseModelOutput(
+                last_hidden_state=encoder_outputs[0],
+                hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None,
+                attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None,
+            )
+
+        # decoder outputs consists of (dec_features, past_key_value, dec_hidden, dec_attn)
+        decoder_outputs = self.decoder(
+            input_ids=decoder_input_ids,
+            attention_mask=decoder_attention_mask,
+            encoder_hidden_states=encoder_outputs[0],
+            encoder_attention_mask=attention_mask,
+            head_mask=decoder_head_mask,
+            cross_attn_head_mask=cross_attn_head_mask,
+            past_key_values=past_key_values,
+            inputs_embeds=decoder_inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        if not return_dict:
+            return decoder_outputs + encoder_outputs
+
+        return Seq2SeqModelOutput(
+            last_hidden_state=decoder_outputs.last_hidden_state,
+            past_key_values=decoder_outputs.past_key_values,
+            decoder_hidden_states=decoder_outputs.hidden_states,
+            decoder_attentions=decoder_outputs.attentions,
+            cross_attentions=decoder_outputs.cross_attentions,
+            encoder_last_hidden_state=encoder_outputs.last_hidden_state,
+            encoder_hidden_states=encoder_outputs.hidden_states,
+            encoder_attentions=encoder_outputs.attentions,
+        )
+
+
+# Copied from transformers.models.m2m_100.modeling_m2m_100.M2M100ForConditionalGeneration->IndicTrans
+class IndicTransForConditionalGeneration(IndicTransPreTrainedModel):
+    base_model_prefix = "model"
+    _tied_weights_keys = None
+    _label_smoothing = 0.0
+
+    def __init__(self, config: IndicTransConfig):
+        super().__init__(config)
+        self.model = IndicTransModel(config)
+        self.lm_head = nn.Linear(
+            config.decoder_embed_dim, config.decoder_vocab_size, bias=False
+        )
+
+        if config.share_decoder_input_output_embed:
+            self.lm_head.weight = self.model.decoder.embed_tokens.weight
+
+        self.post_init()
+
+    def tie_weights(self):
+        pass
+
+    def get_encoder(self):
+        return self.model.get_encoder()
+
+    def get_decoder(self):
+        return self.model.get_decoder()
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+    
+    def set_label_smoothing(self, label_smoothing):
+        self._label_smoothing = label_smoothing
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        decoder_input_ids: Optional[torch.LongTensor] = None,
+        decoder_attention_mask: Optional[torch.LongTensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        decoder_head_mask: Optional[torch.Tensor] = None,
+        cross_attn_head_mask: Optional[torch.Tensor] = None,
+        encoder_outputs: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
+        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        decoder_inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple[torch.Tensor], Seq2SeqLMOutput]:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+        Returns:
+        """
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+
+        if labels is not None:
+            if decoder_input_ids is None:
+                decoder_input_ids = shift_tokens_right(
+                    labels, self.config.pad_token_id, self.config.decoder_start_token_id
+                )
+
+        outputs = self.model(
+            input_ids,
+            attention_mask=attention_mask,
+            decoder_input_ids=decoder_input_ids,
+            encoder_outputs=encoder_outputs,
+            decoder_attention_mask=decoder_attention_mask,
+            head_mask=head_mask,
+            decoder_head_mask=decoder_head_mask,
+            cross_attn_head_mask=cross_attn_head_mask,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            decoder_inputs_embeds=decoder_inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        lm_logits = self.lm_head(outputs[0])
+
+        masked_lm_loss = None
+        if labels is not None:
+            # move labels to the correct device to enable PP
+            labels = labels.to(lm_logits.device)
+            masked_lm_loss = F.cross_entropy(
+                input=lm_logits.view(-1, self.config.decoder_vocab_size),
+                target=labels.view(-1),
+                ignore_index=-100,
+                label_smoothing=self._label_smoothing,
+            )
+
+        if not return_dict:
+            output = (lm_logits,) + outputs[1:]
+            return (
+                ((masked_lm_loss,) + output) if masked_lm_loss is not None else output
+            )
+
+        return Seq2SeqLMOutput(
+            loss=masked_lm_loss,
+            logits=lm_logits,
+            past_key_values=outputs.past_key_values,
+            decoder_hidden_states=outputs.decoder_hidden_states,
+            decoder_attentions=outputs.decoder_attentions,
+            cross_attentions=outputs.cross_attentions,
+            encoder_last_hidden_state=outputs.encoder_last_hidden_state,
+            encoder_hidden_states=outputs.encoder_hidden_states,
+            encoder_attentions=outputs.encoder_attentions,
+        )
+
+    def prepare_inputs_for_generation(
+        self,
+        decoder_input_ids,
+        past_key_values=None,
+        attention_mask=None,
+        head_mask=None,
+        decoder_head_mask=None,
+        cross_attn_head_mask=None,
+        use_cache=None,
+        encoder_outputs=None,
+        **kwargs,
+    ):
+        # cut decoder_input_ids if past is used
+        if past_key_values is not None:
+            decoder_input_ids = decoder_input_ids[:, -1:]
+
+        return {
+            "input_ids": None,  # encoder_outputs is defined. input_ids not needed
+            "encoder_outputs": encoder_outputs,
+            "past_key_values": past_key_values,
+            "decoder_input_ids": decoder_input_ids,
+            "attention_mask": attention_mask,
+            "head_mask": head_mask,
+            "decoder_head_mask": decoder_head_mask,
+            "cross_attn_head_mask": cross_attn_head_mask,
+            "use_cache": use_cache,  # change this to avoid caching (presumably for debugging)
+        }
+
+    @staticmethod
+    def _reorder_cache(past_key_values, beam_idx):
+        reordered_past = ()
+        for layer_past in past_key_values:
+            reordered_past += (
+                tuple(
+                    past_state.index_select(0, beam_idx) for past_state in layer_past
+                ),
+            )
+        return reordered_past

IndicTrans2/huggingface_interface/train_lora.py

@@ -0,0 +1,355 @@
+import os
+import argparse
+import pandas as pd
+from datasets import Dataset
+from sacrebleu.metrics import BLEU, CHRF
+from peft import LoraConfig, get_peft_model
+from IndicTransToolkit import IndicProcessor, IndicDataCollator
+
+from transformers import (
+    Seq2SeqTrainer,
+    Seq2SeqTrainingArguments,
+    AutoModelForSeq2SeqLM,
+    AutoTokenizer,
+    EarlyStoppingCallback,
+)
+
+bleu_metric = BLEU()
+chrf_metric = CHRF()
+
+
+def get_arg_parse():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model",
+        type=str,
+    )
+    parser.add_argument(
+        "--src_lang_list",
+        type=str,
+        help="comma separated list of source languages",
+    )
+    parser.add_argument(
+        "--tgt_lang_list",
+        type=str,
+        help="comma separated list of target languages",
+    )
+    parser.add_argument("--data_dir", type=str)
+    parser.add_argument("--output_dir", type=str)
+    parser.add_argument("--save_steps", type=int, default=1000)
+    parser.add_argument("--eval_steps", type=int, default=1000)
+    parser.add_argument("--batch_size", type=int, default=32)
+    parser.add_argument("--num_train_epochs", type=int, default=100)
+    parser.add_argument("--max_steps", type=int, default=1000000)
+    parser.add_argument("--grad_accum_steps", type=int, default=4)
+    parser.add_argument("--warmup_steps", type=int, default=4000)
+    parser.add_argument("--warmup_ratio", type=int, default=0.0)
+    parser.add_argument("--max_grad_norm", type=float, default=1.0)
+    parser.add_argument("--learning_rate", type=float, default=5e-4)
+    parser.add_argument("--weight_decay", type=float, default=0.0)
+    parser.add_argument("--adam_beta1", type=float, default=0.9)
+    parser.add_argument("--adam_beta2", type=float, default=0.98)
+    parser.add_argument("--dropout", type=float, default=0.0)
+    parser.add_argument("--print_samples", action="store_true")
+    parser.add_argument(
+        "--optimizer",
+        type=str,
+        default="adamw_torch",
+        choices=[
+            "adam_hf",
+            "adamw_torch",
+            "adamw_torch_fused",
+            "adamw_apex_fused",
+            "adafactor",
+        ],
+    )
+    parser.add_argument(
+        "--lr_scheduler",
+        type=str,
+        default="inverse_sqrt",
+        choices=[
+            "inverse_sqrt",
+            "linear",
+            "polynomial",
+            "cosine",
+            "constant",
+            "constant_with_warmup",
+        ],
+    )
+    parser.add_argument("--label_smoothing", type=float, default=0.0)
+    parser.add_argument("--num_workers", type=int, default=8)
+    parser.add_argument("--metric_for_best_model", type=str, default="eval_loss")
+    parser.add_argument("--greater_is_better", action="store_true")
+    parser.add_argument("--lora_target_modules", type=str, default="q_proj,k_proj")
+    parser.add_argument("--lora_dropout", type=float, default=0.1)
+    parser.add_argument("--lora_r", type=int, default=16)
+    parser.add_argument("--lora_alpha", type=int, default=32)
+    parser.add_argument(
+        "--report_to",
+        type=str,
+        default="none",
+        choices=["wandb", "tensorboard", "azure_ml", "none"],
+    )
+    parser.add_argument("--patience", type=int, default=5),
+    parser.add_argument("--threshold", type=float, default=1e-3)
+    return parser
+
+
+def load_and_process_translation_dataset(
+    data_dir,
+    split="train",
+    tokenizer=None,
+    processor=None,
+    src_lang_list=None,
+    tgt_lang_list=None,
+    num_proc=8,
+    seed=42
+):
+    complete_dataset = {
+        "sentence_SRC": [],
+        "sentence_TGT": [],
+    }
+
+    for src_lang in src_lang_list:
+        for tgt_lang in tgt_lang_list:
+            if src_lang == tgt_lang:
+                continue
+            src_path = os.path.join(
+                data_dir, split, f"{src_lang}-{tgt_lang}", f"{split}.{src_lang}"
+            )
+            tgt_path = os.path.join(
+                data_dir, split, f"{src_lang}-{tgt_lang}", f"{split}.{tgt_lang}"
+            )
+            if not os.path.exists(src_path) or not os.path.exists(tgt_path):
+                raise FileNotFoundError(
+                    f"Source ({split}.{src_lang}) or Target ({split}.{tgt_lang}) file not found in {data_dir}"
+                )
+            with open(src_path, encoding="utf-8") as src_file, open(
+                tgt_path, encoding="utf-8"
+            ) as tgt_file:
+                src_lines = src_file.readlines()
+                tgt_lines = tgt_file.readlines()
+
+            # Ensure both files have the same number of lines
+            assert len(src_lines) == len(
+                tgt_lines
+            ), f"Source and Target files have different number of lines for {split}.{src_lang} and {split}.{tgt_lang}"
+
+            complete_dataset["sentence_SRC"] += processor.preprocess_batch(
+                src_lines, src_lang=src_lang, tgt_lang=tgt_lang, is_target=False
+            )
+
+            complete_dataset["sentence_TGT"] += processor.preprocess_batch(
+                tgt_lines, src_lang=tgt_lang, tgt_lang=src_lang, is_target=True
+            )
+
+    complete_dataset = Dataset.from_dict(complete_dataset).shuffle(seed=seed)
+
+    return complete_dataset.map(
+        lambda example: preprocess_fn(
+            example,
+            tokenizer=tokenizer
+        ),
+        batched=True,
+        num_proc=num_proc,
+    )
+
+
+def compute_metrics_factory(
+    tokenizer, metric_dict=None, print_samples=False, n_samples=10
+):
+    def compute_metrics(eval_preds):
+        preds, labels = eval_preds
+
+        labels[labels == -100] = tokenizer.pad_token_id
+        preds[preds == -100] = tokenizer.pad_token_id
+
+        with tokenizer.as_target_tokenizer():
+            preds = [
+                x.strip()
+                for x in tokenizer.batch_decode(
+                    preds, skip_special_tokens=True, clean_up_tokenization_spaces=True
+                )
+            ]
+            labels = [
+                x.strip()
+                for x in tokenizer.batch_decode(
+                    labels, skip_special_tokens=True, clean_up_tokenization_spaces=True
+                )
+            ]
+
+        assert len(preds) == len(
+            labels
+        ), "Predictions and Labels have different lengths"
+
+        df = pd.DataFrame({"Predictions": preds, "References": labels}).sample(
+            n=n_samples
+        )
+
+        if print_samples:
+            for pred, label in zip(df["Predictions"].values, df["References"].values):
+                print(f" | > Prediction: {pred}")
+                print(f" | > Reference: {label}\n")
+
+        return {
+            metric_name: metric.corpus_score(preds, [labels]).score
+            for (metric_name, metric) in metric_dict.items()
+        }
+
+    return compute_metrics
+
+
+def preprocess_fn(example, tokenizer, **kwargs):
+    model_inputs = tokenizer(
+        example["sentence_SRC"], truncation=True, padding=False, max_length=256
+    )
+
+    with tokenizer.as_target_tokenizer():
+        labels = tokenizer(
+            example["sentence_TGT"], truncation=True, padding=False, max_length=256
+        )
+
+    model_inputs["labels"] = labels["input_ids"]
+    return model_inputs
+
+
+def main(args):
+    print(f" | > Loading {args.model} and tokenizer ...")
+    model = AutoModelForSeq2SeqLM.from_pretrained(
+        args.model,
+        trust_remote_code=True,
+        attn_implementation="eager",
+        dropout=args.dropout
+    )
+
+    tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
+    processor = IndicProcessor(inference=False) # pre-process before tokenization
+    
+    data_collator = IndicDataCollator(
+        tokenizer=tokenizer,
+        model=model,
+        padding="longest", # saves padding tokens
+        pad_to_multiple_of=8, # better to have it as 8 when using fp16
+        label_pad_token_id=-100
+    )
+
+    if args.data_dir is not None:
+        train_dataset = load_and_process_translation_dataset(
+            args.data_dir,
+            split="train",
+            tokenizer=tokenizer,
+            processor=processor,
+            src_lang_list=args.src_lang_list.split(","),
+            tgt_lang_list=args.tgt_lang_list.split(","),
+        )
+        print(f" | > Loaded train dataset from {args.data_dir}. Size: {len(train_dataset)} ...")
+
+        eval_dataset = load_and_process_translation_dataset(
+            args.data_dir,
+            split="dev",
+            tokenizer=tokenizer,
+            processor=processor,
+            src_lang_list=args.src_lang_list.split(","),
+            tgt_lang_list=args.tgt_lang_list.split(","),
+        )
+        print(f" | > Loaded eval dataset from {args.data_dir}. Size: {len(eval_dataset)} ...")
+    else:
+        raise ValueError(" | > Data directory not provided")
+
+    lora_config = LoraConfig(
+        r=args.lora_r,
+        bias="none",
+        inference_mode=False,
+        task_type="SEQ_2_SEQ_LM",
+        lora_alpha=args.lora_alpha,
+        lora_dropout=args.lora_dropout,
+        target_modules=args.lora_target_modules.split(","),
+    )
+
+    model.set_label_smoothing(args.label_smoothing)
+
+    model = get_peft_model(model, lora_config)
+    model.print_trainable_parameters()
+
+    print(f" | > Loading metrics factory with BLEU and chrF ...")
+    seq2seq_compute_metrics = compute_metrics_factory(
+        tokenizer=tokenizer,
+        print_samples=args.print_samples,
+        metric_dict={"BLEU": bleu_metric, "chrF": chrf_metric},
+    )
+
+    training_args = Seq2SeqTrainingArguments(
+        output_dir=args.output_dir,
+        do_train=True,
+        do_eval=True,
+        fp16=True, # use fp16 for faster training
+        logging_strategy="steps",
+        evaluation_strategy="steps",
+        save_strategy="steps",
+        logging_steps=100,
+        save_total_limit=1,
+        predict_with_generate=True,
+        load_best_model_at_end=True,
+        max_steps=args.max_steps, # max_steps overrides num_train_epochs
+        per_device_train_batch_size=args.batch_size,
+        per_device_eval_batch_size=args.batch_size,
+        gradient_accumulation_steps=args.grad_accum_steps,
+        eval_accumulation_steps=args.grad_accum_steps,
+        weight_decay=args.weight_decay,
+        adam_beta1=args.adam_beta1,
+        adam_beta2=args.adam_beta2,
+        max_grad_norm=args.max_grad_norm,
+        optim=args.optimizer,
+        lr_scheduler_type=args.lr_scheduler,
+        warmup_ratio=args.warmup_ratio,
+        warmup_steps=args.warmup_steps,
+        learning_rate=args.learning_rate,
+        num_train_epochs=args.num_train_epochs,
+        save_steps=args.save_steps,
+        eval_steps=args.eval_steps,
+        dataloader_num_workers=args.num_workers,
+        metric_for_best_model=args.metric_for_best_model,
+        greater_is_better=args.greater_is_better,
+        report_to=args.report_to,
+        generation_max_length=256,
+        generation_num_beams=5,
+        sortish_sampler=True,
+        group_by_length=True,
+        include_tokens_per_second=True,
+        include_num_input_tokens_seen=True,
+        dataloader_prefetch_factor=2,
+    )
+
+    # Create Trainer instance
+    trainer = Seq2SeqTrainer(
+        model=model,
+        args=training_args,
+        data_collator=data_collator,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        compute_metrics=seq2seq_compute_metrics,
+        callbacks=[
+            EarlyStoppingCallback(
+                early_stopping_patience=args.patience,
+                early_stopping_threshold=args.threshold,
+            )
+        ],
+    )
+
+    print(f" | > Starting training ...")
+
+    try:
+        trainer.train()
+    except KeyboardInterrupt:
+        print(f" | > Training interrupted ...")
+
+    # this will only save the LoRA adapter weights
+    model.save_pretrained(args.output_dir)
+
+
+
+if __name__ == "__main__":
+    parser = get_arg_parse()
+    args = parser.parse_args()
+
+    main(args)

IndicTrans2/huggingface_interface/train_lora.sh

@@ -0,0 +1,35 @@
+export CUDA_VISIBLE_DEVICES=0
+
+data_dir=${1:-"en-indic-exp"}
+model_name=${2:-"ai4bharat/indictrans2-en-indic-dist-200M"}
+output_dir=${3:-"output"}
+src_lang_list=${4:-"eng_Latn"}
+tgt_lang_list=${5:-"asm_Beng,ben_Beng,guj_Gujr,hin_Deva,kan_Knda,mal_Mlym,mar_Deva,npi_Deva,ory_Orya,pan_Guru,tam_Taml,tel_Telu,urd_Arab"}
+
+python3 train_lora.py \
+    --data_dir $data_dir \
+    --model_name $model_name \
+    --output_dir $output_dir \
+    --src_lang_list $src_lang_list \
+    --tgt_lang_list $tgt_lang_list \
+    --save_steps 1000 \
+    --max_steps 1000000 \
+    --batch_size 32 \
+    --grad_accum_steps 4 \
+    --warmup_steps 4000 \
+    --max_grad_norm 1.0 \
+    --learning_rate 2e-4 \
+    --adam_beta1 0.9 \
+    --adam_beta2 0.98 \
+    --optimizer adamw_torch \
+    --lr_scheduler inverse_sqrt \
+    --num_workers 16 \
+    --metric_for_best_model eval_BLEU \
+    --greater_is_better \
+    --patience 10 \
+    --weight_decay 0.01 \
+    --lora_target_modules "q_proj,k_proj" \
+    --lora_dropout 0.1 \
+    --lora_r 16 \
+    --lora_alpha 32 \
+    --print_samples

IndicTrans2/inference/custom_interactive.py

@@ -0,0 +1,304 @@
+# python wrapper for fairseq-interactive command line tool
+
+#!/usr/bin/env python3 -u
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+"""
+Translate raw text with a trained model. Batches data on-the-fly.
+"""
+
+import os
+import ast
+from collections import namedtuple
+
+import torch
+from fairseq import checkpoint_utils, options, tasks, utils
+from fairseq.dataclass.utils import convert_namespace_to_omegaconf
+from fairseq.token_generation_constraints import pack_constraints, unpack_constraints
+from fairseq_cli.generate import get_symbols_to_strip_from_output
+
+import codecs
+
+PWD = os.path.dirname(__file__)
+Batch = namedtuple("Batch", "ids src_tokens src_lengths constraints")
+Translation = namedtuple("Translation", "src_str hypos pos_scores alignments")
+
+
+def make_batches(
+    lines, cfg, task, max_positions, encode_fn, constrainted_decoding=False
+):
+    def encode_fn_target(x):
+        return encode_fn(x)
+
+    if constrainted_decoding:
+        # Strip (tab-delimited) contraints, if present, from input lines,
+        # store them in batch_constraints
+        batch_constraints = [list() for _ in lines]
+        for i, line in enumerate(lines):
+            if "\t" in line:
+                lines[i], *batch_constraints[i] = line.split("\t")
+
+        # Convert each List[str] to List[Tensor]
+        for i, constraint_list in enumerate(batch_constraints):
+            batch_constraints[i] = [
+                task.target_dictionary.encode_line(
+                    encode_fn_target(constraint),
+                    append_eos=False,
+                    add_if_not_exist=False,
+                )
+                for constraint in constraint_list
+            ]
+
+    if constrainted_decoding:
+        constraints_tensor = pack_constraints(batch_constraints)
+    else:
+        constraints_tensor = None
+
+    tokens, lengths = task.get_interactive_tokens_and_lengths(lines, encode_fn)
+
+    itr = task.get_batch_iterator(
+        dataset=task.build_dataset_for_inference(
+            tokens, lengths, constraints=constraints_tensor
+        ),
+        max_tokens=cfg.dataset.max_tokens,
+        max_sentences=cfg.dataset.batch_size,
+        max_positions=max_positions,
+        ignore_invalid_inputs=cfg.dataset.skip_invalid_size_inputs_valid_test,
+    ).next_epoch_itr(shuffle=False)
+    for batch in itr:
+        ids = batch["id"]
+        src_tokens = batch["net_input"]["src_tokens"]
+        src_lengths = batch["net_input"]["src_lengths"]
+        constraints = batch.get("constraints", None)
+
+        yield Batch(
+            ids=ids,
+            src_tokens=src_tokens,
+            src_lengths=src_lengths,
+            constraints=constraints,
+        )
+
+
+class Translator:
+    """
+    Wrapper class to handle the interaction with fairseq model class for translation
+    """
+    
+    def __init__(
+        self, data_dir, checkpoint_path, batch_size=25, constrained_decoding=False
+    ):
+
+        self.constrained_decoding = constrained_decoding
+        self.parser = options.get_generation_parser(interactive=True)
+        # buffer_size is currently not used but we just initialize it to batch
+        # size + 1 to avoid any assertion errors.
+        if self.constrained_decoding:
+            self.parser.set_defaults(
+                path=checkpoint_path,
+                num_workers=-1,
+                constraints="ordered",
+                batch_size=batch_size,
+                buffer_size=batch_size + 1,
+            )
+        else:
+            self.parser.set_defaults(
+                path=checkpoint_path,
+                remove_bpe="subword_nmt",
+                num_workers=-1,
+                batch_size=batch_size,
+                buffer_size=batch_size + 1,
+            )
+        args = options.parse_args_and_arch(self.parser, input_args=[data_dir])
+        # we are explictly setting src_lang and tgt_lang here
+        # generally the data_dir we pass contains {split}-{src_lang}-{tgt_lang}.*.idx files from
+        # which fairseq infers the src and tgt langs(if these are not passed). In deployment we dont
+        # use any idx files and only store the SRC and TGT dictionaries.
+        args.source_lang = "SRC"
+        args.target_lang = "TGT"
+        # since we are truncating sentences to max_seq_len in engine, we can set it to False here
+        args.skip_invalid_size_inputs_valid_test = False
+
+        # we have custom architechtures in this folder and we will let fairseq
+        # import this
+        args.user_dir = os.path.join(PWD, "model_configs")
+        self.cfg = convert_namespace_to_omegaconf(args)
+
+        utils.import_user_module(self.cfg.common)
+
+        if self.cfg.interactive.buffer_size < 1:
+            self.cfg.interactive.buffer_size = 1
+        if self.cfg.dataset.max_tokens is None and self.cfg.dataset.batch_size is None:
+            self.cfg.dataset.batch_size = 1
+
+        assert (
+            not self.cfg.generation.sampling
+            or self.cfg.generation.nbest == self.cfg.generation.beam
+        ), "--sampling requires --nbest to be equal to --beam"
+        assert (
+            not self.cfg.dataset.batch_size
+            or self.cfg.dataset.batch_size <= self.cfg.interactive.buffer_size
+        ), "--batch-size cannot be larger than --buffer-size"
+
+        # Fix seed for stochastic decoding
+        # if self.cfg.common.seed is not None and not self.cfg.generation.no_seed_provided:
+        #     np.random.seed(self.cfg.common.seed)
+        #     utils.set_torch_seed(self.cfg.common.seed)
+
+        # if not self.constrained_decoding:
+        #     self.use_cuda = torch.cuda.is_available() and not self.cfg.common.cpu
+        # else:
+        #     self.use_cuda = False
+
+        self.use_cuda = torch.cuda.is_available() and not self.cfg.common.cpu
+
+        # Setup task, e.g., translation
+        self.task = tasks.setup_task(self.cfg.task)
+
+        # Load ensemble
+        overrides = ast.literal_eval(self.cfg.common_eval.model_overrides)
+        self.models, self._model_args = checkpoint_utils.load_model_ensemble(
+            utils.split_paths(self.cfg.common_eval.path),
+            arg_overrides=overrides,
+            task=self.task,
+            suffix=self.cfg.checkpoint.checkpoint_suffix,
+            strict=(self.cfg.checkpoint.checkpoint_shard_count == 1),
+            num_shards=self.cfg.checkpoint.checkpoint_shard_count,
+        )
+
+        # Set dictionaries
+        self.src_dict = self.task.source_dictionary
+        self.tgt_dict = self.task.target_dictionary
+
+        # Optimize ensemble for generation
+        for model in self.models:
+            if model is None:
+                continue
+            if self.cfg.common.fp16:
+                model.half()
+            if (
+                self.use_cuda
+                and not self.cfg.distributed_training.pipeline_model_parallel
+            ):
+                model.cuda()
+            model.prepare_for_inference_(self.cfg)
+
+        # Initialize generator
+        self.generator = self.task.build_generator(self.models, self.cfg.generation)
+
+        self.tokenizer = None
+        self.bpe = None
+        # # Handle tokenization and BPE
+        # self.tokenizer = self.task.build_tokenizer(self.cfg.tokenizer)
+        # self.bpe = self.task.build_bpe(self.cfg.bpe)
+
+        # Load alignment dictionary for unknown word replacement
+        # (None if no unknown word replacement, empty if no path to align dictionary)
+        self.align_dict = utils.load_align_dict(self.cfg.generation.replace_unk)
+
+        self.max_positions = utils.resolve_max_positions(
+            self.task.max_positions(), *[model.max_positions() for model in self.models]
+        )
+
+    def encode_fn(self, x):
+        if self.tokenizer is not None:
+            x = self.tokenizer.encode(x)
+        if self.bpe is not None:
+            x = self.bpe.encode(x)
+        return x
+
+    def decode_fn(self, x):
+        if self.bpe is not None:
+            x = self.bpe.decode(x)
+        if self.tokenizer is not None:
+            x = self.tokenizer.decode(x)
+        return x
+
+    def translate(self, inputs, constraints=None):
+        if self.constrained_decoding and constraints is None:
+            raise ValueError("Constraints cant be None in constrained decoding mode")
+        if not self.constrained_decoding and constraints is not None:
+            raise ValueError("Cannot pass constraints during normal translation")
+        if constraints:
+            constrained_decoding = True
+            modified_inputs = []
+            for _input, constraint in zip(inputs, constraints):
+                modified_inputs.append(_input + f"\t{constraint}")
+            inputs = modified_inputs
+        else:
+            constrained_decoding = False
+
+        start_id = 0
+        results = []
+        final_translations = []
+        for batch in make_batches(
+            inputs,
+            self.cfg,
+            self.task,
+            self.max_positions,
+            self.encode_fn,
+            constrained_decoding,
+        ):
+            bsz = batch.src_tokens.size(0)
+            src_tokens = batch.src_tokens
+            src_lengths = batch.src_lengths
+            constraints = batch.constraints
+            if self.use_cuda:
+                src_tokens = src_tokens.cuda()
+                src_lengths = src_lengths.cuda()
+                if constraints is not None:
+                    constraints = constraints.cuda()
+
+            sample = {
+                "net_input": {
+                    "src_tokens": src_tokens,
+                    "src_lengths": src_lengths,
+                },
+            }
+
+            translations = self.task.inference_step(
+                self.generator, self.models, sample, constraints=constraints
+            )
+
+            list_constraints = [[] for _ in range(bsz)]
+            if constrained_decoding:
+                list_constraints = [unpack_constraints(c) for c in constraints]
+            for i, (id, hypos) in enumerate(zip(batch.ids.tolist(), translations)):
+                src_tokens_i = utils.strip_pad(src_tokens[i], self.tgt_dict.pad())
+                constraints = list_constraints[i]
+                results.append(
+                    (
+                        start_id + id,
+                        src_tokens_i,
+                        hypos,
+                        {
+                            "constraints": constraints,
+                        },
+                    )
+                )
+
+        # sort output to match input order
+        for id_, src_tokens, hypos, _ in sorted(results, key=lambda x: x[0]):
+            src_str = ""
+            if self.src_dict is not None:
+                src_str = self.src_dict.string(
+                    src_tokens, self.cfg.common_eval.post_process
+                )
+
+            # Process top predictions
+            for hypo in hypos[: min(len(hypos), self.cfg.generation.nbest)]:
+                hypo_tokens, hypo_str, alignment = utils.post_process_prediction(
+                    hypo_tokens=hypo["tokens"].int().cpu(),
+                    src_str=src_str,
+                    alignment=hypo["alignment"],
+                    align_dict=self.align_dict,
+                    tgt_dict=self.tgt_dict,
+
+                    extra_symbols_to_ignore=get_symbols_to_strip_from_output(
+                        self.generator
+                    ),
+                )
+                detok_hypo_str = self.decode_fn(hypo_str)
+                final_translations.append(detok_hypo_str)
+        return final_translations

IndicTrans2/inference/download.py

@@ -0,0 +1,5 @@
+import urduhack
+urduhack.download()
+
+import nltk
+nltk.download('punkt')

IndicTrans2/inference/engine.py

@@ -0,0 +1,472 @@
+import hashlib
+import os
+import uuid
+from typing import List, Tuple, Union, Dict
+
+import regex as re
+import sentencepiece as spm
+from indicnlp.normalize import indic_normalize
+from indicnlp.tokenize import indic_detokenize, indic_tokenize
+from indicnlp.tokenize.sentence_tokenize import DELIM_PAT_NO_DANDA, sentence_split
+from indicnlp.transliterate import unicode_transliterate
+from mosestokenizer import MosesSentenceSplitter
+from nltk.tokenize import sent_tokenize
+from sacremoses import MosesDetokenizer, MosesPunctNormalizer, MosesTokenizer
+from tqdm import tqdm
+
+from .flores_codes_map_indic import flores_codes, iso_to_flores
+from .normalize_punctuation import punc_norm
+from .normalize_regex_inference import EMAIL_PATTERN, normalize
+
+
+def split_sentences(paragraph: str, lang: str) -> List[str]:
+    """
+    Splits the input text paragraph into sentences. It uses `moses` for English and
+    `indic-nlp` for Indic languages.
+
+    Args:
+        paragraph (str): input text paragraph.
+        lang (str): flores language code.
+
+    Returns:
+        List[str] -> list of sentences.
+    """
+    if lang == "eng_Latn":
+        with MosesSentenceSplitter(flores_codes[lang]) as splitter:
+            sents_moses = splitter([paragraph])
+        sents_nltk = sent_tokenize(paragraph)
+        if len(sents_nltk) < len(sents_moses):
+            sents = sents_nltk
+        else:
+            sents = sents_moses
+        return [sent.replace("\xad", "") for sent in sents]
+    else:
+        return sentence_split(paragraph, lang=flores_codes[lang], delim_pat=DELIM_PAT_NO_DANDA)
+
+
+def add_token(sent: str, src_lang: str, tgt_lang: str, delimiter: str = " ") -> str:
+    """
+    Add special tokens indicating source and target language to the start of the input sentence.
+    The resulting string will have the format: "`{src_lang} {tgt_lang} {input_sentence}`".
+
+    Args:
+        sent (str): input sentence to be translated.
+        src_lang (str): flores lang code of the input sentence.
+        tgt_lang (str): flores lang code in which the input sentence will be translated.
+        delimiter (str): separator to add between language tags and input sentence (default: " ").
+
+    Returns:
+        str: input sentence with the special tokens added to the start.
+    """
+    return src_lang + delimiter + tgt_lang + delimiter + sent
+
+
+def apply_lang_tags(sents: List[str], src_lang: str, tgt_lang: str) -> List[str]:
+    """
+    Add special tokens indicating source and target language to the start of the each input sentence.
+    Each resulting input sentence will have the format: "`{src_lang} {tgt_lang} {input_sentence}`".
+
+    Args:
+        sent (str): input sentence to be translated.
+        src_lang (str): flores lang code of the input sentence.
+        tgt_lang (str): flores lang code in which the input sentence will be translated.
+
+    Returns:
+        List[str]: list of input sentences with the special tokens added to the start.
+    """
+    tagged_sents = []
+    for sent in sents:
+        tagged_sent = add_token(sent.strip(), src_lang, tgt_lang)
+        tagged_sents.append(tagged_sent)
+    return tagged_sents
+
+
+def truncate_long_sentences(
+    sents: List[str], placeholder_entity_map_sents: List[Dict]
+) -> Tuple[List[str], List[Dict]]:
+    """
+    Truncates the sentences that exceed the maximum sequence length.
+    The maximum sequence for the IndicTrans2 model is limited to 256 tokens.
+
+    Args:
+        sents (List[str]): list of input sentences to truncate.
+
+    Returns:
+        Tuple[List[str], List[Dict]]: tuple containing the list of sentences with truncation applied and the updated placeholder entity maps.
+    """
+    MAX_SEQ_LEN = 256
+    new_sents = []
+    placeholders = []
+
+    for j, sent in enumerate(sents):
+        words = sent.split()
+        num_words = len(words)
+        if num_words > MAX_SEQ_LEN:
+            sents = []
+            i = 0
+            while i <= len(words):
+                sents.append(" ".join(words[i : i + MAX_SEQ_LEN]))
+                i += MAX_SEQ_LEN
+            placeholders.extend([placeholder_entity_map_sents[j]] * (len(sents)))
+            new_sents.extend(sents)
+        else:
+            placeholders.append(placeholder_entity_map_sents[j])
+            new_sents.append(sent)
+    return new_sents, placeholders
+
+
+class Model:
+    """
+    Model class to run the IndicTransv2 models using python interface.
+    """
+
+    def __init__(
+        self,
+        ckpt_dir: str,
+        device: str = "cuda",
+        input_lang_code_format: str = "flores",
+        model_type: str = "ctranslate2",
+    ):
+        """
+        Initialize the model class.
+
+        Args:
+            ckpt_dir (str): path of the model checkpoint directory.
+            device (str, optional): where to load the model (defaults: cuda).
+        """
+        self.ckpt_dir = ckpt_dir
+        self.en_tok = MosesTokenizer(lang="en")
+        self.en_normalizer = MosesPunctNormalizer()
+        self.en_detok = MosesDetokenizer(lang="en")
+        self.xliterator = unicode_transliterate.UnicodeIndicTransliterator()
+
+        print("Initializing sentencepiece model for SRC and TGT")
+        self.sp_src = spm.SentencePieceProcessor(
+            model_file=os.path.join(ckpt_dir, "vocab", "model.SRC")
+        )
+        self.sp_tgt = spm.SentencePieceProcessor(
+            model_file=os.path.join(ckpt_dir, "vocab", "model.TGT")
+        )
+
+        self.input_lang_code_format = input_lang_code_format
+
+        print("Initializing model for translation")
+        # initialize the model
+        if model_type == "ctranslate2":
+            import ctranslate2
+
+            self.translator = ctranslate2.Translator(
+                self.ckpt_dir, device=device
+            )  # , compute_type="auto")
+            self.translate_lines = self.ctranslate2_translate_lines
+        elif model_type == "fairseq":
+            from .custom_interactive import Translator
+
+            self.translator = Translator(
+                data_dir=os.path.join(self.ckpt_dir, "final_bin"),
+                checkpoint_path=os.path.join(self.ckpt_dir, "model", "checkpoint_best.pt"),
+                batch_size=100,
+            )
+            self.translate_lines = self.fairseq_translate_lines
+        else:
+            raise NotImplementedError(f"Unknown model_type: {model_type}")
+
+    def ctranslate2_translate_lines(self, lines: List[str]) -> List[str]:
+        tokenized_sents = [x.strip().split(" ") for x in lines]
+        translations = self.translator.translate_batch(
+            tokenized_sents,
+            max_batch_size=9216,
+            batch_type="tokens",
+            max_input_length=160,
+            max_decoding_length=256,
+            beam_size=5,
+        )
+        translations = [" ".join(x.hypotheses[0]) for x in translations]
+        return translations
+
+    def fairseq_translate_lines(self, lines: List[str]) -> List[str]:
+        return self.translator.translate(lines)
+
+    def paragraphs_batch_translate__multilingual(self, batch_payloads: List[tuple]) -> List[str]:
+        """
+        Translates a batch of input paragraphs (including pre/post processing)
+        from any language to any language.
+
+        Args:
+            batch_payloads (List[tuple]): batch of long input-texts to be translated, each in format: (paragraph, src_lang, tgt_lang)
+
+        Returns:
+            List[str]: batch of paragraph-translations in the respective languages.
+        """
+        paragraph_id_to_sentence_range = []
+        global__sents = []
+        global__preprocessed_sents = []
+        global__preprocessed_sents_placeholder_entity_map = []
+
+        for i in range(len(batch_payloads)):
+            paragraph, src_lang, tgt_lang = batch_payloads[i]
+            if self.input_lang_code_format == "iso":
+                src_lang, tgt_lang = iso_to_flores[src_lang], iso_to_flores[tgt_lang]
+
+            batch = split_sentences(paragraph, src_lang)
+            global__sents.extend(batch)
+
+            preprocessed_sents, placeholder_entity_map_sents = self.preprocess_batch(
+                batch, src_lang, tgt_lang
+            )
+
+            global_sentence_start_index = len(global__preprocessed_sents)
+            global__preprocessed_sents.extend(preprocessed_sents)
+            global__preprocessed_sents_placeholder_entity_map.extend(placeholder_entity_map_sents)
+            paragraph_id_to_sentence_range.append(
+                (global_sentence_start_index, len(global__preprocessed_sents))
+            )
+
+        translations = self.translate_lines(global__preprocessed_sents)
+
+        translated_paragraphs = []
+        for paragraph_id, sentence_range in enumerate(paragraph_id_to_sentence_range):
+            tgt_lang = batch_payloads[paragraph_id][2]
+            if self.input_lang_code_format == "iso":
+                tgt_lang = iso_to_flores[tgt_lang]
+
+            postprocessed_sents = self.postprocess(
+                translations[sentence_range[0] : sentence_range[1]],
+                global__preprocessed_sents_placeholder_entity_map[
+                    sentence_range[0] : sentence_range[1]
+                ],
+                tgt_lang,
+            )
+            translated_paragraph = " ".join(postprocessed_sents)
+            translated_paragraphs.append(translated_paragraph)
+
+        return translated_paragraphs
+
+    # translate a batch of sentences from src_lang to tgt_lang
+    def batch_translate(self, batch: List[str], src_lang: str, tgt_lang: str) -> List[str]:
+        """
+        Translates a batch of input sentences (including pre/post processing)
+        from source language to target language.
+
+        Args:
+            batch (List[str]): batch of input sentences to be translated.
+            src_lang (str): flores source language code.
+            tgt_lang (str): flores target language code.
+
+        Returns:
+            List[str]: batch of translated-sentences generated by the model.
+        """
+
+        assert isinstance(batch, list)
+
+        if self.input_lang_code_format == "iso":
+            src_lang, tgt_lang = iso_to_flores[src_lang], iso_to_flores[tgt_lang]
+
+        preprocessed_sents, placeholder_entity_map_sents = self.preprocess_batch(
+            batch, src_lang, tgt_lang
+        )
+        translations = self.translate_lines(preprocessed_sents)
+        return self.postprocess(translations, placeholder_entity_map_sents, tgt_lang)
+
+    # translate a paragraph from src_lang to tgt_lang
+    def translate_paragraph(self, paragraph: str, src_lang: str, tgt_lang: str) -> str:
+        """
+        Translates an input text paragraph (including pre/post processing)
+        from source language to target language.
+
+        Args:
+            paragraph (str): input text paragraph to be translated.
+            src_lang (str): flores source language code.
+            tgt_lang (str): flores target language code.
+
+        Returns:
+            str: paragraph translation generated by the model.
+        """
+
+        assert isinstance(paragraph, str)
+
+        if self.input_lang_code_format == "iso":
+            flores_src_lang = iso_to_flores[src_lang]
+        else:
+            flores_src_lang = src_lang
+
+        sents = split_sentences(paragraph, flores_src_lang)
+        postprocessed_sents = self.batch_translate(sents, src_lang, tgt_lang)
+        translated_paragraph = " ".join(postprocessed_sents)
+
+        return translated_paragraph
+
+    def preprocess_batch(self, batch: List[str], src_lang: str, tgt_lang: str) -> List[str]:
+        """
+        Preprocess an array of sentences by normalizing, tokenization, and possibly transliterating it. It also tokenizes the
+        normalized text sequences using sentence piece tokenizer and also adds language tags.
+
+        Args:
+            batch (List[str]): input list of sentences to preprocess.
+            src_lang (str): flores language code of the input text sentences.
+            tgt_lang (str): flores language code of the output text sentences.
+
+        Returns:
+            Tuple[List[str], List[Dict]]: a tuple of list of preprocessed input text sentences and also a corresponding list of dictionary
+                mapping placeholders to their original values.
+        """
+        preprocessed_sents, placeholder_entity_map_sents = self.preprocess(batch, lang=src_lang)
+        tokenized_sents = self.apply_spm(preprocessed_sents)
+        tokenized_sents, placeholder_entity_map_sents = truncate_long_sentences(
+            tokenized_sents, placeholder_entity_map_sents
+        )
+        tagged_sents = apply_lang_tags(tokenized_sents, src_lang, tgt_lang)
+        return tagged_sents, placeholder_entity_map_sents
+
+    def apply_spm(self, sents: List[str]) -> List[str]:
+        """
+        Applies sentence piece encoding to the batch of input sentences.
+
+        Args:
+            sents (List[str]): batch of the input sentences.
+
+        Returns:
+            List[str]: batch of encoded sentences with sentence piece model
+        """
+        return [" ".join(self.sp_src.encode(sent, out_type=str)) for sent in sents]
+
+    def preprocess_sent(
+        self,
+        sent: str,
+        normalizer: Union[MosesPunctNormalizer, indic_normalize.IndicNormalizerFactory],
+        lang: str,
+    ) -> Tuple[str, Dict]:
+        """
+        Preprocess an input text sentence by normalizing, tokenization, and possibly transliterating it.
+
+        Args:
+            sent (str): input text sentence to preprocess.
+            normalizer (Union[MosesPunctNormalizer, indic_normalize.IndicNormalizerFactory]): an object that performs normalization on the text.
+            lang (str): flores language code of the input text sentence.
+
+        Returns:
+            Tuple[str, Dict]: A tuple containing the preprocessed input text sentence and a corresponding dictionary
+            mapping placeholders to their original values.
+        """
+        iso_lang = flores_codes[lang]
+        sent = punc_norm(sent, iso_lang)
+        sent, placeholder_entity_map = normalize(sent)
+
+        transliterate = True
+        if lang.split("_")[1] in ["Arab", "Aran", "Olck", "Mtei", "Latn"]:
+            transliterate = False
+
+        if iso_lang == "en":
+            processed_sent = " ".join(
+                self.en_tok.tokenize(self.en_normalizer.normalize(sent.strip()), escape=False)
+            )
+        elif transliterate:
+            # transliterates from the any specific language to devanagari
+            # which is why we specify lang2_code as "hi".
+            processed_sent = self.xliterator.transliterate(
+                " ".join(
+                    indic_tokenize.trivial_tokenize(normalizer.normalize(sent.strip()), iso_lang)
+                ),
+                iso_lang,
+                "hi",
+            ).replace(" ् ", "्")
+        else:
+            # we only need to transliterate for joint training
+            processed_sent = " ".join(
+                indic_tokenize.trivial_tokenize(normalizer.normalize(sent.strip()), iso_lang)
+            )
+
+        return processed_sent, placeholder_entity_map
+
+    def preprocess(self, sents: List[str], lang: str):
+        """
+        Preprocess an array of sentences by normalizing, tokenization, and possibly transliterating it.
+
+        Args:
+            batch (List[str]): input list of sentences to preprocess.
+            lang (str): flores language code of the input text sentences.
+
+        Returns:
+            Tuple[List[str], List[Dict]]: a tuple of list of preprocessed input text sentences and also a corresponding list of dictionary
+                mapping placeholders to their original values.
+        """
+        processed_sents, placeholder_entity_map_sents = [], []
+
+        if lang == "eng_Latn":
+            normalizer = None
+        else:
+            normfactory = indic_normalize.IndicNormalizerFactory()
+            normalizer = normfactory.get_normalizer(flores_codes[lang])
+
+        for sent in sents:
+            sent, placeholder_entity_map = self.preprocess_sent(sent, normalizer, lang)
+            processed_sents.append(sent)
+            placeholder_entity_map_sents.append(placeholder_entity_map)
+
+        return processed_sents, placeholder_entity_map_sents
+
+    def postprocess(
+        self,
+        sents: List[str],
+        placeholder_entity_map: List[Dict],
+        lang: str,
+        common_lang: str = "hin_Deva",
+    ) -> List[str]:
+        """
+        Postprocesses a batch of input sentences after the translation generations.
+
+        Args:
+            sents (List[str]): batch of translated sentences to postprocess.
+            placeholder_entity_map (List[Dict]): dictionary mapping placeholders to the original entity values.
+            lang (str): flores language code of the input sentences.
+            common_lang (str, optional): flores language code of the transliterated language (defaults: hin_Deva).
+
+        Returns:
+            List[str]: postprocessed batch of input sentences.
+        """
+
+        lang_code, script_code = lang.split("_")
+        # SPM decode
+        for i in range(len(sents)):
+            # sent_tokens = sents[i].split(" ")
+            # sents[i] = self.sp_tgt.decode(sent_tokens)
+
+            sents[i] = sents[i].replace(" ", "").replace("▁", " ").strip()
+
+            # Fixes for Perso-Arabic scripts
+            # TODO: Move these normalizations inside indic-nlp-library
+            if script_code in {"Arab", "Aran"}:
+                # UrduHack adds space before punctuations. Since the model was trained without fixing this issue, let's fix it now
+                sents[i] = sents[i].replace(" ؟", "؟").replace(" ۔", "۔").replace(" ،", "،")
+                # Kashmiri bugfix for palatalization: https://github.com/AI4Bharat/IndicTrans2/issues/11
+                sents[i] = sents[i].replace("ٮ۪", "ؠ")
+
+        assert len(sents) == len(placeholder_entity_map)
+
+        for i in range(0, len(sents)):
+            for key in placeholder_entity_map[i].keys():
+                sents[i] = sents[i].replace(key, placeholder_entity_map[i][key])
+
+        # Detokenize and transliterate to native scripts if applicable
+        postprocessed_sents = []
+
+        if lang == "eng_Latn":
+            for sent in sents:
+                postprocessed_sents.append(self.en_detok.detokenize(sent.split(" ")))
+        else:
+            for sent in sents:
+                outstr = indic_detokenize.trivial_detokenize(
+                    self.xliterator.transliterate(
+                        sent, flores_codes[common_lang], flores_codes[lang]
+                    ),
+                    flores_codes[lang],
+                )
+                
+                # Oriya bug: indic-nlp-library produces ଯ଼ instead of ୟ when converting from Devanagari to Odia
+                # TODO: Find out what's the issue with unicode transliterator for Oriya and fix it
+                if lang_code == "ory":
+                    outstr = outstr.replace("ଯ଼", 'ୟ')
+
+                postprocessed_sents.append(outstr)
+
+        return postprocessed_sents

IndicTrans2/inference/flores_codes_map_indic.py

@@ -0,0 +1,83 @@
+"""
+FLORES language code mapping to 2 letter ISO language code for compatibility 
+with Indic NLP Library (https://github.com/anoopkunchukuttan/indic_nlp_library)
+"""
+flores_codes = {
+    "asm_Beng": "as",
+    "awa_Deva": "hi",
+    "ben_Beng": "bn",
+    "bho_Deva": "hi",
+    "brx_Deva": "hi",
+    "doi_Deva": "hi",
+    "eng_Latn": "en",
+    "gom_Deva": "kK",
+    "guj_Gujr": "gu",
+    "hin_Deva": "hi",
+    "hne_Deva": "hi",
+    "kan_Knda": "kn",
+    "kas_Arab": "ur",
+    "kas_Deva": "hi",
+    "kha_Latn": "en",
+    "lus_Latn": "en",
+    "mag_Deva": "hi",
+    "mai_Deva": "hi",
+    "mal_Mlym": "ml",
+    "mar_Deva": "mr",
+    "mni_Beng": "bn",
+    "mni_Mtei": "hi",
+    "npi_Deva": "ne",
+    "ory_Orya": "or",
+    "pan_Guru": "pa",
+    "san_Deva": "hi",
+    "sat_Olck": "or",
+    "snd_Arab": "ur",
+    "snd_Deva": "hi",
+    "tam_Taml": "ta",
+    "tel_Telu": "te",
+    "urd_Arab": "ur",
+}
+
+
+flores_to_iso = {
+    "asm_Beng": "as",
+    "awa_Deva": "awa",
+    "ben_Beng": "bn",
+    "bho_Deva": "bho",
+    "brx_Deva": "brx",
+    "doi_Deva": "doi",
+    "eng_Latn": "en",
+    "gom_Deva": "gom",
+    "guj_Gujr": "gu",
+    "hin_Deva": "hi",
+    "hne_Deva": "hne",
+    "kan_Knda": "kn",
+    "kas_Arab": "ksa",
+    "kas_Deva": "ksd",
+    "kha_Latn": "kha",
+    "lus_Latn": "lus",
+    "mag_Deva": "mag",
+    "mai_Deva": "mai",
+    "mal_Mlym": "ml",
+    "mar_Deva": "mr",
+    "mni_Beng": "mnib",
+    "mni_Mtei": "mnim",
+    "npi_Deva": "ne",
+    "ory_Orya": "or",
+    "pan_Guru": "pa",
+    "san_Deva": "sa",
+    "sat_Olck": "sat",
+    "snd_Arab": "sda",
+    "snd_Deva": "sdd",
+    "tam_Taml": "ta",
+    "tel_Telu": "te",
+    "urd_Arab": "ur",
+}
+
+iso_to_flores = {iso_code: flores_code for flores_code, iso_code in flores_to_iso.items()}
+# Patch for digraphic langs.
+iso_to_flores["ks"] = "kas_Arab"
+iso_to_flores["ks_Deva"] = "kas_Deva"
+iso_to_flores["mni"] = "mni_Mtei"
+iso_to_flores["mni_Beng"] = "mni_Beng"
+iso_to_flores["sd"] = "snd_Arab"
+iso_to_flores["sd_Deva"] = "snd_Deva"

IndicTrans2/inference/indic_num_map.py

@@ -0,0 +1,117 @@
+"""
+A dictionary mapping intended to normalize the numerals in Indic languages from 
+native script to Roman script. This is done to ensure that the figures / numbers 
+mentioned in native script are perfectly preserved during translation.
+"""
+INDIC_NUM_MAP = {
+    "\u09e6": "0",
+    "0": "0",
+    "\u0ae6": "0",
+    "\u0ce6": "0",
+    "\u0966": "0",
+    "\u0660": "0",
+    "\uabf0": "0",
+    "\u0b66": "0",
+    "\u0a66": "0",
+    "\u1c50": "0",
+    "\u06f0": "0",
+    "\u09e7": "1",
+    "1": "1",
+    "\u0ae7": "1",
+    "\u0967": "1",
+    "\u0ce7": "1",
+    "\u06f1": "1",
+    "\uabf1": "1",
+    "\u0b67": "1",
+    "\u0a67": "1",
+    "\u1c51": "1",
+    "\u0c67": "1",
+    "\u09e8": "2",
+    "2": "2",
+    "\u0ae8": "2",
+    "\u0968": "2",
+    "\u0ce8": "2",
+    "\u06f2": "2",
+    "\uabf2": "2",
+    "\u0b68": "2",
+    "\u0a68": "2",
+    "\u1c52": "2",
+    "\u0c68": "2",
+    "\u09e9": "3",
+    "3": "3",
+    "\u0ae9": "3",
+    "\u0969": "3",
+    "\u0ce9": "3",
+    "\u06f3": "3",
+    "\uabf3": "3",
+    "\u0b69": "3",
+    "\u0a69": "3",
+    "\u1c53": "3",
+    "\u0c69": "3",
+    "\u09ea": "4",
+    "4": "4",
+    "\u0aea": "4",
+    "\u096a": "4",
+    "\u0cea": "4",
+    "\u06f4": "4",
+    "\uabf4": "4",
+    "\u0b6a": "4",
+    "\u0a6a": "4",
+    "\u1c54": "4",
+    "\u0c6a": "4",
+    "\u09eb": "5",
+    "5": "5",
+    "\u0aeb": "5",
+    "\u096b": "5",
+    "\u0ceb": "5",
+    "\u06f5": "5",
+    "\uabf5": "5",
+    "\u0b6b": "5",
+    "\u0a6b": "5",
+    "\u1c55": "5",
+    "\u0c6b": "5",
+    "\u09ec": "6",
+    "6": "6",
+    "\u0aec": "6",
+    "\u096c": "6",
+    "\u0cec": "6",
+    "\u06f6": "6",
+    "\uabf6": "6",
+    "\u0b6c": "6",
+    "\u0a6c": "6",
+    "\u1c56": "6",
+    "\u0c6c": "6",
+    "\u09ed": "7",
+    "7": "7",
+    "\u0aed": "7",
+    "\u096d": "7",
+    "\u0ced": "7",
+    "\u06f7": "7",
+    "\uabf7": "7",
+    "\u0b6d": "7",
+    "\u0a6d": "7",
+    "\u1c57": "7",
+    "\u0c6d": "7",
+    "\u09ee": "8",
+    "8": "8",
+    "\u0aee": "8",
+    "\u096e": "8",
+    "\u0cee": "8",
+    "\u06f8": "8",
+    "\uabf8": "8",
+    "\u0b6e": "8",
+    "\u0a6e": "8",
+    "\u1c58": "8",
+    "\u0c6e": "8",
+    "\u09ef": "9",
+    "9": "9",
+    "\u0aef": "9",
+    "\u096f": "9",
+    "\u0cef": "9",
+    "\u06f9": "9",
+    "\uabf9": "9",
+    "\u0b6f": "9",
+    "\u0a6f": "9",
+    "\u1c59": "9",
+    "\u0c6f": "9",
+}

IndicTrans2/inference/model_configs/__init__.py

@@ -0,0 +1 @@
+from . import custom_transformer

IndicTrans2/inference/model_configs/custom_transformer.py

@@ -0,0 +1,82 @@
+from fairseq.models import register_model_architecture
+from fairseq.models.transformer import base_architecture
+
+
+@register_model_architecture("transformer", "transformer_2x")
+def transformer_big(args):
+    args.encoder_embed_dim = getattr(args, "encoder_embed_dim", 1024)
+    args.encoder_ffn_embed_dim = getattr(args, "encoder_ffn_embed_dim", 4096)
+    args.encoder_attention_heads = getattr(args, "encoder_attention_heads", 16)
+    args.encoder_normalize_before = getattr(args, "encoder_normalize_before", False)
+    args.decoder_embed_dim = getattr(args, "decoder_embed_dim", 1024)
+    args.decoder_ffn_embed_dim = getattr(args, "decoder_ffn_embed_dim", 4096)
+    args.decoder_attention_heads = getattr(args, "decoder_attention_heads", 16)
+    base_architecture(args)
+
+
+@register_model_architecture("transformer", "transformer_4x")
+def transformer_huge(args):
+    args.encoder_embed_dim = getattr(args, "encoder_embed_dim", 1536)
+    args.encoder_ffn_embed_dim = getattr(args, "encoder_ffn_embed_dim", 4096)
+    args.encoder_attention_heads = getattr(args, "encoder_attention_heads", 16)
+    args.encoder_normalize_before = getattr(args, "encoder_normalize_before", False)
+    args.decoder_embed_dim = getattr(args, "decoder_embed_dim", 1536)
+    args.decoder_ffn_embed_dim = getattr(args, "decoder_ffn_embed_dim", 4096)
+    args.decoder_attention_heads = getattr(args, "decoder_attention_heads", 16)
+    base_architecture(args)
+
+
+@register_model_architecture("transformer", "transformer_9x")
+def transformer_xlarge(args):
+    args.encoder_embed_dim = getattr(args, "encoder_embed_dim", 2048)
+    args.encoder_ffn_embed_dim = getattr(args, "encoder_ffn_embed_dim", 8192)
+    args.encoder_attention_heads = getattr(args, "encoder_attention_heads", 16)
+    args.encoder_normalize_before = getattr(args, "encoder_normalize_before", False)
+    args.decoder_embed_dim = getattr(args, "decoder_embed_dim", 2048)
+    args.decoder_ffn_embed_dim = getattr(args, "decoder_ffn_embed_dim", 8192)
+    args.decoder_attention_heads = getattr(args, "decoder_attention_heads", 16)
+    base_architecture(args)
+
+
+@register_model_architecture("transformer", "transformer_12e12d_9xeq")
+def transformer_vxlarge(args):
+    args.encoder_embed_dim = getattr(args, "encoder_embed_dim", 1536)
+    args.encoder_ffn_embed_dim = getattr(args, "encoder_ffn_embed_dim", 4096)
+    args.encoder_attention_heads = getattr(args, "encoder_attention_heads", 16)
+    args.encoder_normalize_before = getattr(args, "encoder_normalize_before", False)
+    args.decoder_embed_dim = getattr(args, "decoder_embed_dim", 1536)
+    args.decoder_ffn_embed_dim = getattr(args, "decoder_ffn_embed_dim", 4096)
+    args.decoder_attention_heads = getattr(args, "decoder_attention_heads", 16)
+    args.encoder_layers = getattr(args, "encoder_layers", 12)
+    args.decoder_layers = getattr(args, "decoder_layers", 12)
+    base_architecture(args)
+
+
+@register_model_architecture("transformer", "transformer_18_18")
+def transformer_deep(args):
+    args.encoder_embed_dim = getattr(args, "encoder_embed_dim", 1024)
+    args.encoder_ffn_embed_dim = getattr(args, "encoder_ffn_embed_dim", 8 * 1024)
+    args.encoder_attention_heads = getattr(args, "encoder_attention_heads", 16)
+    args.encoder_normalize_before = getattr(args, "encoder_normalize_before", True)
+    args.decoder_normalize_before = getattr(args, "decoder_normalize_before", True)
+    args.decoder_embed_dim = getattr(args, "decoder_embed_dim", 1024)
+    args.decoder_ffn_embed_dim = getattr(args, "decoder_ffn_embed_dim", 8 * 1024)
+    args.decoder_attention_heads = getattr(args, "decoder_attention_heads", 16)
+    args.encoder_layers = getattr(args, "encoder_layers", 18)
+    args.decoder_layers = getattr(args, "decoder_layers", 18)
+    base_architecture(args)
+
+
+@register_model_architecture("transformer", "transformer_24_24")
+def transformer_xdeep(args):
+    args.encoder_embed_dim = getattr(args, "encoder_embed_dim", 1024)
+    args.encoder_ffn_embed_dim = getattr(args, "encoder_ffn_embed_dim", 8 * 1024)
+    args.encoder_attention_heads = getattr(args, "encoder_attention_heads", 16)
+    args.encoder_normalize_before = getattr(args, "encoder_normalize_before", True)
+    args.decoder_normalize_before = getattr(args, "decoder_normalize_before", True)
+    args.decoder_embed_dim = getattr(args, "decoder_embed_dim", 1024)
+    args.decoder_ffn_embed_dim = getattr(args, "decoder_ffn_embed_dim", 8 * 1024)
+    args.decoder_attention_heads = getattr(args, "decoder_attention_heads", 16)
+    args.encoder_layers = getattr(args, "encoder_layers", 24)
+    args.decoder_layers = getattr(args, "decoder_layers", 24)
+    base_architecture(args)

IndicTrans2/inference/normalize-punctuation.perl

@@ -0,0 +1,90 @@
+#!/usr/bin/env perl
+#
+# This file is part of moses.  Its use is licensed under the GNU Lesser General
+# Public License version 2.1 or, at your option, any later version.
+
+use warnings;
+use strict;
+
+my $language = "en";
+my $PENN = 0;
+
+while (@ARGV) {
+    $_ = shift;
+    /^-b$/ && ($| = 1, next); # not buffered (flush each line)
+    /^-l$/ && ($language = shift, next);
+    /^[^\-]/ && ($language = $_, next);
+  	/^-penn$/ && ($PENN = 1, next);
+}
+
+while(<STDIN>) {
+    s/\r//g;
+    # remove extra spaces
+    s/\(/ \(/g;
+    s/\)/\) /g; s/ +/ /g;
+    s/\) ([\.\!\:\?\;\,])/\)$1/g;
+    s/\( /\(/g;
+    s/ \)/\)/g;
+    s/(\d) \%/$1\%/g;
+    s/ :/:/g;
+    s/ ;/;/g;
+    # normalize unicode punctuation
+    if ($PENN == 0) {
+      s/\`/\'/g;
+      s/\'\'/ \" /g;
+    }
+
+    s/„/\"/g;
+    s/“/\"/g;
+    s/”/\"/g;
+    s/–/-/g;
+    s/—/ - /g; s/ +/ /g;
+    s/´/\'/g;
+    s/([a-z])‘([a-z])/$1\'$2/gi;
+    s/([a-z])’([a-z])/$1\'$2/gi;
+    s/‘/\'/g;
+    s/‚/\'/g;
+    s/’/\"/g;
+    s/''/\"/g;
+    s/´´/\"/g;
+    s/…/.../g;
+    # French quotes
+    s/ « / \"/g;
+    s/« /\"/g;
+    s/«/\"/g;
+    s/ » /\" /g;
+    s/ »/\"/g;
+    s/»/\"/g;
+    # handle pseudo-spaces
+    s/ \%/\%/g;
+    s/nº /nº /g;
+    s/ :/:/g;
+    s/ ºC/ ºC/g;
+    s/ cm/ cm/g;
+    s/ \?/\?/g;
+    s/ \!/\!/g;
+    s/ ;/;/g;
+    s/, /, /g; s/ +/ /g;
+
+    # English "quotation," followed by comma, style
+    if ($language eq "en") {
+	s/\"([,\.]+)/$1\"/g;
+    }
+    # Czech is confused
+    elsif ($language eq "cs" || $language eq "cz") {
+    }
+    # German/Spanish/French "quotation", followed by comma, style
+    else {
+	s/,\"/\",/g;	
+	s/(\.+)\"(\s*[^<])/\"$1$2/g; # don't fix period at end of sentence
+    }
+
+
+    if ($language eq "de" || $language eq "es" || $language eq "cz" || $language eq "cs" || $language eq "fr") {
+	s/(\d) (\d)/$1,$2/g;
+    }
+    else {
+	s/(\d) (\d)/$1.$2/g;
+    }
+    print $_;
+}

IndicTrans2/inference/normalize_punctuation.py

@@ -0,0 +1,60 @@
+# IMPORTANT NOTE: DO NOT DIRECTLY EDIT THIS FILE
+# This file was manually ported from `normalize-punctuation.perl`
+# TODO: Only supports English, add others
+
+import regex as re
+multispace_regex = re.compile("[ ]{2,}")
+multidots_regex = re.compile(r"\.{2,}")
+end_bracket_space_punc_regex = re.compile(r"\) ([\.!:?;,])")
+digit_space_percent = re.compile(r"(\d) %")
+double_quot_punc = re.compile(r"\"([,\.]+)")
+digit_nbsp_digit = re.compile(r"(\d) (\d)")
+
+def punc_norm(text, lang="en"):
+    text = text.replace('\r', '') \
+                .replace('(', " (") \
+                .replace(')', ") ") \
+                \
+                .replace("( ", "(") \
+                .replace(" )", ")") \
+                \
+                .replace(" :", ':') \
+                .replace(" ;", ';') \
+                .replace('`', "'") \
+                \
+                .replace('„', '"') \
+                .replace('“', '"') \
+                .replace('”', '"') \
+                .replace('–', '-') \
+                .replace('—', " - ") \
+                .replace('´', "'") \
+                .replace('‘', "'") \
+                .replace('‚', "'") \
+                .replace('’', "'") \
+                .replace("''", "\"") \
+                .replace("´´", '"') \
+                .replace('…', "...") \
+                .replace(" « ", " \"") \
+                .replace("« ", '"') \
+                .replace('«', '"') \
+                .replace(" » ", "\" ") \
+                .replace(" »", '"') \
+                .replace('»', '"') \
+                .replace(" %", '%') \
+                .replace("nº ", "nº ") \
+                .replace(" :", ':') \
+                .replace(" ºC", " ºC") \
+                .replace(" cm", " cm") \
+                .replace(" ?", '?') \
+                .replace(" !", '!') \
+                .replace(" ;", ';') \
+                .replace(", ", ", ") \
+                
+    
+    text = multispace_regex.sub(' ', text)
+    text = multidots_regex.sub('.', text)
+    text = end_bracket_space_punc_regex.sub(r")\1", text)
+    text = digit_space_percent.sub(r"\1%", text)
+    text = double_quot_punc.sub(r'\1"', text) # English "quotation," followed by comma, style
+    text = digit_nbsp_digit.sub(r"\1.\2", text) # What does it mean?
+    return text.strip(' ')

IndicTrans2/inference/normalize_punctuation.sh

@@ -0,0 +1,33 @@
+set -euo pipefail
+
+root=$(dirname $0)
+
+lang_map_path=$root/utils.map_token_lang.tsv
+
+usage () {
+    echo "usage: $0 lang" >&2
+    exit 1
+}
+
+[ $# -eq 1 ] || usage
+
+lang=$1
+
+declare -A lang_map
+
+while read line; do
+    key=$(cut -f1 <<< "$line")
+    val=$(cut -f2 <<< "$line")
+    lang_map[$key]=$val
+done < $lang_map_path
+
+if [ -v "lang_map[$lang]" ]; then
+    lang=${lang_map[$lang]}
+elif [ -v "lang_map[${lang:0:3}]" ]; then
+    lang=${lang_map[${lang:0:3}]}
+else
+    echo "undefined mapping: ${lang}, falling back to: en" >&2
+    lang=en
+fi
+
+perl $root/normalize-punctuation.perl $lang

IndicTrans2/inference/normalize_regex_inference.py

@@ -0,0 +1,105 @@
+from typing import Tuple
+import regex as re
+import sys
+from tqdm import tqdm
+from .indic_num_map import INDIC_NUM_MAP
+
+
+URL_PATTERN = r'\b(?<![\w/.])(?:(?:https?|ftp)://)?(?:(?:[\w-]+\.)+(?!\.))(?:[\w/\-?#&=%.]+)+(?!\.\w+)\b'
+EMAIL_PATTERN = r'[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}'
+# handles dates, time, percentages, proportion, ratio, etc
+NUMERAL_PATTERN = r"(~?\d+\.?\d*\s?%?\s?-?\s?~?\d+\.?\d*\s?%|~?\d+%|\d+[-\/.,:']\d+[-\/.,:'+]\d+(?:\.\d+)?|\d+[-\/.:'+]\d+(?:\.\d+)?)"
+# handles upi, social media handles and hashtags
+OTHER_PATTERN = r'[A-Za-z0-9]*[#|@]\w+'
+
+
+def normalize_indic_numerals(line: str):
+    """
+    Normalize the numerals in Indic languages from native script to Roman script (if present).
+    
+    Args:
+        line (str): an input string with Indic numerals to be normalized.
+    
+    Returns:
+        str: an input string with the all Indic numerals normalized to Roman script.
+    """
+    return "".join([INDIC_NUM_MAP.get(c, c) for c in line])
+
+
+def wrap_with_placeholders(text: str, patterns: list) -> Tuple[str, dict]:
+    """
+    Wraps substrings with matched patterns in the given text with placeholders and returns
+    the modified text along with a mapping of the placeholders to their original value.
+    
+    Args:
+        text (str): an input string which needs to be wrapped with the placeholders.
+        pattern (list): list of patterns to search for in the input string.
+    
+    Returns:
+        Tuple[str, dict]: a tuple containing the modified text and a dictionary mapping 
+            placeholders to their original values.
+    """
+    serial_no = 1
+    
+    placeholder_entity_map = dict()
+    
+    for pattern in patterns:
+        matches = set(re.findall(pattern, text))
+        
+        # wrap common match with placeholder tags
+        for match in matches:
+            if pattern==URL_PATTERN :
+                #Avoids false positive URL matches for names with initials.
+                temp = match.replace(".",'')
+                if len(temp)<4:
+                    continue
+            if pattern==NUMERAL_PATTERN :
+                #Short numeral patterns do not need placeholder based handling.
+                temp = match.replace(" ",'').replace(".",'').replace(":",'')
+                if len(temp)<4:
+                    continue
+            
+            #Set of Translations of "ID" in all the suppported languages have been collated.            
+            #This has been added to deal with edge cases where placeholders might get translated. 
+            indic_failure_cases = ['آی ڈی ', 'ꯑꯥꯏꯗꯤ', 'आईडी', 'आई . डी . ', 'ऐटि', 'آئی ڈی ', 'ᱟᱭᱰᱤ ᱾', 'आयडी', 'ऐडि', 'आइडि']         
+            placeholder = "<ID{}>".format(serial_no)
+            alternate_placeholder = "< ID{} >".format(serial_no)                    
+            placeholder_entity_map[placeholder] = match
+            placeholder_entity_map[alternate_placeholder] = match
+            
+            for i in indic_failure_cases:
+                placeholder_temp = "<{}{}>".format(i,serial_no)
+                placeholder_entity_map[placeholder_temp] = match
+                placeholder_temp = "< {}{} >".format(i, serial_no)
+                placeholder_entity_map[placeholder_temp] = match
+                placeholder_temp = "< {} {} >".format(i, serial_no)
+                placeholder_entity_map[placeholder_temp] = match
+            
+            text = text.replace(match, placeholder)
+            serial_no+=1
+    
+    text = re.sub("\s+", " ", text)
+    
+    #Regex has failure cases in trailing "/" in URLs, so this is a workaround. 
+    text = text.replace(">/",">")
+        
+    return text, placeholder_entity_map
+
+
+def normalize(text: str, patterns: list = [EMAIL_PATTERN, URL_PATTERN, NUMERAL_PATTERN, OTHER_PATTERN]) -> Tuple[str, dict]:
+    """
+    Normalizes and wraps the spans of input string with placeholder tags. It first normalizes
+    the Indic numerals in the input string to Roman script. Later, it uses the input string with normalized
+    Indic numerals to wrap the spans of text matching the pattern with placeholder tags.
+    
+    Args:
+        text (str): input string.
+        pattern (list): list of patterns to search for in the input string.
+    
+    Returns:
+        Tuple[str, dict]: a tuple containing the modified text and a dictionary mapping 
+            placeholders to their original values.
+    """
+    text = normalize_indic_numerals(text.strip("\n"))
+    text, placeholder_entity_map  = wrap_with_placeholders(text, patterns)
+    return text, placeholder_entity_map

IndicTrans2/inference/requirements.txt

@@ -0,0 +1,11 @@
+git+https://github.com/anoopkunchukuttan/indic_nlp_library
+git+https://github.com/pytorch/fairseq
+sacremoses
+pandas
+mock
+nltk
+sacrebleu
+urduhack[tf]
+mosestokenizer
+ctranslate2
+sentencepiece

IndicTrans2/inference/triton_server/Dockerfile

@@ -0,0 +1,25 @@
+ARG BASE_IMAGE=nvcr.io/nvidia/tritonserver:22.12-py3
+FROM ${BASE_IMAGE}
+
+# Ensure apt-get won't prompt for selecting options
+ENV DEBIAN_FRONTEND=noninteractive
+ENV PYTHONIOENCODING=utf8
+
+WORKDIR /home
+
+WORKDIR /home/indicTrans2
+COPY requirements.txt .
+RUN pip install -r requirements.txt
+
+COPY download.py .
+RUN python3 download.py
+
+COPY . ./inference
+
+WORKDIR /home/
+COPY ./triton_server/triton_repo ./triton_repo
+
+CMD ["tritonserver", "--model-repository=/home/triton_repo", "--log-verbose=2", "--strict-model-config=false", "--http-port=8000", "--grpc-port=8001", "--metrics-port=8002"]
+EXPOSE 8000
+EXPOSE 8001
+EXPOSE 8002

IndicTrans2/inference/triton_server/README.md

@@ -0,0 +1,22 @@
+# Triton server
+
+## Building the image
+
+```
+cd indicTrans2/inference/
+docker build -f triton_server/Dockerfile -t indictrans2_triton .
+```
+
+## Running the container
+
+Place the `en-indic` and `indic-en` checkpoint folders into `indicTrans2/checkpoints` directory
+
+Then start the server by:
+```
+docker run --shm-size=256m --gpus=1 --rm -v ${PWD}/../checkpoints/:/models/checkpoints -p 8000:8000 -t indictrans2_triton
+```
+
+## Sample client
+
+- Do `pip install tritonclient[all] gevent` first.
+- Then `python3 triton_server/client.py`

IndicTrans2/inference/triton_server/azure_ml/README.md

@@ -0,0 +1,56 @@
+# Deployment on Azure Machine Learning
+
+## Pre-requisites
+
+```
+cd inference/triton_server
+```
+
+Set the environment for AML:
+```
+export RESOURCE_GROUP=Dhruva-prod
+export WORKSPACE_NAME=dhruva--central-india
+export DOCKER_REGISTRY=dhruvaprod
+```
+
+Also remember to edit the `yml` files accordingly.
+
+## Registering the model
+
+```
+az ml model create --file azure_ml/model.yml --resource-group $RESOURCE_GROUP --workspace-name $WORKSPACE_NAME
+```
+
+## Pushing the docker image to Container Registry
+
+```
+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
+```
+
+## Creating the execution environment
+
+```
+az ml environment create -f azure_ml/environment.yml -g $RESOURCE_GROUP -w $WORKSPACE_NAME
+```
+
+## Publishing the endpoint for online inference
+
+```
+az ml online-endpoint create -f azure_ml/endpoint.yml -g $RESOURCE_GROUP -w $WORKSPACE_NAME
+```
+
+Now from the Azure Portal, open the Container Registry, and grant ACR_PULL permission for the above endpoint, so that it is allowed to download the docker image.
+
+## Attaching a deployment
+
+```
+az ml online-deployment create -f azure_ml/deployment.yml --all-traffic -g $RESOURCE_GROUP -w $WORKSPACE_NAME
+```
+
+## Testing if inference works
+
+1. From Azure ML Studio, go to the "Consume" tab, and get the endpoint domain (without `https://` or trailing `/`) and an authentication key.
+2. In `client.py`, enable `ENABLE_SSL = True`, and then set the `ENDPOINT_URL` variable as well as `Authorization` value inside `HTTP_HEADERS`.
+3. Run `python3 client.py`

IndicTrans2/inference/triton_server/azure_ml/deployment.yml

@@ -0,0 +1,13 @@
+$schema: https://azuremlschemas.azureedge.net/latest/managedOnlineDeployment.schema.json
+name: ai4b-indictransv2--t4-piv--gpu
+endpoint_name: ai4b-indictransv2--t4
+model: azureml:indictrans-v2--models:1
+model_mount_path: /models
+environment: azureml:triton-indictrans-v2-env:1
+instance_type: Standard_NC4as_T4_v3
+instance_count: 1
+request_settings:
+  request_timeout_ms: 90000
+  max_concurrent_requests_per_instance: 100
+  max_queue_wait_ms: 2000
+app_insights_enabled: true

IndicTrans2/inference/triton_server/azure_ml/endpoint.yml

@@ -0,0 +1,3 @@
+$schema: https://azuremlschemas.azureedge.net/latest/managedOnlineEndpoint.schema.json
+name: ai4b-indictransv2--t4
+auth_mode: key

IndicTrans2/inference/triton_server/azure_ml/environment.yml

@@ -0,0 +1,14 @@
+$schema: https://azuremlschemas.azureedge.net/latest/environment.schema.json
+name: triton-indictrans-v2-env
+image: dhruvaprod.azurecr.io/nmt/triton-indictrans-v2:latest
+version: 1
+inference_config:
+  liveness_route:
+    path: /v2/health/live
+    port: 8000
+  readiness_route:
+    path: /v2/health/ready
+    port: 8000
+  scoring_route:
+    path: /
+    port: 8000

IndicTrans2/inference/triton_server/azure_ml/model.yml

@@ -0,0 +1,5 @@
+$schema: https://azuremlschemas.azureedge.net/latest/model.schema.json
+name: indictrans-v2--models
+version: 1
+path: ../../../checkpoints
+type: triton_model

IndicTrans2/inference/triton_server/client.py

@@ -0,0 +1,55 @@
+import tritonclient.http as http_client
+from tritonclient.utils import *
+import numpy as np
+
+ENABLE_SSL = False
+ENDPOINT_URL = 'localhost:8000'
+HTTP_HEADERS = {"Authorization": "Bearer __PASTE_KEY_HERE__"}
+
+# Connect to the server
+if ENABLE_SSL:
+    import gevent.ssl
+    triton_http_client = http_client.InferenceServerClient(
+        url=ENDPOINT_URL, verbose=False,
+        ssl=True, ssl_context_factory=gevent.ssl._create_default_https_context,
+    )
+else:
+    triton_http_client = http_client.InferenceServerClient(
+        url=ENDPOINT_URL, verbose=False,
+    )
+
+print("Is server ready - {}".format(triton_http_client.is_server_ready(headers=HTTP_HEADERS)))
+
+def get_string_tensor(string_values, tensor_name):
+    string_obj = np.array(string_values, dtype="object")
+    input_obj = http_client.InferInput(tensor_name, string_obj.shape, np_to_triton_dtype(string_obj.dtype))
+    input_obj.set_data_from_numpy(string_obj)
+    return input_obj
+
+def get_translation_input_for_triton(texts: list, src_lang: str, tgt_lang: str):
+    return [
+        get_string_tensor([[text] for text in texts], "INPUT_TEXT"),
+        get_string_tensor([[src_lang]] * len(texts), "INPUT_LANGUAGE_ID"),
+        get_string_tensor([[tgt_lang]] * len(texts), "OUTPUT_LANGUAGE_ID"),
+    ]
+
+# Prepare input and output tensors
+input_sentences = ["Hello world, I am Ram and I am from Ayodhya.", "How are you Ravan bro?"]
+inputs = get_translation_input_for_triton(input_sentences, "en", "hi")
+output0 = http_client.InferRequestedOutput("OUTPUT_TEXT")
+
+# Send request
+response = triton_http_client.infer(
+    "nmt",
+    model_version='1',
+    inputs=inputs,
+    outputs=[output0],
+    headers=HTTP_HEADERS,
+)#.get_response()
+
+# Decode the response
+output_batch = response.as_numpy('OUTPUT_TEXT').tolist()
+for input_sentence, translation in zip(input_sentences, output_batch):
+    print()
+    print(input_sentence)
+    print(translation[0].decode("utf-8"))

IndicTrans2/inference/triton_server/triton_repo/nmt/1/model.py

@@ -0,0 +1,167 @@
+import os
+import sys
+import json
+import numpy as np
+import triton_python_backend_utils as pb_utils
+
+PWD = os.path.dirname(__file__)
+
+INFERENCE_MODULE_DIR = "/home/indicTrans2/"
+sys.path.insert(0, INFERENCE_MODULE_DIR)
+from inference.engine import Model, iso_to_flores
+INDIC_LANGUAGES = set(iso_to_flores)
+
+ALLOWED_DIRECTION_STRINGS = {"en-indic", "indic-en", "indic-indic"}
+FORCE_PIVOTING = False
+DEFAULT_PIVOT_LANG = "en"
+
+class TritonPythonModel:
+    def initialize(self, args):
+        self.model_config = json.loads(args['model_config'])
+        self.model_instance_device_id = json.loads(args['model_instance_device_id'])
+        self.output_name = "OUTPUT_TEXT"
+        self.output_dtype = pb_utils.triton_string_to_numpy(
+            pb_utils.get_output_config_by_name(self.model_config, self.output_name)["data_type"])
+        
+
+        # checkpoints_root_dir = os.path.join(PWD, "checkpoints")
+        checkpoints_root_dir = "/models/checkpoints"
+        checkpoint_folders = [ f.path for f in os.scandir(checkpoints_root_dir) if f.is_dir() ]
+        # The assumption is that, each folder name is `<src_direction>-to-<tgt_direction>`
+
+        if not checkpoint_folders:
+            raise RuntimeError(f"No checkpoint folders in: {checkpoints_root_dir}")
+
+        self.models = {}
+        for checkpoint_folder in checkpoint_folders:
+            direction_string = os.path.basename(checkpoint_folder)
+            assert direction_string in ALLOWED_DIRECTION_STRINGS, f"Checkpoint folder-name `{direction_string}` not allowed"
+            self.models[direction_string] = Model(os.path.join(checkpoint_folder, "ct2_fp16_model"), input_lang_code_format="iso", model_type="ctranslate2")
+            # self.models[direction_string] = Model(checkpoint_folder, input_lang_code_format="iso", model_type="fairseq")
+        
+        self.pivot_lang = None
+        if "en-indic" in self.models and "indic-en" in self.models:
+            if  "indic-indic" not in self.models:
+                self.pivot_lang = DEFAULT_PIVOT_LANG
+            elif FORCE_PIVOTING:
+                del self.models["indic-indic"]
+                self.pivot_lang = DEFAULT_PIVOT_LANG
+    
+    def get_direction_string(self, input_language_id, output_language_id):
+        direction_string = None
+        if input_language_id == DEFAULT_PIVOT_LANG and output_language_id in INDIC_LANGUAGES:
+            direction_string = "en-indic"
+        elif input_language_id in INDIC_LANGUAGES:
+            if output_language_id == DEFAULT_PIVOT_LANG:
+                direction_string = "indic-en"
+            elif output_language_id in INDIC_LANGUAGES:
+                direction_string = "indic-indic"
+        return direction_string
+
+    def get_model(self, input_language_id, output_language_id):
+        direction_string = self.get_direction_string(input_language_id, output_language_id)
+        
+        if direction_string in self.models:
+            return self.models[direction_string]
+        raise RuntimeError(f"Language-pair not supported: {input_language_id}-{output_language_id}")
+
+    def execute(self,requests):
+        # print("REQ_COUNT", len(requests))
+        modelwise_batches = {}
+        responses = []
+        for request_id, request in enumerate(requests):
+            input_text_batch = pb_utils.get_input_tensor_by_name(request, "INPUT_TEXT").as_numpy()
+            input_language_id_batch = pb_utils.get_input_tensor_by_name(request, "INPUT_LANGUAGE_ID").as_numpy()
+            output_language_id_batch = pb_utils.get_input_tensor_by_name(request, "OUTPUT_LANGUAGE_ID").as_numpy()
+            
+            input_text_batch = [input_text[0].decode("utf-8", "ignore") for input_text in input_text_batch]
+            input_language_id_batch = [input_language_id[0].decode("utf-8", "ignore") for input_language_id in input_language_id_batch]
+            output_language_id_batch = [output_language_id[0].decode("utf-8", "ignore") for output_language_id in output_language_id_batch]
+
+            responses.append([['']] * len(input_text_batch))
+
+            for input_id, (input_text, input_language_id, output_language_id) in enumerate(zip(input_text_batch, input_language_id_batch, output_language_id_batch)):
+                direction_string = self.get_direction_string(input_language_id, output_language_id)
+                if direction_string not in self.models:
+                    if direction_string == "indic-indic" and self.pivot_lang:
+                        pass
+                    else:
+                        raise RuntimeError(f"Language-pair not supported: {input_language_id}-{output_language_id}")
+                
+                if direction_string not in modelwise_batches:
+                    modelwise_batches[direction_string] = {
+                        "payloads": [],
+                        "text_id_to_req_id_input_id": [],
+                    }
+                
+                modelwise_batches[direction_string]["payloads"].append([input_text, input_language_id, output_language_id])
+                modelwise_batches[direction_string]["text_id_to_req_id_input_id"].append((request_id, input_id))
+
+        for direction_string, batch in modelwise_batches.items():
+            if direction_string == "indic-indic" and self.pivot_lang:
+                model = self.get_model("hi", self.pivot_lang)
+                original_langs = []
+                for i in range(len(batch["payloads"])):
+                    original_langs.append(batch["payloads"][i][2])
+                    batch["payloads"][i][2] = self.pivot_lang
+
+                pivot_texts = model.paragraphs_batch_translate__multilingual(batch["payloads"])
+
+                for i in range(len(batch["payloads"])):
+                    batch["payloads"][i][0] = pivot_texts[i]
+                    batch["payloads"][i][1] = self.pivot_lang
+                    batch["payloads"][i][2] = original_langs[i]
+                
+                model = self.get_model(self.pivot_lang, "hi")
+                translations = model.paragraphs_batch_translate__multilingual(batch["payloads"])
+            else:
+                model = self.models[direction_string]
+                translations = model.paragraphs_batch_translate__multilingual(batch["payloads"])
+                # translations = ["bro"] * len(batch["payloads"])
+            
+            for translation, (request_id, output_id) in zip(translations, batch["text_id_to_req_id_input_id"]):
+                responses[request_id][output_id] = [translation]
+        
+        for i in range(len(responses)):
+            responses[i] = pb_utils.InferenceResponse(output_tensors=[
+                pb_utils.Tensor(
+                    self.output_name,
+                    np.array(responses[i], dtype=self.output_dtype),
+                )
+            ])
+        return responses
+    
+    def execute_sequential(self,requests):
+        # print("REQ_COUNT", len(requests))
+        responses = []
+        for request in requests:
+            input_text_batch = pb_utils.get_input_tensor_by_name(request, "INPUT_TEXT").as_numpy()
+            input_language_id_batch = pb_utils.get_input_tensor_by_name(request, "INPUT_LANGUAGE_ID").as_numpy()
+            output_language_id_batch = pb_utils.get_input_tensor_by_name(request, "OUTPUT_LANGUAGE_ID").as_numpy()
+            
+            input_text_batch = [input_text[0].decode("utf-8", "ignore") for input_text in input_text_batch]
+            input_language_id_batch = [input_language_id[0].decode("utf-8", "ignore") for input_language_id in input_language_id_batch]
+            output_language_id_batch = [output_language_id[0].decode("utf-8", "ignore") for output_language_id in output_language_id_batch]
+
+            generated_outputs = []
+
+            for input_text, input_language_id, output_language_id in zip(input_text_batch, input_language_id_batch, output_language_id_batch):
+                if self.pivot_lang and (input_language_id != self.pivot_lang and output_language_id != self.pivot_lang):
+                    model = self.get_model(input_language_id, self.pivot_lang)
+                    pivot_text = model.translate_paragraph(input_text, input_language_id, self.pivot_lang)
+                    
+                    model = self.get_model(self.pivot_lang, output_language_id)
+                    translation = model.translate_paragraph(pivot_text, self.pivot_lang, output_language_id)
+                else:
+                    model = self.get_model(input_language_id, output_language_id)
+                    translation = model.translate_paragraph(input_text, input_language_id, output_language_id)
+                generated_outputs.append([translation])
+
+            inference_response = pb_utils.InferenceResponse(output_tensors=[
+                pb_utils.Tensor(
+                    self.output_name,
+                    np.array(generated_outputs, dtype=self.output_dtype),
+                )
+            ])
+            responses.append(inference_response)
+        return responses

IndicTrans2/inference/triton_server/triton_repo/nmt/config.pbtxt

@@ -0,0 +1,32 @@
+backend: "python"
+max_batch_size: 512
+input [{
+  name: "INPUT_TEXT"
+  data_type: TYPE_STRING
+  dims: 1
+},
+{
+  name: "INPUT_LANGUAGE_ID"
+  data_type: TYPE_STRING
+  dims: 1
+},
+{
+  name: "OUTPUT_LANGUAGE_ID"
+  data_type: TYPE_STRING
+  dims: 1
+}]
+  
+output {
+  name: "OUTPUT_TEXT"
+  data_type: TYPE_STRING
+  dims: 1
+}
+
+dynamic_batching { 
+  
+}
+
+instance_group [{
+ count: 1
+ kind: KIND_GPU
+}]

IndicTrans2/inference/utils.map_token_lang.tsv

@@ -0,0 +1,26 @@
+asm_Beng	hi
+ben_Beng	hi
+brx_Deva	hi
+doi_Deva	hi
+gom_Deva	hi
+eng_Latn	en
+guj_Gujr	hi
+hin_Deva	hi
+kan_Knda	hi
+kas_Arab	ar
+kas_Deva	hi
+mai_Deva	hi
+mar_Deva	hi
+mal_Mlym	hi
+mni_Beng	hi
+mni_Mtei	en
+npi_Deva	hi
+ory_Orya	hi
+pan_Guru	hi
+san_Deva	hi
+sat_Olck	hi
+snd_Arab	ar
+snd_Deva	hi
+tam_Taml	hi
+tel_Telu	hi
+urd_Arab	ar