QuantBridge/energy-news-classifier-ner-multitask

Input headline
      |
BertTokenizer  (do_lower_case=True, max_length=128)
      |
DistilBERT encoder  (6 layers · 768 dim · 12 heads · ~67M params)
[weights from QuantBridge/energy-intelligence-multitask-custom-ner]
      |
      +──────────────────────────────────────────┐
      |                                          |
  all token hidden states                   [CLS] hidden state
      |                                          |
  Dropout(0.1)                         Linear(768→768) + ReLU
      |                                     Dropout(0.2)
  Linear(768→19)                        Linear(768→10)
      |                                          |
  NER logits                             CLS logits
  argmax → BIO entity tags           sigmoid → topic probabilities

pip install transformers torch

import torch
import numpy as np
from transformers import AutoTokenizer, AutoConfig, AutoModel

MODEL_ID = "QuantBridge/energy-intelligence-multitask"

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
model     = AutoModel.from_pretrained(MODEL_ID, trust_remote_code=True)
model.eval()

def sigmoid(x):
    return 1 / (1 + np.exp(-x))

def predict(text: str, cls_threshold: float = 0.20):
    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=128)
    inputs.pop("token_type_ids", None)   # DistilBERT does not use these

    with torch.no_grad():
        output = model(**inputs)

    tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])

    # ── Named Entity Recognition ──────────────────────────────────────────
    ner_id2label = {int(k): v for k, v in model.config.ner_id2label.items()}
    tag_ids = output.ner_logits[0].argmax(-1).tolist()

    entities = []
    current = None
    for token, tag_id in zip(tokens, tag_ids):
        if token in ("[CLS]", "[SEP]", "[PAD]"):
            if current: entities.append(current); current = None
            continue
        tag = ner_id2label[tag_id]
        if tag.startswith("B-"):
            if current: entities.append(current)
            current = {"text": token.replace("##", ""), "type": tag[2:]}
        elif tag.startswith("I-") and current:
            current["text"] += token[2:] if token.startswith("##") else f" {token}"
        else:
            if current: entities.append(current); current = None
    if current: entities.append(current)

    # ── Topic Classification ──────────────────────────────────────────────
    cls_id2label = {int(k): v for k, v in model.config.cls_id2label.items()}
    probs  = sigmoid(output.cls_logits[0].numpy())
    topics = {cls_id2label[i]: float(probs[i]) for i in range(len(probs))}
    active = {lbl: p for lbl, p in topics.items() if p >= cls_threshold}

    return entities, active

# Example
headline = "Russia cuts natural gas flows to Poland and Bulgaria following payment dispute"
entities, topics = predict(headline)

print("Entities found:")
for e in entities:
    print(f"  [{e['type']}]  {e['text']}")

print("\nActive topic labels:")
for topic, score in sorted(topics.items(), key=lambda x: -x[1]):
    print(f"  {topic}: {score:.3f}")

Entities found:
  [COUNTRY]   Russia
  [COUNTRY]   Poland
  [COUNTRY]   Bulgaria
  [COMMODITY] natural gas

Active topic labels:
  politics: 0.362
  macro: 0.357

def get_entities(text: str) -> list[dict]:
    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=128)
    inputs.pop("token_type_ids", None)
    with torch.no_grad():
        output = model(**inputs)
    tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
    ner_id2label = {int(k): v for k, v in model.config.ner_id2label.items()}
    tag_ids = output.ner_logits[0].argmax(-1).tolist()
    # ... (decode as shown above)

def get_topic_scores(text: str) -> dict[str, float]:
    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=128)
    inputs.pop("token_type_ids", None)
    with torch.no_grad():
        output = model(**inputs)
    cls_id2label = {int(k): v for k, v in model.config.cls_id2label.items()}
    probs = sigmoid(output.cls_logits[0].numpy())
    return {cls_id2label[i]: float(probs[i]) for i in range(len(probs))}

energy-intelligence-multitask/
  configuration_energy_multitask.py   # EnergyMultitaskConfig (DistilBertConfig subclass)
  modeling_energy_multitask.py        # EnergyMultitaskModel  (two-head architecture)
  config.json                         # Serialised config with auto_map
  model.safetensors                   # Combined weights (~256 MB)
  tokenizer.json                      # Fast tokenizer
  tokenizer_config.json               # Tokenizer settings

Raw news headline
       ↓
EnergyMultitaskModel (this model)
       ↓
  entities  ──────────────────→  who / what / where
  topic labels  ──────────────→  energy / risk / trade / macro / ...
       ↓
Structured intelligence signal for downstream analysis

@misc{quantbridge2025multitask,
  author       = {QuantBridge},
  title        = {Energy Intelligence Multitask Model (NER + Classification)},
  year         = {2025},
  publisher    = {Hugging Face},
  howpublished = {\url{https://huggingface.co/QuantBridge/energy-intelligence-multitask}},
}