Upload eval/train_and_eval.py with huggingface_hub

eval/train_and_eval.py

@@ -0,0 +1,144 @@
+"""FINSENSE stage-1 trainer — deberta-v3-base on Financial PhraseBank 50agree.
+
+Research-verified recipe (FINSENSE.md): generic encoder, no DAPT, our own
+published stratified split (FPB has no official one), macro-F1 primary,
+confusion matrix logged (pos/neutral confusion = the known ceiling).
+
+Usage: python3 finsense_train.py [--seed 42] [--smoke] [--base microsoft/deberta-v3-base]
+Runs on MPS (Mac) or CUDA. Smoke = 200 steps sanity.
+"""
+import argparse
+import json
+import os
+
+import io
+import zipfile
+
+import numpy as np
+import torch
+from huggingface_hub import hf_hub_download
+from sklearn.metrics import accuracy_score, confusion_matrix, f1_score
+from transformers import (AutoModelForSequenceClassification, AutoTokenizer,
+                          Trainer, TrainingArguments)
+
+p = argparse.ArgumentParser()
+p.add_argument("--seed", type=int, default=42)
+p.add_argument("--smoke", action="store_true")
+p.add_argument("--base", default="microsoft/deberta-v3-base")
+p.add_argument("--corpus", default=None,
+               help="optional finsense_arm_*.jsonl — replaces the FPB train split "
+                    "as training data (val/test splits stay FPB, untouched)")
+args = p.parse_args()
+
+_arm = ("-" + os.path.basename(args.corpus).replace("finsense_arm_", "arm").replace(".jsonl", "")) if args.corpus else ""
+OUT = os.path.expanduser(f"~/finsense_runs/{args.base.split('/')[-1]}-s{args.seed}{_arm}")
+os.makedirs(OUT, exist_ok=True)
+LABELS = ["negative", "neutral", "positive"]  # FPB label order 0/1/2
+
+# parse the raw FPB zip directly (script-datasets unsupported, no parquet branch)
+zp = hf_hub_download("takala/financial_phrasebank", "data/FinancialPhraseBank-v1.0.zip",
+                     repo_type="dataset")
+LAB2ID = {"negative": 0, "neutral": 1, "positive": 2}
+rows = []
+with zipfile.ZipFile(zp) as z:
+    with z.open("FinancialPhraseBank-v1.0/Sentences_50Agree.txt") as fh:
+        for line in io.TextIOWrapper(fh, encoding="latin-1"):
+            line = line.strip()
+            if not line or "@" not in line:
+                continue
+            sent, lab = line.rsplit("@", 1)
+            rows.append({"sentence": sent.strip(), "label": LAB2ID[lab.strip()]})
+print(f"FPB 50agree rows: {len(rows)}", flush=True)
+
+# published split: stratified 80/10/10, split-seed FIXED at 42 regardless of
+# training seed (same rule as Parable corpus splits)
+import random as _rnd
+_rnd.Random(42).shuffle(rows)
+by_label = {i: [] for i in range(3)}
+for r in rows:
+    by_label[r["label"]].append(r)
+train, val, test = [], [], []
+for lab, rows in by_label.items():
+    n = len(rows)
+    test += rows[: int(n * 0.10)]
+    val += rows[int(n * 0.10): int(n * 0.20)]
+    train += rows[int(n * 0.20):]
+print(f"split: train={len(train)} val={len(val)} test={len(test)}", flush=True)
+
+if args.corpus:
+    train = [json.loads(l) for l in open(args.corpus)]
+    import random as _r2
+    _r2.Random(args.seed).shuffle(train)
+    print(f"corpus override: {args.corpus} -> train={len(train)}", flush=True)
+
+_trc = {"trust_remote_code": True} if "NeoBERT" in args.base else {}
+tok = AutoTokenizer.from_pretrained(args.base, **_trc)
+
+def enc(rows):
+    from datasets import Dataset
+    d = Dataset.from_list(rows)
+    return d.map(lambda b: tok(b["sentence"], truncation=True, max_length=128),
+                 batched=True, remove_columns=["sentence"])
+
+torch.manual_seed(args.seed)
+model = AutoModelForSequenceClassification.from_pretrained(
+    args.base, num_labels=3, **_trc,
+    id2label=dict(enumerate(LABELS)), label2id={l: i for i, l in enumerate(LABELS)})
+
+def metrics(pred):
+    y, yhat = pred.label_ids, pred.predictions.argmax(-1)
+    return {"accuracy": accuracy_score(y, yhat),
+            "macro_f1": f1_score(y, yhat, average="macro")}
+
+from transformers import TrainerCallback
+
+class MPSCacheFlush(TrainerCallback):
+    """torch-MPS leaks on deberta attention kernels; flush or the kernel
+    OOM-kills the process silently around step ~100 on a 16GB machine."""
+    def on_step_end(self, args, state, control, **kw):
+        if state.global_step % 25 == 0 and torch.backends.mps.is_available():
+            torch.mps.empty_cache()
+
+import transformers as _tf
+_tok_kw = {"processing_class": tok} if int(_tf.__version__.split(".")[0]) >= 5 else {"tokenizer": tok}
+trainer = Trainer(
+    model=model,
+    **_tok_kw,
+    callbacks=[MPSCacheFlush()],
+    train_dataset=enc(train),
+    eval_dataset=enc(val),
+    compute_metrics=metrics,
+    args=TrainingArguments(
+        output_dir=f"{OUT}/ckpt",
+        num_train_epochs=1 if args.smoke else 5,
+        max_steps=50 if args.smoke else -1,
+        per_device_train_batch_size=8, gradient_accumulation_steps=2, per_device_eval_batch_size=16,
+        learning_rate=(1e-5 if "deberta" in args.base else 2e-5), warmup_ratio=0.1, weight_decay=0.01,
+        eval_strategy="epoch" if not args.smoke else "no",
+        save_strategy="epoch" if not args.smoke else "no",
+        load_best_model_at_end=not args.smoke,
+        metric_for_best_model="macro_f1",
+        logging_steps=25, seed=args.seed, report_to=[],
+        use_cpu=False,
+        # deberta-v3's disentangled attention overflows under fp16 -> class
+        # collapse; transformers v5 auto-enables mixed precision on CUDA
+        fp16=False, bf16=False,
+    ),
+)
+trainer.train()
+
+# final: held-out test (never used for selection)
+pred = trainer.predict(enc(test))
+y, yhat = pred.label_ids, pred.predictions.argmax(-1)
+res = {"base": args.base, "seed": args.seed, "smoke": args.smoke,
+       "test_accuracy": float(accuracy_score(y, yhat)),
+       "test_macro_f1": float(f1_score(y, yhat, average="macro")),
+       "confusion_matrix": confusion_matrix(y, yhat).tolist(),
+       "labels": LABELS, "split": "stratified 80/10/10, split-seed 42",
+       "corpus": args.corpus or "fpb-train",
+       "dataset": "financial_phrasebank sentences_50agree"}
+json.dump(res, open(f"{OUT}/result.json", "w"), indent=1)
+print("FINSENSE_RESULT", json.dumps(res), flush=True)
+if not args.smoke:
+    model.save_pretrained(f"{OUT}/model")
+    tok.save_pretrained(f"{OUT}/model")