Training a Custom Financial Sentiment Model for Stock Market Prediction Using Hugging Face
Why domain-specific fine-tuning beats a general sentiment model
A general-purpose sentiment model trained on movie or product reviews misreads financial language: "shares fell less than expected" is favorable, "beats on revenue, misses on guidance" is mixed, and "record profits amid layoffs" is genuinely ambiguous. This tutorial fine-tunes FinBERT (a BERT variant pretrained on financial text) on a labeled headline dataset for three-class sentiment (positive/neutral/negative), which is the standard framing for using sentiment as one input signal alongside price/volume data — not a standalone prediction of price direction.
1. Technical architecture
┌────────────────────┐ ┌───────────────────┐ ┌────────────────────────┐
│ Labeled headline │──▶│ Tokenization & │──▶│ Fine-tuning job │
│ dataset (CSV/JSONL) │ │ train/val/test split │ │ (FinBERT + classification │
│ │ │ │ │ head, Trainer API) │
└────────────────────┘ └───────────────────┘ └────────────┬───────────┘
▼
┌────────────────────────┐
│ Evaluation (F1, confusion │
│ matrix per class) │
└────────────┬───────────┘
▼
┌────────────────────┐ ┌───────────────────┐ ┌────────────────────────┐
│ Daily headline pull │──▶│ FastAPI inference │──▶│ Aggregated daily sentiment │
│ (news API, batched) │ │ endpoint │ │ score per ticker (stored │
│ │ │ │ │ for downstream signal use) │
└────────────────────┘ └───────────────────┘ └────────────────────────┘
2. Complete code implementation
2.1 Data loading and tokenization
# sentiment/dataset.py
import pandas as pd
from datasets import Dataset, DatasetDict
from transformers import AutoTokenizer
from sklearn.model_selection import train_test_split
MODEL_CHECKPOINT = "ProsusAI/finbert"
LABEL2ID = {"negative": 0, "neutral": 1, "positive": 2}
ID2LABEL = {v: k for k, v in LABEL2ID.items()}
def load_and_split(csv_path: str, test_size: float = 0.15, val_size: float = 0.15) -> DatasetDict:
df = pd.read_csv(csv_path) # expected columns: headline, label ('positive'/'neutral'/'negative')
df = df.dropna(subset=["headline", "label"])
df = df[df["label"].isin(LABEL2ID.keys())]
if df.empty:
raise ValueError("No valid rows after filtering — check CSV columns and label values")
df["labels"] = df["label"].map(LABEL2ID)
train_df, temp_df = train_test_split(df, test_size=(test_size + val_size), stratify=df["labels"], random_state=42)
val_df, test_df = train_test_split(temp_df, test_size=test_size / (test_size + val_size), stratify=temp_df["labels"], random_state=42)
return DatasetDict({
"train": Dataset.from_pandas(train_df[["headline", "labels"]].reset_index(drop=True)),
"validation": Dataset.from_pandas(val_df[["headline", "labels"]].reset_index(drop=True)),
"test": Dataset.from_pandas(test_df[["headline", "labels"]].reset_index(drop=True)),
})
def tokenize_dataset(dataset_dict: DatasetDict) -> DatasetDict:
tokenizer = AutoTokenizer.from_pretrained(MODEL_CHECKPOINT)
def tokenize_fn(batch):
return tokenizer(batch["headline"], truncation=True, padding="max_length", max_length=64)
return dataset_dict.map(tokenize_fn, batched=True)
2.2 Fine-tuning script
# sentiment/train.py
import numpy as np
from transformers import (
AutoModelForSequenceClassification,
TrainingArguments,
Trainer,
EarlyStoppingCallback,
)
from sklearn.metrics import f1_score, precision_recall_fscore_support
from dataset import load_and_split, tokenize_dataset, MODEL_CHECKPOINT, ID2LABEL, LABEL2ID
OUTPUT_DIR = "checkpoints/finbert-sentiment"
def compute_metrics(eval_pred):
logits, labels = eval_pred
preds = np.argmax(logits, axis=-1)
precision, recall, f1, _ = precision_recall_fscore_support(labels, preds, average="macro", zero_division=0)
return {"macro_f1": f1, "precision": precision, "recall": recall}
def main():
raw = load_and_split("data/financial_headlines.csv")
tokenized = tokenize_dataset(raw)
model = AutoModelForSequenceClassification.from_pretrained(
MODEL_CHECKPOINT,
num_labels=3,
id2label=ID2LABEL,
label2id=LABEL2ID,
)
args = TrainingArguments(
output_dir=OUTPUT_DIR,
eval_strategy="epoch",
save_strategy="epoch",
learning_rate=2e-5,
per_device_train_batch_size=16,
per_device_eval_batch_size=32,
num_train_epochs=5,
weight_decay=0.01,
load_best_model_at_end=True,
metric_for_best_model="macro_f1",
logging_steps=25,
report_to="none",
)
trainer = Trainer(
model=model,
args=args,
train_dataset=tokenized["train"],
eval_dataset=tokenized["validation"],
compute_metrics=compute_metrics,
callbacks=[EarlyStoppingCallback(early_stopping_patience=2)],
)
trainer.train()
test_results = trainer.evaluate(tokenized["test"])
print("Test set results:", test_results)
trainer.save_model(f"{OUTPUT_DIR}/final")
tokenized["train"].info # no-op, keeps tokenizer object referenced for clarity
if __name__ == "__main__":
main()
2.3 Serving API
# sentiment/serve.py
import logging
from functools import lru_cache
from datetime import date
from collections import defaultdict
import torch
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForSequenceClassification
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("sentiment-api")
MODEL_DIR = "checkpoints/finbert-sentiment/final"
app = FastAPI(title="Financial Sentiment API")
tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_DIR)
model.eval()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
class HeadlineBatch(BaseModel):
ticker: str
headlines: list[str]
class ScoredHeadline(BaseModel):
headline: str
label: str
score: float
MAX_BATCH_SIZE = 64
@app.post("/score", response_model=list[ScoredHeadline])
def score_headlines(payload: HeadlineBatch):
if not payload.headlines:
raise HTTPException(status_code=400, detail="headlines list must not be empty")
if len(payload.headlines) > MAX_BATCH_SIZE:
raise HTTPException(status_code=400, detail=f"batch size exceeds limit of {MAX_BATCH_SIZE}")
try:
inputs = tokenizer(payload.headlines, truncation=True, padding=True, max_length=64, return_tensors="pt").to(device)
with torch.no_grad():
logits = model(**inputs).logits
probs = torch.softmax(logits, dim=-1)
top_scores, top_labels = torch.max(probs, dim=-1)
except RuntimeError as exc:
logger.exception("Inference failed")
raise HTTPException(status_code=500, detail="model inference error") from exc
id2label = model.config.id2label
return [
ScoredHeadline(headline=h, label=id2label[int(top_labels[i])], score=float(top_scores[i]))
for i, h in enumerate(payload.headlines)
]
@app.get("/daily-score/{ticker}")
def daily_score(ticker: str):
"""Aggregates today's cached headline scores into a single -1..+1 signal.
In production, back this with a real store (Redis/Postgres) instead of memory."""
scores = _daily_cache.get((ticker.upper(), date.today().isoformat()))
if not scores:
raise HTTPException(status_code=404, detail=f"No scored headlines cached for {ticker} today")
weight = {"negative": -1, "neutral": 0, "positive": 1}
weighted = sum(weight[s["label"]] * s["score"] for s in scores) / len(scores)
return {"ticker": ticker.upper(), "date": date.today().isoformat(), "sentiment_score": round(weighted, 4), "n_headlines": len(scores)}
_daily_cache: dict[tuple[str, str], list[dict]] = defaultdict(list)
3. Step-by-step configuration guide
- Install dependencies:
pip install transformers datasets scikit-learn pandas torch fastapi uvicorn --break-system-packages - Prepare
data/financial_headlines.csvwithheadline,labelcolumns. Public starting points include the Financial PhraseBank dataset — verify its license terms for your use case before commercial use, and supplement it with your own labeled data for better domain match. - Train:
python sentiment/train.py. Watchmacro_f1on the validation set each epoch —EarlyStoppingCallbackwill stop training if it stalls for 2 consecutive epochs. - Inspect the test metrics printed at the end of training; a macro F1 below roughly 0.75 on financial headlines usually means the label set needs cleaning or more training examples for the minority class (often "neutral", which tends to be under-represented).
- Serve the model:
uvicorn sentiment.serve:app --host 0.0.0.0 --port 8001 - Smoke test:
curl -X POST http://localhost:8001/score \ -H "Content-Type: application/json" \ -d '{"ticker": "AAPL", "headlines": ["Apple beats revenue estimates but issues cautious guidance"]}' - Wire the daily aggregation into your data pipeline — score each day's headlines as they arrive and append to
_daily_cache(or a real datastore), then read/daily-score/{ticker}from downstream signal-generation code such as the trading bot in the companion Alpaca tutorial.
4. Error handling and edge cases
- Class imbalance (financial headlines skew heavily toward "neutral") is addressed by using
average="macro"incompute_metricsrather than accuracy, so the reported score doesn't hide poor performance on the minority positive/negative classes. - Empty or oversized batch requests are explicitly rejected with 400 errors in
score_headlines, including a hard cap (MAX_BATCH_SIZE) to bound GPU memory use and keep p99 latency predictable under load. - Tokenizer truncation — headlines are capped at 64 tokens (
max_length=64), which comfortably covers virtually all real headlines; if you later feed full article bodies through the same endpoint, raise this and re-tune your batch size to avoid OOM. - CUDA OOM during inference — wrap the forward pass in
try/except RuntimeError, log the full trace server-side, and return a generic 500 to the client rather than exposing internal tensor shapes or memory info. - Stale cached daily scores — the naive in-memory
_daily_cachehere is for demonstration only; in production, key by trading day (not calendar day, to handle weekends/holidays correctly) and expire/reset it explicitly rather than letting it grow unbounded. - Label leakage during data splitting —
train_test_splitusesstratify=df["labels"]to preserve class balance across train/val/test, and splitting happens once before any tokenization or augmentation, preventing near-duplicate headlines from leaking across the split boundary if your raw data has duplicates (deduplicate the source CSV first if that's a risk). - Model overconfidence — softmax scores from a fine-tuned classifier are not true probabilities; if you use
sentiment_scoreas a trading signal input, calibrate it (e.g., temperature scaling on a held-out set) rather than treating raw softmax output as calibrated confidence.