Update README.md

README.md

@@ -61,13 +61,21 @@ Please refer to the [github repository](https://github.com/AI4Bharat/IndicTrans2
 ```python
 import torch
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
-from IndicTransTokenizer import IndicProcessor
-
+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 = "eng_Latn", "hin_Deva"
 model_name = "ai4bharat/indictrans2-en-indic-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)
 
@@ -78,12 +86,8 @@ input_sentences = [
     "My friend has invited me to his birthday party, and I will give him a gift.",
 ]
 
-src_lang, tgt_lang = "eng_Latn", "hin_Deva"
-
 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,
@@ -120,8 +124,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