Views
No views yet
| Metric | Score |
|---|---|
| Accuracy | 0.846 |
| F1-macro | 0.681 |
| Class | Precision | Recall | F1 | Support |
|---|---|---|---|---|
| Financial/Earnings | 0.809 | 0.885 | 0.845 | 148 |
| Operations/Business | 0.805 | 0.899 | 0.849 | 197 |
| Accidents/Safety | 0.333 | 0.200 | 0.250 | 5 |
| Labor/Layoffs | 0.909 | 0.870 | 0.889 | 23 |
| Legal/Regulatory | 0.875 | 0.875 | 0.875 | 56 |
| Fraud/Misconduct | 0.333 | 0.143 | 0.200 | 7 |
| Other/External | 0.914 | 0.806 | 0.857 | 252 |
1import os, json, sys, importlib.util, torch
2from typing import Dict, List
3from huggingface_hub import hf_hub_download
4from transformers import AutoTokenizer
5
6REPO = "MasSolutions/Alpaxa-FinBERT-Tone-7class"
7
8TOPIC_LABELS: List[str] = [
9 "Financial/Earnings", # 0
10 "Operations/Business", # 1
11 "Accidents/Safety", # 2
12 "Labor/Layoffs", # 3
13 "Legal/Regulatory", # 4
14 "Fraud/Misconduct", # 5
15 "Other/External", # 6
16]
17SENTIMENT_LABELS = ["Neutral", "Positive", "Negative"]
18
19def sanitize_cfg(cfg: Dict) -> Dict:
20 """Keep only keys accepted by the model constructor."""
21 allowed = {
22 "encoder_ckpt",
23 "sentiment_ckpt",
24 "num_labels",
25 "use_mean_pool",
26 "use_focal",
27 "focal_gamma",
28 }
29 return {k: v for k, v in cfg.items() if k in allowed}
30
31# 1) Download the modeling file and import it dynamically
32model_py = hf_hub_download(REPO, "modeling_multi_head_finbert.py")
33spec = importlib.util.spec_from_file_location("modeling_multi_head_finbert", model_py)
34mod = importlib.util.module_from_spec(spec)
35sys.modules["modeling_multi_head_finbert"] = mod
36spec.loader.exec_module(mod)
37
38MultiHeadFinBERT = mod.MultiHeadFinBERT
39
40# 2) Load config / weights / tokenizer
41cfg_path = hf_hub_download(REPO, "config.json")
42state_path = None
43try:
44 state_path = hf_hub_download(REPO, "model.safetensors")
45 use_safetensors = True
46except Exception:
47 state_path = hf_hub_download(REPO, "model_state.pth")
48 use_safetensors = False
49
50with open(cfg_path) as f:
51 cfg_raw = json.load(f)
52cfg = sanitize_cfg(cfg_raw)
53
54tok = AutoTokenizer.from_pretrained(REPO)
55
56# 3) Build and load model
57model = MultiHeadFinBERT(**cfg)
58if use_safetensors:
59 from safetensors.torch import load_file as load_safetensors
60 state = load_safetensors(state_path, device="cpu")
61else:
62 state = torch.load(state_path, map_location="cpu")
63model.load_state_dict(state, strict=True)
64model.eval()
65
66# 4) Inference
67texts = ["Delta flight crashed. 198 dead 2 missing."] # Test Headline
68batch = tok(texts, return_tensors="pt", truncation=True, max_length=128, padding=True)
69batch.pop("token_type_ids", None)
70
71with torch.no_grad():
72 out = model(**batch)
73topic_probs = out["logits"].softmax(-1) # [B, 7]
74senti_probs = out["sentiment_logits"].softmax(-1) # [B, 3]
75
76# 5) Pretty print with labels
77for i, text in enumerate(texts):
78 tp = topic_probs[i]
79 sp = senti_probs[i]
80
81 pred_topic_idx = int(tp.argmax().item())
82 pred_senti_idx = int(sp.argmax().item())
83
84 print(f"\nText: {text}")
85 print(f"→ Topic: {TOPIC_LABELS[pred_topic_idx]} (p={tp[pred_topic_idx]:.4f})")
86 print(f"→ Sentiment: {SENTIMENT_LABELS[pred_senti_idx]} "
87 f"(neg={sp[0]:.3f}, neu={sp[1]:.3f}, pos={sp[2]:.3f})")
88
89 # Top-3 topics
90 topk = torch.topk(tp, k=min(3, tp.numel()))
91 print("Top-3 topic probs:")
92 for idx, prob in zip(topk.indices.tolist(), topk.values.tolist()):
93 print(f" {TOPIC_LABELS[idx]:<20} {prob:.4f}")