cointegrated/rubert-tiny-sentiment-balanced

La siguiente función estima el sentimiento del texto dado:

import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

model_checkpoint = 'cointegrated/rubert-tiny-sentiment-balanced'
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
model = AutoModelForSequenceClassification.from_pretrained(model_checkpoint)
if torch.cuda.is_available():
    model.cuda()

def get_sentiment(text, return_type='label'):
    """ Calcular el sentimiento de un texto. `return_type` puede ser 'label', 'score' o 'proba' """
    with torch.no_grad():
        inputs = tokenizer(text, return_tensors='pt', truncation=True, padding=True).to(model.device)
        proba = torch.sigmoid(model(**inputs).logits).cpu().numpy()[0]
    if return_type == 'label':
        return model.config.id2label[proba.argmax()]
    elif return_type == 'score':
        return proba.dot([-1, 0, 1])
    return proba

text = 'Какая гадость эта ваша заливная рыба!'
# clasificar el texto
print(get_sentiment(text, 'label')) # negativo
# puntuar el texto en la escala de -1 (muy negativo) a +1 (muy positivo)
print(get_sentiment(text, 'score')) # -0.5894946306943893
# calcular probabilidades de todas las etiquetas
print(get_sentiment(text, 'proba')) # [0.7870447  0.4947824  0.19755007]