Update README.md

README.md

@@ -64,17 +64,22 @@ Please refer to the [github repository](https://github.com/AI4Bharat/IndicTrans2
 
 ```python
 import torch
-from transformers import (
-    AutoModelForSeq2SeqLM,
-    AutoTokenizer,
-)
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 from IndicTransToolkit import IndicProcessor
+# recommended to run this on a gpu with flash_attn installed
+# don't set attn_implemetation if you don't have flash_attn
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
-
+src_lang, tgt_lang = "hin_Deva", "eng_Latn"
 model_name = "ai4bharat/indictrans2-indic-en-1B"
 tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
 
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name, trust_remote_code=True)
+model = AutoModelForSeq2SeqLM.from_pretrained(
+    model_name, 
+    trust_remote_code=True, 
+    torch_dtype=torch.float16, # performance might slightly vary for bfloat16
+    attn_implementation="flash_attention_2"
+).to(DEVICE)
 
 ip = IndicProcessor(inference=True)
 
@@ -85,16 +90,12 @@ input_sentences = [
     "मेरे मित्र ने मुझे उसके जन्मदिन की पार्टी में बुलाया है, और मैं उसे एक तोहफा दूंगा।",
 ]
 
-src_lang, tgt_lang = "hin_Deva", "eng_Latn"
-
 batch = ip.preprocess_batch(
     input_sentences,
     src_lang=src_lang,
     tgt_lang=tgt_lang,
 )
 
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-
 # Tokenize the sentences and generate input encodings
 inputs = tokenizer(
     batch,
@@ -131,7 +132,11 @@ for input_sentence, translation in zip(input_sentences, translations):
     print(f"{tgt_lang}: {translation}")
 ```
 
+### 📢 Long Context IT2 Models
 
+- New RoPE based IndicTrans2 models which are capable of handling sequence lengths **upto 2048 tokens** are available [here](https://huggingface.co/collections/prajdabre/indictrans2-rope-6742ddac669a05db0804db35).
+- These models can be used by just changing the `model_name` parameter. Please read the model card of the RoPE-IT2 models for more information about the generation.
+- It is recommended to run these models with `flash_attention_2` for efficient generation. 
 
 ### Citation