FinBERT fine-tuning on move labels: spec, implementation, results

agents/aurora_processing.py

@@ -75,6 +75,24 @@ def _assign_label(pct_change, threshold):
 # Processing node
 # ---------------------------------------------------------------------------
 
+def _assign_split(df, dataset_end=None):
+    """Add a time-based train/test `split` column: earliest 80% of dates = train,
+    rest = test. Date-based (not random) so no future information leaks into
+    training. `dataset_end` (YYYY-MM-DD) optionally drops later rows first —
+    e.g. "2019-12-31" keeps the COVID regime out of the test window (see
+    docs/superpowers/specs/2026-07-02-finbert-finetune-design.md)."""
+    if dataset_end:
+        before = len(df)
+        df = df[pd.to_datetime(df["date"]) <= pd.to_datetime(dataset_end)].reset_index(drop=True)
+        print(f"  Dataset ended at {dataset_end}: dropped {before - len(df):,} later rows")
+    split_cut = pd.to_datetime(df["date"]).quantile(0.8)
+    df["split"] = ["train" if d <= split_cut else "test"
+                   for d in pd.to_datetime(df["date"])]
+    print(f"  Split at {split_cut.date()}: "
+          f"{(df['split'] == 'train').sum():,} train / {(df['split'] == 'test').sum():,} test")
+    return df
+
+
 def processing_node(state: PipelineState) -> dict:
     """LangGraph node — runs the full processing pipeline and returns the
     output file path as `processed_data_path` in the state."""
@@ -145,10 +163,13 @@ def processing_node(state: PipelineState) -> dict:
     ].reset_index(drop=True)
     print(f"  Outliers dropped: {before_outliers - len(df):,}. Final: {len(df):,} rows")
 
+    # 5b. Time-based train/test split (and optional dataset end date).
+    df = _assign_split(df, state.get("dataset_end"))
+
     # 6. Export (Handoff 1 schema from data_contracts.md)
     output_cols = [
         "article_id", "date", "ticker", "article_title",
-        "price_t", "price_t1", "pct_change", "label",
+        "price_t", "price_t1", "pct_change", "label", "split",
     ]
     df[output_cols].to_csv(out_path, index=False)
     print(f"[processing] Written to {out_path}")

agents/finetune_finbert.py

