Add training script: SetFit one-vs-rest multi-label classifier on GIZ vulnerability data

train_setfit.py

@@ -0,0 +1,134 @@
+import os
+os.environ.setdefault("USE_TF", "0")
+os.environ.setdefault("USE_TORCH", "1")
+os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
+
+import json
+import time
+
+import numpy as np
+import pandas as pd
+import trackio
+
+def _trackio_safe(fn):
+    def wrapper(*a, **k):
+        try:
+            return fn(*a, **k)
+        except Exception as e:
+            print(f"trackio.{fn.__name__} failed (non-fatal): {e}")
+    return wrapper
+trackio.log = _trackio_safe(trackio.log)
+trackio.init = _trackio_safe(trackio.init)
+trackio.finish = _trackio_safe(trackio.finish)
+
+from datasets import Dataset
+from huggingface_hub import HfApi
+from sklearn.metrics import f1_score
+from sklearn.model_selection import train_test_split
+from setfit import SetFitModel, TrainingArguments, Trainer
+
+SMOKE = os.environ.get("SMOKE", "0") == "1"
+PARQUET_URL = "https://huggingface.co/datasets/GIZ/vulnerability_training_data_full/resolve/refs%2Fconvert%2Fparquet/default/train/0000.parquet"
+LABELS = [
+    "Agricultural communities", "Coastal communities", "Ethnic, racial or other minorities",
+    "Fishery communities", "Informal sector workers", "Members of indigenous and local communities",
+    "Migrants and displaced persons", "Older persons", "Other", "Persons living in poverty",
+    "Persons with disabilities", "Persons with pre-existing health conditions",
+    "Residents of drought-prone regions", "Rural populations", "Sexual minorities (LGBTQI+)",
+    "Urban populations", "Women and other genders",
+]
+REPO_ID = "peter2000/setfit-vulnerability-groups"
+
+def load_data():
+    df = pd.read_parquet(PARQUET_URL)
+    assert len(df) == 475, f"expected 475 rows, got {len(df)}"
+    Y = df[LABELS].values.astype(np.int64)
+    nlab = Y.sum(1)
+    idx_tr, idx_te = train_test_split(
+        np.arange(len(df)), test_size=0.2, random_state=42, stratify=np.minimum(nlab, 3)
+    )
+    return df, Y, np.asarray(idx_tr), np.asarray(idx_te)
+
+def ece(conf, correct, n_bins=15):
+    conf = np.asarray(conf, dtype=np.float64)
+    corr = np.asarray(correct, dtype=np.float64)
+    bins = np.linspace(0.0, 1.0, n_bins + 1)
+    e = 0.0
+    for lo, hi in zip(bins[:-1], bins[1:]):
+        m = (conf > lo) & (conf <= hi)
+        if m.sum() > 0:
+            e += m.mean() * abs(corr[m].mean() - conf[m].mean())
+    return float(e)
+
+def evaluate(Y_true, P_pred, threshold=0.5):
+    pred = (P_pred >= threshold).astype(int)
+    per_label = f1_score(Y_true, pred, average=None, zero_division=0)
+    macro = float(f1_score(Y_true, pred, average="macro", zero_division=0))
+    micro = float(f1_score(Y_true, pred, average="micro", zero_division=0))
+    conf = np.where(pred == 1, P_pred, 1.0 - P_pred)
+    corr = (pred == Y_true).astype(np.float64)
+    return {
+        "macro_f1": macro,
+        "micro_f1": micro,
+        "ece": ece(conf, corr),
+        "subset_accuracy": float(((pred == Y_true).all(axis=1)).mean()),
+        "per_label_f1": {LABELS[i]: round(float(per_label[i]), 4) for i in range(len(LABELS))},
+    }
+
+def main():
+    df, Y, idx_tr, idx_te = load_data()
+    texts = df["text"].tolist()
+    X_tr = [texts[i] for i in idx_tr]
+    X_te = [texts[i] for i in idx_te]
+    Y_tr, Y_te = Y[idx_tr], Y[idx_te]
+    print(f"train={len(X_tr)} test={len(X_te)} labels={len(LABELS)}")
+
+    if SMOKE:
+        X_tr, Y_tr = X_tr[:40], Y_tr[:40]
+        X_te, Y_te = X_te[:10], Y_te[:10]
+
+    model = SetFitModel.from_pretrained(
+        "sentence-transformers/paraphrase-mpnet-base-v2",
+        multi_target_strategy="one-vs-rest",
+        labels=LABELS,
+    )
+    model.model_body.max_seq_length = 256
+
+    train_ds = Dataset.from_dict({
+        "text": X_tr,
+        "label": [list(map(int, y)) for y in Y_tr],
+    })
+
+    args = TrainingArguments(
+        batch_size=16,
+        num_epochs=1,
+        num_iterations=20 if not SMOKE else 2,
+        seed=42,
+    )
+    trainer = Trainer(model=model, args=args, train_dataset=train_ds)
+    trainer.train()
+
+    P_te = np.asarray(model.predict_proba(X_te))
+    metrics = evaluate(Y_te, P_te)
+    print(json.dumps(metrics, indent=2))
+
+    trackio.init(project="vulnerability-multilabel-classifier", space_id="peter2000/vulnerability-multilabel-classifier-trackio")
+    trackio.log({"setfit_macro_f1": metrics["macro_f1"], "setfit_micro_f1": metrics["micro_f1"], "setfit_ece": metrics["ece"]}, step=1)
+    trackio.log({f"setfit_f1/{k}": v for k, v in metrics["per_label_f1"].items()}, step=1)
+
+    trainer.push_to_hub(repo_id=REPO_ID)
+    api = HfApi(token=os.environ.get("HF_TOKEN"))
+    api.upload_file(
+        path_or_fileobj=json.dumps(metrics, indent=2).encode(),
+        path_in_repo="metrics.json",
+        repo_id=REPO_ID,
+        repo_type="model",
+        commit_message="Add evaluation metrics (macro-F1, micro-F1, per-label F1, ECE)",
+    )
+    trackio.finish()
+    print("DONE")
+
+if __name__ == "__main__":
+    t0 = time.time()
+    main()
+    print(f"elapsed {time.time()-t0:.0f}s")