lkurakht/pt_br_classifier

1import streamlit as st2import torch3import numpy as np4import json5import typing as tp6 7import torch.nn.functional as F8from torch import Tensor9 10from datasets import ClassLabel11import transformers12from transformers import BertForSequenceClassification13from transformers import BertForSequenceClassification, AutoTokenizer14 15st.markdown("## Portuguese European and Brazilian dialect classifier")16st.markdown("[You can see the difference between dialects here](https://en.wikipedia.org/wiki/Portuguese_language#Writing_system)")17text = st.text_input('## Text:')18 19tokenizer = AutoTokenizer.from_pretrained('adalbertojunior/distilbert-portuguese-cased', do_lower_case=False)20 21classes = ['pt', 'pt_br']22 23class_label = ClassLabel(names=classes)24 25 26@st.cache27def get_model():28    return BertForSequenceClassification.from_pretrained(29   './pt_br_model',30    num_labels = 2,31    output_attentions = False,32    output_hidden_states = False,33    )34 35 36model = get_model()37 38 39@torch.inference_mode()40def print_results():41    input_tensor = tokenizer(text, padding=True, truncation=True, max_length=256, add_special_tokens=True, return_tensors="pt")42    logits = model(**input_tensor).logits43    probabilities = F.softmax(logits, dim=1).flatten().tolist()44    maxidx = np.argmax(probabilities)45    results = f"### {classes[maxidx]} score: {probabilities[maxidx]*100}%"46    st.markdown('## Results:')47    st.markdown(results)48 49if text:50    print_results()51