Update README.md

README.md

@@ -13,7 +13,7 @@ widget:
   example_title: "Felipe Casali Silva, Teste, Rio de Janeiro, RJ"
 ---
 
-# lgpd_pii_identifier : Financial BERT PT BR
+# lgpd_pii_identifier : LGPD PII Identifier
 
 lgpd_pii_identifier is a pre-trained NLP model to identify sensitive data in the scope of LGPD (Lei Geral de Proteção de Dados)
 
@@ -32,7 +32,7 @@ data according to their businness needs, and governance rules.
 In order to use the model, you need to get the HuggingFace auth token. You can get it [here](https://huggingface.co/settings/token).
 
 ```python
-from transformers import AutoTokenizer, BertForSequenceClassification
+from transformers import DistilBertModel, DistilBertTokenizer
 import numpy as np
   
 pred_mapper = {
@@ -42,8 +42,8 @@ pred_mapper = {
     3: "estado"
   }
 
-tokenizer = AutoTokenizer.from_pretrained("FelipeCasali-USP/lgpd_pii_identifier")
-lgpd_pii_identifier = BertForSequenceClassification.from_pretrained("FelipeCasali-USP/lgpd_pii_identifier")
+tokenizer = DistilBertTokenizer.from_pretrained("FelipeCasali-USP/lgpd_pii_identifier")
+lgpd_pii_identifier = DistilBertModel.from_pretrained("FelipeCasali-USP/lgpd_pii_identifier")
 
 tokens = tokenizer(["String to be analized"], return_tensors="pt",
                     padding=True, truncation=True, max_length=512)