lkurakht/pt_br_classifier

1import torch2 3import json4import typing as tp5 6import torch.nn.functional as F7from torch import Tensor8 9from datasets import ClassLabel10import transformers11from transformers import BertForSequenceClassification12from transformers import BertForSequenceClassification, AutoTokenizer13import numpy as np14 15tokenizer = AutoTokenizer.from_pretrained('adalbertojunior/distilbert-portuguese-cased', do_lower_case=False)16 17classes = ['pt','pt_br']18class_label = ClassLabel(names=classes)19 20def get_model():21    return BertForSequenceClassification.from_pretrained(22   './pt_br_model',23    num_labels = 2,24    output_attentions = False,25    output_hidden_states = False,26    )27 28model = get_model()29text = 'hello'30input_tensor = tokenizer(text, padding=True, truncation=True, max_length=256, add_special_tokens=True, return_tensors="pt")31 32logits=model(**input_tensor).logits33probabilities = F.softmax(logits, dim=1).flatten().tolist()34maxidx = np.argmax(probabilities)35print(classes[maxidx], probabilities[maxidx])