@@ -0,0 +1,182 @@
+"""Fine-tune FinBERT on the up/down/neutral move labels (standalone, run once).
+
+Not part of the pipeline loop. Trains on processed_data.csv's split=train rows
+(the last 10% of train *dates* held out for validation), evaluates on split=test,
+saves the best-validation model to outputs/finbert_finetuned/ and a
+finetune_report.json reproducibility record next to it. The generated classifier
+(agents/nadi_classifier.py) picks the saved model up automatically on its next run.
+
+Design + motivation: docs/superpowers/specs/2026-07-02-finbert-finetune-design.md.
+
+    uv run agents/finetune_finbert.py --data data/processed_data.csv
+"""
+
+import argparse
+import json
+import os
+import random
+
+import pandas as pd
+import torch
+from torch.utils.data import DataLoader, Dataset
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+
+MODEL = "ProsusAI/finbert"
+LABELS = ["up", "down", "neutral"]
+LABEL2ID = {name: i for i, name in enumerate(LABELS)}
+HF_TOKEN = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
+
+
+class HeadlineDataset(Dataset):
+    def __init__(self, titles, labels, tokenizer, max_length):
+        self.enc = tokenizer(list(titles), truncation=True, max_length=max_length,
+                             padding="max_length", return_tensors="pt")
+        self.labels = torch.tensor([LABEL2ID[label] for label in labels])
+
+    def __len__(self):
+        return len(self.labels)
+
+    def __getitem__(self, i):
+        return {"input_ids": self.enc["input_ids"][i],
+                "attention_mask": self.enc["attention_mask"][i],
+                "labels": self.labels[i]}
+
+
+def _device():
+    if torch.backends.mps.is_available():
+        return torch.device("mps")
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    return torch.device("cpu")
+
+
+def _split_frames(df: pd.DataFrame):
+    """train/val/test frames. Validation is the last 10% of *train* dates so it
+    is time-ordered like the real test set (no random leakage)."""
+    if "split" not in df.columns:
+        raise SystemExit(
+            "processed_data.csv has no 'split' column — regenerate it with the "
+            "updated Processing agent (agents/aurora_processing.py) first."
+        )
+    train_all = df[df["split"] == "train"]
+    test = df[df["split"] == "test"]
+    val_cut = pd.to_datetime(train_all["date"]).quantile(0.9)
+    is_val = pd.to_datetime(train_all["date"]) > val_cut
+    return train_all[~is_val], train_all[is_val], test
+
+
+@torch.no_grad()
+def evaluate(model, loader, device):
+    model.eval()
+    correct_total, per_class = 0, {i: [0, 0] for i in range(len(LABELS))}  # hits, count
+    for batch in loader:
+        preds = model(input_ids=batch["input_ids"].to(device),
+                      attention_mask=batch["attention_mask"].to(device)).logits.argmax(dim=1).cpu()
+        for pred, true in zip(preds, batch["labels"]):
+            per_class[int(true)][1] += 1
+            if pred == true:
+                per_class[int(true)][0] += 1
+                correct_total += 1
+    n = sum(count for _, count in per_class.values())
+    return (correct_total / n if n else 0.0,
+            {LABELS[i]: round(hits / count, 4) if count else 0.0
+             for i, (hits, count) in per_class.items()})
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
+    ap.add_argument("--data", default="data/processed_data.csv")
+    ap.add_argument("--out-dir", default="outputs/finbert_finetuned")
+    ap.add_argument("--epochs", type=int, default=3)
+    ap.add_argument("--batch-size", type=int, default=16)
+    ap.add_argument("--lr", type=float, default=2e-5)
+    ap.add_argument("--max-length", type=int, default=128)
+    ap.add_argument("--seed", type=int, default=42)
+    ap.add_argument("--limit", type=int, default=None,
+                    help="cap rows per split (smoke tests)")
+    args = ap.parse_args()
+
+    random.seed(args.seed)
+    torch.manual_seed(args.seed)
+
+    df = pd.read_csv(args.data)
+    train, val, test = _split_frames(df)
+    if args.limit:
+        train, val, test = train.head(args.limit), val.head(args.limit), test.head(args.limit)
+    print(f"[finetune] rows: {len(train)} train / {len(val)} val / {len(test)} test "
+          f"| train dates {train['date'].min()}..{train['date'].max()}")
+
+    device = _device()
+    tokenizer = AutoTokenizer.from_pretrained(MODEL, token=HF_TOKEN)
+    model = AutoModelForSequenceClassification.from_pretrained(
+        MODEL, token=HF_TOKEN, num_labels=len(LABELS),
+        id2label={i: name for i, name in enumerate(LABELS)}, label2id=LABEL2ID,
+        ignore_mismatched_sizes=True,  # fresh 3-way head over the sentiment one
+    ).to(device)
+
+    def loader(frame, shuffle=False):
+        return DataLoader(
+            HeadlineDataset(frame["article_title"], frame["label"], tokenizer, args.max_length),
+            batch_size=args.batch_size, shuffle=shuffle)
+
+    train_loader, val_loader, test_loader = loader(train, shuffle=True), loader(val), loader(test)
+
+    # Weighted CE so the neutral-heavy train distribution doesn't collapse the
+    # up/down classes to noise.
+    counts = train["label"].value_counts()
+    weights = torch.tensor([len(train) / (len(LABELS) * counts.get(name, 1))
+                            for name in LABELS], dtype=torch.float).to(device)
+    loss_fn = torch.nn.CrossEntropyLoss(weight=weights)
+    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr)
+
+    best_val, history = -1.0, []
+    os.makedirs(args.out_dir, exist_ok=True)
+    for epoch in range(1, args.epochs + 1):
+        model.train()
+        running = 0.0
+        for step, batch in enumerate(train_loader, 1):
+            optimizer.zero_grad()
+            logits = model(input_ids=batch["input_ids"].to(device),
+                           attention_mask=batch["attention_mask"].to(device)).logits
+            loss = loss_fn(logits, batch["labels"].to(device))
+            loss.backward()
+            optimizer.step()
+            running += loss.item()
+            if step % 100 == 0:
+                print(f"  epoch {epoch} step {step}/{len(train_loader)} "
+                      f"loss {running / step:.4f}")
+        val_acc, _ = evaluate(model, val_loader, device)
+        history.append({"epoch": epoch, "train_loss": round(running / max(len(train_loader), 1), 4),
+                        "val_accuracy": round(val_acc, 4)})
+        print(f"[finetune] epoch {epoch}: val accuracy {val_acc:.4f}")
+        if val_acc > best_val:  # keep the best-validation weights, not the last
+            best_val = val_acc
+            model.save_pretrained(args.out_dir)
+            tokenizer.save_pretrained(args.out_dir)
+
+    # Score the SAVED (best-val) model on the held-out test split.
+    best = AutoModelForSequenceClassification.from_pretrained(args.out_dir).to(device)
+    test_acc, test_class = evaluate(best, test_loader, device)
+    print(f"[finetune] test accuracy {test_acc:.4f} per-class {test_class}")
+
+    report = {
+        "base_model": MODEL,
+        "model_dir": args.out_dir,
+        "test_accuracy": round(test_acc, 4),
+        "test_class_accuracy": test_class,
+        "best_val_accuracy": round(best_val, 4),
+        "epochs": history,
+        "train_rows": len(train), "val_rows": len(val), "test_rows": len(test),
+        "train_date_range": [str(train["date"].min()), str(train["date"].max())],
+        "test_date_range": [str(test["date"].min()), str(test["date"].max())],
+        "params": {"epochs": args.epochs, "batch_size": args.batch_size, "lr": args.lr,
+                   "max_length": args.max_length, "seed": args.seed, "limit": args.limit},
+    }
+    report_path = os.path.join(os.path.dirname(args.out_dir) or ".", "finetune_report.json")
+    with open(report_path, "w", encoding="utf-8") as f:
+        json.dump(report, f, indent=2)
+    print(f"[finetune] report written to {report_path}")
+
+
+if __name__ == "__main__":
+    main()

