Djacon's picture
Update README.md
1161779
|
raw
history blame
2.64 kB
metadata
license: mit
language:
  - ru
tags:
  - russian
  - classification
  - emotion
  - emotion-detection
  - emotion-recognition
  - multiclass
widget:
  - text: Как дела?
  - text: Дурак твой дед
  - text: Только попробуй!!!
  - text: Не хочу в школу(
  - text: Сейчас ровно час дня
  - text: >-
      А ты уверен, что эти полоски снизу не врут? Точно уверен? Вот прям 100
      процентов?
datasets:
  - Djacon/ru_go_emotions

First - you should prepare few functions to talk to model

import torch
from transformers import BertForSequenceClassification, AutoTokenizer

LABELS_RU = ['нейтрально', 'радость', 'грусть', 'гнев', 'интерес', 'удивление', 'отвращение', 'страх', 'вина', 'стыд']
tokenizer = AutoTokenizer.from_pretrained('Djacon/rubert-tiny2-russian-emotion-detection')
model = BertForSequenceClassification.from_pretrained('Djacon/rubert-tiny2-russian-emotion-detection')

# Predicting emotion in text
@torch.no_grad()
def predict_emotion(text: str) -> str:
    inputs = tokenizer(text, truncation=True, return_tensors='pt')
    inputs = inputs.to(model.device)
    
    outputs = model(**inputs)
    
    pred = torch.nn.functional.sigmoid(outputs.logits)
    pred = pred.argmax(dim=1)
    
    return LABELS_RU[pred[0]]


# Probabilistic prediction of emotion in a text
@torch.no_grad()
def predict_emotions(text: str) -> dict:
    inputs = tokenizer(text, truncation=True, return_tensors='pt')
    inputs = inputs.to(model.device)

    outputs = model(**inputs)

    pred = torch.nn.functional.sigmoid(outputs.logits)

    emotions_list = {}
    for i in range(len(pred[0].tolist())):
        emotions_list[LABELS_RU[i]] = round(pred[0].tolist()[i], 4)
    return emotions_list

And then - just gently ask a model to predict your emotion

simple_prediction = predict_emotion("Какой же сегодня прекрасный день, братья")
not_simple_prediction = predict_emotions("Какой же сегодня прекрасный день, братья")

print(simple_prediction)
print(not_simple_prediction)
# радость
# {'нейтрально': 0.1985, 'радость': 0.7419, 'грусть': 0.0261, 'гнев': 0.0295, 'интерес': 0.1983, 'удивление': 0.4305, 'отвращение': 0.0082, 'страх': 0.008, 'вина': 0.0046, 'стыд': 0.0097}

Citations

@misc{Djacon,
  author = {Djacon},
  year = {2023},
  publisher = {Hugging Face},
  journal = {Hugging Face Hub},
}