News topic classifier (e5-base)

Fine-tuned intfloat/multilingual-e5-base for multi-label classification of news articles into 13 topics: Business And Finance, Market And Exchange, Demographics, Scientific Exploration, Labour Dispute, Disaster, Conflicts And War, Politics, Crime, Theft, Assault, Corporate Crime, Fraud. No label means "Other".

Usage

  • Input: "passage: " + title + "\n" + content, truncated to 512 tokens.
  • Output: sigmoid per category, compared with the per-category threshold in thresholds.json (not 0.5).
  • If Theft or Assault is predicted, add Crime as well.
import json

import torch
from huggingface_hub import hf_hub_download
from transformers import AutoModelForSequenceClassification, AutoTokenizer

REPO = "maxafanasiev/news-topic-classifier-e5-base"  # private: run `hf auth login` first
PREFIX = "passage: "
PARENTS = {"Theft": "Crime", "Assault": "Crime"}

tokenizer = AutoTokenizer.from_pretrained(REPO)
model = AutoModelForSequenceClassification.from_pretrained(REPO).eval()
thresholds = json.load(open(hf_hub_download(REPO, "thresholds.json")))
categories = list(thresholds)  # same order as the model outputs


@torch.no_grad()
def classify(articles: list[dict], max_len: int = 512) -> list[list[str]]:
    texts = [f"{PREFIX}{a['title']}\n{a['content']}" for a in articles]
    batch = tokenizer(texts, truncation=True, max_length=max_len, padding=True, return_tensors="pt")
    probs = torch.sigmoid(model(**batch).logits)
    results = []
    for row in probs:
        labels = [c for c, p in zip(categories, row.tolist()) if p >= thresholds[c]]
        labels += [PARENTS[c] for c in labels if c in PARENTS and PARENTS[c] not in labels]
        results.append(labels or ["Other"])
    return results


print(classify([
    {"title": "Police arrest two men over jewellery store robbery", "content": "Officers detained the suspects after ..."},
    {"title": "El banco central sube los tipos de interés", "content": "La decisión responde a la inflación ..."},
]))
# e.g. [['Crime', 'Theft'], ['Business And Finance']]

Classify in batches of 16–32 articles; max_len=256 is about twice as fast at a small quality cost (see Results).

Download once, use without Hugging Face

  1. Download the files (needs access to this private repo once), with any of:
    • CLI: hf auth login, then hf download maxafanasiev/news-topic-classifier-e5-base --local-dir ./news-topic-classifier-e5-base
    • git: git lfs install && git clone https://huggingface.co/maxafanasiev/news-topic-classifier-e5-base (use a token as the password)
    • browser: the "Files and versions" tab, download every file into one folder
    • one archive: download news-topic-classifier-e5-base.tar.gz from "Files and versions" (or hf download maxafanasiev/news-topic-classifier-e5-base news-topic-classifier-e5-base.tar.gz --local-dir .), then tar xzf news-topic-classifier-e5-base.tar.gz; it unpacks into ./news-topic-classifier-e5-base/, ready for the code below. Windows 10+ has tar built in.
  2. Keep these files together: model.safetensors, config.json, tokenizer.json, tokenizer_config.json, thresholds.json.
  3. Load from the folder; HF_HUB_OFFLINE=1 makes sure nothing is fetched from the internet:
import json
import os
from pathlib import Path

os.environ["HF_HUB_OFFLINE"] = "1"

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

MODEL_DIR = Path("./news-topic-classifier-e5-base")
PREFIX = "passage: "
PARENTS = {"Theft": "Crime", "Assault": "Crime"}

tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_DIR).eval()
thresholds = json.loads((MODEL_DIR / "thresholds.json").read_text())
categories = list(thresholds)


@torch.no_grad()
def classify(articles: list[dict], max_len: int = 512) -> list[list[str]]:
    texts = [f"{PREFIX}{a['title']}\n{a['content']}" for a in articles]
    batch = tokenizer(texts, truncation=True, max_length=max_len, padding=True, return_tensors="pt")
    probs = torch.sigmoid(model(**batch).logits)
    results = []
    for row in probs:
        labels = [c for c, p in zip(categories, row.tolist()) if p >= thresholds[c]]
        labels += [PARENTS[c] for c in labels if c in PARENTS and PARENTS[c] not in labels]
        results.append(labels or ["Other"])
    return results

Requirements: torch and transformers (no huggingface_hub login needed at this point).

Results (gold set, CPU Apple M4 Pro, fp32)

Max tokens micro-F1 (random gold, 163) exact set match CPU latency
512 0.860 0.804 42 ms/article
256 0.836 0.785 20 ms/article

Training data: ~8k news articles labeled by an LLM teacher (Sonnet 5.5); gold built from three blind LLM labelers plus a judge, not hand-checked. Training data is private and not included.

Downloads last month
26
Safetensors
Model size
0.3B params
Tensor type
F32
·
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Model tree for maxafanasiev/news-topic-classifier-e5-base

Finetuned
(187)
this model