Multi-label emotion classifier for Russian text based on Izard's 10 basic emotions,
fine-tuned with QLoRA on
RuIzardEmotions.
1import json
2import torch
3import torch.nn as nn
4from transformers import AutoTokenizer, AutoModel, BitsAndBytesConfig
5from transformers.modeling_outputs import SequenceClassifierOutput
6from peft import PeftModel
7
8class BertWithClassifier(nn.Module):
9 def __init__(self, encoder, hidden_size, num_labels):
10 super().__init__()
11 self.encoder = encoder
12 self.dropout = nn.Dropout(0.1)
13 self.classifier = nn.Linear(hidden_size, num_labels)
14
15 def forward(self, input_ids=None, attention_mask=None,
16 token_type_ids=None, **kwargs):
17 out = self.encoder(
18 input_ids=input_ids,
19 attention_mask=attention_mask,
20 token_type_ids=token_type_ids,
21 )
22 pooled = self.dropout(out.last_hidden_state[:, 0, :].float())
23 return SequenceClassifierOutput(logits=self.classifier(pooled))
24
25
26REPO = "ilyali034/rubert-emotion-ru-large"
27
28with open("emotion_config.json") as f:
29 cfg = json.load(f)
30
31tokenizer = AutoTokenizer.from_pretrained("ai-forever/ruBert-large")
32
33base = AutoModel.from_pretrained(
34 "ai-forever/ruBert-large",
35 quantization_config=BitsAndBytesConfig(load_in_4bit=True),
36 device_map="auto",
37)
38base = PeftModel.from_pretrained(base, REPO + "/lora_adapter")
39
40model = BertWithClassifier(base, base.config.hidden_size, len(cfg["labels"]))
41model.classifier.load_state_dict(torch.load("classifier.pt", map_location="cpu"))
42model.eval()
43
44def predict(text: str) -> dict:
45 inputs = tokenizer(
46 text,
47 return_tensors="pt",
48 truncation=True,
49 max_length=128,
50 padding=True,
51 ).to("cuda")
52 with torch.no_grad():
53 probs = torch.sigmoid(model(**inputs).logits).cpu().numpy()[0]
54 thresholds = list(cfg["thresholds"].values())
55 return {
56 lbl: round(float(p), 4)
57 for lbl, p, thr in zip(cfg["labels"], probs, thresholds)
58 if p > thr
59 }
60
61print(predict("Я очень рад этой новости!"))
62# {'joy': 0.8231, 'enthusiasm': 0.6714}
63
64print(predict("Мне стыдно за своё поведение, я чувствую себя виноватым"))
65# {'guilt': 0.7102, 'shame': 0.5891}