Views
No views yet
1from transformers import DistilBertModel, DistilBertTokenizer
2import numpy as np
3
4pred_mapper = {
5 0: "cnpj",
6 1: "cpf",
7 2: "nome",
8 3: "estado"
9 }
10
11tokenizer = DistilBertTokenizer.from_pretrained("FelipeCasali-USP/lgpd_pii_identifier")
12lgpd_pii_identifier = DistilBertModel.from_pretrained("FelipeCasali-USP/lgpd_pii_identifier")
13
14tokens = tokenizer(["String to be analized"], return_tensors="pt",
15 padding=True, truncation=True, max_length=512)
16lgpd_pii_identifier_outputs = lgpd_pii_identifier(**tokens)
17
18preds = [pred_mapper[np.argmax(pred)] for pred in lgpd_pii_identifier_outputs.logits.cpu().detach().numpy()]