Update README.md (#4)

- Update README.md (baeeca6f7732e3ea0c5105856f71ed3702d9fc1a)


Co-authored-by: Varun Gumma <[email protected]>

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"
 
-
+src_lang, tgt_lang = "hin_Deva", "tam_Taml"
 model_name = "ai4bharat/indictrans2-indic-indic-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)
 
@@ -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,10 +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: