RuBERT Emotion Classifier (QLoRA)

Multi-label emotion classifier for Russian text based on Izard's 10 basic emotions, fine-tuned with QLoRA on RuIzardEmotions.

Base model: ai-forever/ruBert-large
Method: QLoRA (4-bit NF4 + LoRA r=8)
Labels: joy, sadness, anger, enthusiasm, surprise, disgust, fear, guilt, shame, neutral

Metrics (validation set, after per-class threshold tuning)

Per-class breakdown

Shame is the hardest class due to low support (153 samples) and high overlap with guilt.

Usage

import json
import torch
import torch.nn as nn
from transformers import AutoTokenizer, AutoModel, BitsAndBytesConfig
from transformers.modeling_outputs import SequenceClassifierOutput
from peft import PeftModel

class BertWithClassifier(nn.Module):
    def __init__(self, encoder, hidden_size, num_labels):
        super().__init__()
        self.encoder = encoder
        self.dropout = nn.Dropout(0.1)
        self.classifier = nn.Linear(hidden_size, num_labels)

    def forward(self, input_ids=None, attention_mask=None,
                token_type_ids=None, **kwargs):
        out = self.encoder(
            input_ids=input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
        )
        pooled = self.dropout(out.last_hidden_state[:, 0, :].float())
        return SequenceClassifierOutput(logits=self.classifier(pooled))


REPO = "ilyali034/rubert-emotion-ru-large"

with open("emotion_config.json") as f:
    cfg = json.load(f)

tokenizer = AutoTokenizer.from_pretrained("ai-forever/ruBert-large")

base = AutoModel.from_pretrained(
    "ai-forever/ruBert-large",
    quantization_config=BitsAndBytesConfig(load_in_4bit=True),
    device_map="auto",
)
base = PeftModel.from_pretrained(base, REPO + "/lora_adapter")

model = BertWithClassifier(base, base.config.hidden_size, len(cfg["labels"]))
model.classifier.load_state_dict(torch.load("classifier.pt", map_location="cpu"))
model.eval()

def predict(text: str) -> dict:
    inputs = tokenizer(
        text,
        return_tensors="pt",
        truncation=True,
        max_length=128,
        padding=True,
    ).to("cuda")
    with torch.no_grad():
        probs = torch.sigmoid(model(**inputs).logits).cpu().numpy()[0]
    thresholds = list(cfg["thresholds"].values())
    return {
        lbl: round(float(p), 4)
        for lbl, p, thr in zip(cfg["labels"], probs, thresholds)
        if p > thr
    }

print(predict("Я очень рад этой новости!"))
# {'joy': 0.8231, 'enthusiasm': 0.6714}

print(predict("Мне стыдно за своё поведение, я чувствую себя виноватым"))
# {'guilt': 0.7102, 'shame': 0.5891}

Training configuration

Optimal thresholds

Files