Update README.md (#2)

- Update README.md (1a225950ddfd74cea6e10bd23e6dcdfac2946a9b)


Co-authored-by: Varun Gumma <VarunGumma@users.noreply.huggingface.co>

README.md

@@ -61,17 +61,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-dist-200M"
 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)
 
@@ -82,8 +87,6 @@ input_sentences = [
     "मेरे मित्र ने मुझे उसके जन्मदिन की पार्टी में बुलाया है, और मैं उसे एक तोहफा दूंगा।",
 ]
 
-src_lang, tgt_lang = "hin_Deva", "eng_Latn"
-
 batch = ip.preprocess_batch(
     input_sentences,
     src_lang=src_lang,
@@ -128,9 +131,10 @@ for input_sentence, translation in zip(input_sentences, translations):
     print(f"{tgt_lang}: {translation}")
 ```
 
-**Note: IndicTrans2 is now compatible with AutoTokenizer, however you need to use IndicProcessor from [IndicTransTokenizer](https://github.com/VarunGumma/IndicTransTokenizer) for preprocessing before tokenization.**
-
-
+### 📢 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