Update README.md

README.md

@@ -57,17 +57,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"
 
-
-model_name = "ai4bharat/indictrans2-indic-indic-dist-320M"
+src_lang, tgt_lang = "hin_Deva", "tam_Taml"
+model_name = "ai4bharat/indictrans2-indic-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,8 +83,6 @@ input_sentences = [
     "मेरे मित्र ने मुझे उसके जन्मदिन की पार्टी में बुलाया है, और मैं उसे एक तोहफा दूंगा।",
 ]
 
-src_lang, tgt_lang = "hin_Deva", "tam_Taml"
-
 batch = ip.preprocess_batch(
     input_sentences,
     src_lang=src_lang,
@@ -124,9 +127,6 @@ 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.**
-
-
 ### Citation
 
 If you consider using our work then please cite using: