from transformers import BertTokenizer, AutoTokenizer, BertForTokenClassification, pipeline
import torch

tokenizer = BertTokenizer.from_pretrained('neuralmind/bert-base-portuguese-cased')
model = BertForTokenClassification.from_pretrained('neuralmind/bert-base-portuguese-cased')

inputs = tokenizer("Ivo viu a uva", return_tensors="pt")
labels = torch.tensor([1] * inputs["input_ids"].size(1)).unsqueeze(0)  # Batch size 1

print(inputs)

outputs = model(**inputs, labels=labels)
loss, scores = outputs[:2]

print ("\n\nNER")

MODEL_NAME = "celiudos/legal-bert-lgpd"
#MODEL_NAME = "neuralmind/bert-base-portuguese-cased"
tokenizer = AutoTokenizer.from_pretrained(
    MODEL_NAME,
    model_max_length=512,
)

pipe = pipeline(
    "ner",
    tokenizer=tokenizer,
    model=MODEL_NAME,
    stride=100,
    aggregation_strategy="first",
    device=0 if torch.cuda.is_available() else -1,
)

print(
pipe(
    "Anotação de Responsabilidade Técnica Nº 1055330634101 de 12 de janeiro de 2013 relativa à Lei Federal Nº 531. Trata-se de representação referente a possível falsificação documentação técnica registrada pelo CREA-SP, feita pelo senhor Francis Pantele da Cozzi, CPF: 412.612.341-32, telefone (31) 951358433, email fran@bol.com, atinente à sua contratação pela senhora Marinalva Bete Raz, CPF: 049.567.041-22, telefone (61) 9412 3333, mulher branca, opinião política conservadora, religião evangélica. Marinalva Bete Raz reclama por indenização por danos morais no dia 14.05.2013 no valor de R$ 82.662,00 (Oitenta e dois mil, seiscentos e sessenta e dois reais) relacionado ao endereço IP 192.168.01 e ao endereço constante no CEP 59123-222, Rua dos Pioneiros, nº 450, Jardim Esmeralda, Campo Grande, MS."
))

print(
pipe(
    "Reunião na casa de Braga Netto discutiu ações para gerar caos social, afirmou Mauro Cid"
))