Views
No views yet
Appendix D: Model Card of the preprint for further details on model training, intended use, data, metrics, limitations and recommendations.1import torch
2from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
3from IndicTransToolkit.processor import IndicProcessor
4# recommended to run this on a gpu with flash_attn installed
5# don't set attn_implemetation if you don't have flash_attn
6DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
7
8src_lang, tgt_lang = "hin_Deva", "eng_Latn"
9model_name = "ai4bharat/indictrans2-indic-en-1B"
10tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
11
12model = AutoModelForSeq2SeqLM.from_pretrained(
13 model_name,
14 trust_remote_code=True,
15 torch_dtype=torch.float16, # performance might slightly vary for bfloat16
16 attn_implementation="flash_attention_2"
17).to(DEVICE)
18
19ip = IndicProcessor(inference=True)
20
21input_sentences = [
22 "जब मैं छोटा था, मैं हर रोज़ पार्क जाता था।",
23 "हमने पिछले सप्ताह एक नई फिल्म देखी जो कि बहुत प्रेरणादायक थी।",
24 "अगर तुम मुझे उस समय पास मिलते, तो हम बाहर खाना खाने चलते।",
25 "मेरे मित्र ने मुझे उसके जन्मदिन की पार्टी में बुलाया है, और मैं उसे एक तोहफा दूंगा।",
26]
27
28batch = ip.preprocess_batch(
29 input_sentences,
30 src_lang=src_lang,
31 tgt_lang=tgt_lang,
32)
33
34# Tokenize the sentences and generate input encodings
35inputs = tokenizer(
36 batch,
37 truncation=True,
38 padding="longest",
39 return_tensors="pt",
40 return_attention_mask=True,
41).to(DEVICE)
42
43# Generate translations using the model
44with torch.no_grad():
45 generated_tokens = model.generate(
46 **inputs,
47 use_cache=True,
48 min_length=0,
49 max_length=256,
50 num_beams=5,
51 num_return_sequences=1,
52 )
53
54# Decode the generated tokens into text
55generated_tokens = tokenizer.batch_decode(
56 generated_tokens,
57 skip_special_tokens=True,
58 clean_up_tokenization_spaces=True,
59)
60
61# Postprocess the translations, including entity replacement
62translations = ip.postprocess_batch(generated_tokens, lang=tgt_lang)
63
64for input_sentence, translation in zip(input_sentences, translations):
65 print(f"{src_lang}: {input_sentence}")
66 print(f"{tgt_lang}: {translation}")model_name parameter. Please read the model card of the RoPE-IT2 models for more information about the generation.flash_attention_2 for efficient generation.@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={}
}