agents/jack_manager.py

@@ -235,21 +235,30 @@ def write_retune(state: ManagerState) -> dict:
     return {"decision_log": [f"iteration {state['iteration']}: wrote retune_request.json"]}
 
 
-def proceed(state: ManagerState) -> dict:
-    """sample_for_explanation.csv (Handoff 4) — drawn from predictions_test.csv.
-    Terminal: hands off to Freddi and waits for explanations.csv."""
+def write_sample(predictions_path: str, sample_size: int = 300) -> int:
+    """Write sample_for_explanation.csv (Handoff 4) from a predictions file.
+    Shared by the proceed node and the pipeline's best-iteration restore, so the
+    sample format has exactly one definition."""
     import pandas as pd
 
-    preds = pd.read_csv(state.get("predictions_path", "mock_data/predictions_test.csv"))
+    preds = pd.read_csv(predictions_path)
     sample = (preds[["article_id", "article_title", "predicted_label", "label",
                      "confidence", "prob_up", "prob_down", "prob_neutral"]]
               .rename(columns={"label": "actual_label"}))
-    n = min(len(sample), state.get("sample_size", 300))
+    n = min(len(sample), sample_size)
     if n < len(sample):                                   # only subsample when needed
         sample = sample.sample(n=n, random_state=42)      # representative + reproducible
-    _write_decision(state)
     os.makedirs(OUTPUT_DIR, exist_ok=True)
     sample.to_csv(os.path.join(OUTPUT_DIR, "sample_for_explanation.csv"), index=False)
+    return n
+
+
+def proceed(state: ManagerState) -> dict:
+    """sample_for_explanation.csv (Handoff 4) — drawn from predictions_test.csv.
+    Terminal: hands off to Freddi and waits for explanations.csv."""
+    _write_decision(state)
+    n = write_sample(state.get("predictions_path", "mock_data/predictions_test.csv"),
+                     state.get("sample_size", 300))
     return {"decision_log": [f"iteration {state['iteration']}: wrote sample_for_explanation.csv ({n} rows)"]}
 
 

