Update README.md

README.md

@@ -30,7 +30,7 @@ This model is initialized with BERTweet and trained with an MLM objective.
 This pre-trained language model **can be fine-tunned to any downstream task (e.g. classification)**.
 
 ```python
-from transformers import BertTokenizer, BertForMaskedLM, pipeline
+from transformers import AutoModel, AutoTokenizer, pipeline
 import torch
 
 # choose GPU if available
@@ -40,8 +40,8 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 pretrained_LM_path = "kornosk/polibertweet-mlm"
 
 # load model
-tokenizer = BertTokenizer.from_pretrained(pretrained_LM_path)
-model = BertForMaskedLM.from_pretrained(pretrained_LM_path)
+tokenizer = AutoTokenizer.from_pretrained(pretrained_LM_path)
+model = AutoModel.from_pretrained(pretrained_LM_path)
 
 # fill mask
 example = "Trump is the <mask> of USA"