energy-news-classifier-ner-multitask

Energy Intelligence Multitask Model

QuantBridge / energy-intelligence-multitask

A single DistilBERT model with a shared encoder and two task heads for energy and financial news analysis. One forward pass returns both named entities and topic labels simultaneously.

---

Architecture

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

---

NER Label Space — 19 BIO tags

Each type uses standard BIO tagging: B-<TYPE> starts a span, I-<TYPE> continues it, O marks non-entities.

---

Classification Label Space — 10 topic labels

Note on classification scores: The classification head was trained on AG News + Reuters + Kaggle — datasets dominated by general business and macro content. Domain-specific labels (energy, shipping, risk, regulation, stocks) score lower as a result. The relative ranking of scores is semantically meaningful even when raw values are low. See Limitations.

---

Test Results

Evaluated on 40 real-world energy & financial news headlines across 9 domain groups (ENERGY, GEOPOLITICAL, SHIPPING, TRADE, MACRO, CORPORATE, REGULATION, TECHNOLOGY, STOCKS, RISK).

NER Results

Entity type frequency:

Key NER observations:

• COMMODITY is the top entity type — the model reliably extracts energy goods (oil, crude, LNG, natural gas, hydrogen) and commodities (aluminum, solar panels)
• COUNTRY and ORGANIZATION fire consistently across all domain groups
• COMPANY detection is accurate: correctly identifies both energy majors (ExxonMobil, Shell, BP) and non-energy companies (Tesla, Maersk, Vestas)
• LOCATION captures geopolitically important hotspots correctly (Red Sea, Strait of Hormuz, Gulf of Mexico, North Sea)
• MARKET fires on "S&P 500" but misses "Brent" and "WTI" — likely a tokenisation artefact where these are split sub-words during BertTokenizer processing
• PERSON and INFRASTRUCTURE did not fire on this specific test set; these types are present in the model's label vocabulary and will activate on appropriate inputs

Classification Results (threshold = 0.20)

Label activation frequency across 40 headlines:

Domain-group heatmap (>>> = group average score ≥ 0.35):

Key classification observations:

• macro is the dominant label across all domain groups — a direct consequence of training data composition (AG News World category and Kaggle both map heavily to macro)
• politics fires on TECHNOLOGY and GEOPOLITICAL groups, which is semantically reasonable (government energy policy, sanctions)
• Domain-specific labels (energy, shipping, risk, regulation, stocks) score consistently low — these categories are underrepresented in training data
• The score ranking is meaningful: for ENERGY headlines, energy consistently ranks above trade, shipping, and regulation even when below threshold — the model has learned the correct relative associations

---

Usage

Important: This model uses custom architecture files. Always pass trust_remote_code=True.

Installation

pip install transformers torch

Full inference — NER + Classification

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}")

Expected output:

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

Active topic labels:
  politics: 0.362
  macro: 0.357

NER only — decode all entity spans

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)

Classification only — get all label scores

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))}

---

Training Details

Encoder

Transferred from QuantBridge/energy-intelligence-multitask-custom-ner — DistilBERT fine-tuned on energy and financial news for BIO entity recognition.

NER Head

Weights transferred directly from the NER backbone (classifier.* → ner_classifier.*). No additional NER training was performed.

Classification Head

Trained separately from scratch on a merged corpus:

Training split: 80% train / 10% validation / 10% test, seed 42.

Hyperparameters:

---

Model Files

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

---

Limitations

• English only — trained exclusively on English-language news text.
• Classification data bias — training corpora (AG News, Kaggle) are dominated by business and macro content. Domain-specific labels (energy, shipping, risk, regulation, stocks) score lower across the board and may not cross common thresholds even when semantically correct. A recommended threshold for this model is 0.20 rather than the default 0.50.
• NER on headlines — the NER head was fine-tuned on short news headlines; performance may be lower on long-form documents.
• Max length — inputs are truncated to 128 tokens. Longer texts should be chunked.
• PERSON / INFRASTRUCTURE — these entity types exist in the label vocabulary but fired less frequently on financial news headlines compared to COMPANY, COUNTRY, and COMMODITY.
• Not for trading — this model is intended as an intelligence tagging layer, not for real-time trading or financial decision-making.

---

Intended Use

This model is the tagging layer in an energy intelligence pipeline:

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

---

Related Models

• QuantBridge/energy-intelligence-multitask-custom-ner — NER backbone (encoder source)
• QuantBridge/energy-news-classifier-ner-multitask — Classification-only model (single head)

---

Citation

@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}},
}