agents/nadi_classifier.py

@@ -32,6 +32,7 @@
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
 MODEL = "ProsusAI/finbert"
+MODEL_DIR = {model_dir}  # fine-tuned weights dir; None -> pretrained sentiment head
 MAX_LENGTH = {max_length}
 THRESHOLD = {threshold}
 FOCUS_LABELS = {focus_labels}
@@ -42,12 +43,18 @@
 # faster downloads); fall back to unauthenticated access when it is absent.
 HF_TOKEN = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
 
-tokenizer = AutoTokenizer.from_pretrained(MODEL, token=HF_TOKEN)
-model = AutoModelForSequenceClassification.from_pretrained(MODEL, token=HF_TOKEN)
+if MODEL_DIR and os.path.isdir(MODEL_DIR):
+    # Fine-tuned on our move labels: id2label is already up/down/neutral.
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)
+    model = AutoModelForSequenceClassification.from_pretrained(MODEL_DIR)
+    ID2OURS = {{i: model.config.id2label[i].lower() for i in model.config.id2label}}
+else:
+    # Pretrained sentiment head: translate sentiment -> move direction.
+    tokenizer = AutoTokenizer.from_pretrained(MODEL, token=HF_TOKEN)
+    model = AutoModelForSequenceClassification.from_pretrained(MODEL, token=HF_TOKEN)
+    ID2OURS = {{i: SENTIMENT_TO_LABEL[model.config.id2label[i].lower()] for i in model.config.id2label}}
 model.eval()
 
-ID2OURS = {{i: SENTIMENT_TO_LABEL[model.config.id2label[i].lower()] for i in model.config.id2label}}
-
 def classify(title: str) -> dict:
     inputs = tokenizer(title, return_tensors="pt", truncation=True, max_length=MAX_LENGTH)
     with torch.no_grad():
@@ -90,11 +97,19 @@ def main(src: str, dst: str) -> None:
     if not rows:
         raise ValueError("Source file is empty")
 
-    out_cols = list(rows[0].keys()) + [
+    # When Aurora supplies a train/test split, predict only the held-out test
+    # rows (training rows were seen by the fine-tuned model). Without a split
+    # column every row is treated as test (pre-split behavior).
+    if "split" in rows[0]:
+        rows = [r for r in rows if r["split"] == "test"]
+        if not rows:
+            raise ValueError("No split=test rows in source file")
+
+    # split goes last per the contract, wherever it sat in the input.
+    out_cols = [c for c in rows[0].keys() if c != "split"] + [
         "predicted_label", "confidence", "prob_up", "prob_down", "prob_neutral", "split"
     ]
-
-    # Process all rows and mark them as test split
+
     for row in rows:
         pred_data = classify(row["article_title"])
         row.update(pred_data)
@@ -131,11 +146,19 @@ def generate_code(state: PipelineState) -> dict:
     code_path = state.get("classifier_code_path") or os.path.join(OUTPUT_DIR, "classifier.py")
     os.makedirs(os.path.dirname(code_path) or ".", exist_ok=True)
 
+    # Prefer fine-tuned weights when a training run has produced them (see
+    # agents/finetune_finbert.py); otherwise the generated script falls back to
+    # the pretrained sentiment head.
+    model_dir = state.get("model_dir") or os.path.join(OUTPUT_DIR, "finbert_finetuned")
+    if not os.path.isdir(model_dir):
+        model_dir = None
+
     formatted_code = CLASSIFIER_TEMPLATE.format(
         threshold=threshold,
         max_length=max_length,
         focus_labels=repr(focus_labels),
-        boost_factor=boost_factor
+        boost_factor=boost_factor,
+        model_dir=repr(model_dir)
     )
 
     with open(code_path, "w", encoding="utf-8") as f:
@@ -154,7 +177,7 @@ def generate_code(state: PipelineState) -> dict:
         f.write(formatted_code)
 
     metadata = {
-        "model_name": "ProsusAI/finbert",
+        "model_name": model_dir or "ProsusAI/finbert",
         "fine_tuning_params": {
             "threshold": threshold,
             "max_length": max_length,