odaria/rubert_news_classifier

RuBERT news classifier

This model is trained to classify news articles by category (those include Business, Culture, World, Society, Politics, Consumer market, Incidents, Sports, Telecommunications, Finance, Economics - 10 categories in total).

Details

Description

• Language: Russian
• License: MIT
• Finetuned from model: https://huggingface.co/cointegrated/rubert-tiny2

Model Sources

The model was fine-tuned for a news classification project that can be found in the following repository. The project also includes SVM and Logistic Regression models trained for the same classification task on the same dataset.

Direct Use

import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

tokenizer = AutoTokenizer.from_pretrained("your-username/news-classifier-models")
model = AutoModelForSequenceClassification.from_pretrained("your-username/news-classifier-models")
model.eval()

CATEGORIES = ['Бизнес', 'Культура', 'Мир', 'Общество', 'Политика', 'Потребительский рынок', 'Происшествия', 'Спорт', 'Телекоммуникации', 'Экономика']

def predict(text: str) -> dict:
    inputs = tokenizer(
        text,
        max_length=512,
        truncation=True,
        padding="max_length",
        return_tensors="pt",
    )
    with torch.no_grad():
        logits = model(**inputs).logits

    probs    = torch.softmax(logits, dim=-1).squeeze()
    pred_idx = probs.argmax().item()

    return {
        "category":   CATEGORIES[pred_idx],
        "confidence": round(probs[pred_idx].item(), 3),
        "all_probs":  {cat: round(p.item(), 3) for cat, p in zip(CATEGORIES, probs)},
    }

result = predict("Центральный банк повысил ключевую ставку до 21 процента.")
print(result)

Training Data

The model is trained only on Russian news dataset containing 10000+ news articles from Kommersant newspaper (https://www.kommersant.ru/). Each news category included around 1000 articles each. Due to copyright risks the training dataset can not be published.

Metrics

• Accuracy = 0.811
• F1 weighted = 0.807
• Precision weighted = 0.809
• Recall weighted = 0.811