File size: 3,996 Bytes
8153a62 36a680b 8153a62 8d69eed 8153a62 36a680b 8d69eed 36a680b 8d69eed 36a680b 8153a62 8d69eed 8153a62 36a680b 8153a62 36a680b 8153a62 36a680b 8153a62 36a680b 8153a62 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | from __future__ import annotations
import argparse
from collections import Counter, defaultdict
from dataclasses import dataclass
from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline
from mailsort.dataset import DEFAULT_DATASET_ID, DEFAULT_DATASET_REVISION, load_labeled_emails
@dataclass(frozen=True)
class Config:
dataset_id: str
dataset_revision: str | None
model_id_or_path: str
subfolder: str | None
split: str
max_samples: int | None
def _build_text(subject: str, body: str) -> str:
subject = "" if subject is None else str(subject)
body = "" if body is None else str(body)
if subject and body:
return f"Subject: {subject}\n\nBody: {body}"
return subject or body
def _parse_args() -> Config:
p = argparse.ArgumentParser(description="Evaluate a model (local or Hub) against a HF dataset split.")
p.add_argument("--dataset-id", default=DEFAULT_DATASET_ID)
p.add_argument(
"--dataset-revision",
default=None,
help=f"Révision Git du dataset (défaut: {DEFAULT_DATASET_REVISION} = dernière version).",
)
p.add_argument(
"--model",
default="FlowRank/mailSort",
help="Chemin local OU repo Hugging Face (ex. FlowRank/mailSort).",
)
p.add_argument("--subfolder", default="model", help="Sous-dossier des artefacts (ex. model).")
p.add_argument("--split", default="test", help="Which split to evaluate (e.g. test).")
p.add_argument("--max-samples", type=int, default=None, help="Limit evaluation to N samples.")
a = p.parse_args()
return Config(
dataset_id=a.dataset_id,
dataset_revision=a.dataset_revision,
model_id_or_path=a.model,
subfolder=a.subfolder,
split=a.split,
max_samples=a.max_samples,
)
def main() -> int:
cfg = _parse_args()
ds = load_labeled_emails(cfg.dataset_id, revision=cfg.dataset_revision)
if cfg.split not in ds:
raise SystemExit(f"Split '{cfg.split}' not found. Available: {list(ds.keys())}")
rows = ds[cfg.split]
if cfg.max_samples is not None:
rows = rows.select(range(min(cfg.max_samples, len(rows))))
model_kwargs = {}
if cfg.subfolder:
model_kwargs["subfolder"] = cfg.subfolder
tokenizer = AutoTokenizer.from_pretrained(cfg.model_id_or_path, **model_kwargs)
model = AutoModelForSequenceClassification.from_pretrained(cfg.model_id_or_path, **model_kwargs)
clf = pipeline(
"text-classification",
model=model,
tokenizer=tokenizer,
truncation=True,
)
correct = 0
total = 0
per_label = Counter()
per_label_ok = Counter()
confusion = defaultdict(Counter) # true -> pred -> count
for ex in rows:
text = _build_text(ex.get("subject"), ex.get("body"))
true_label = str(ex["label"])
pred = clf(text, top_k=1)[0]["label"]
total += 1
per_label[true_label] += 1
confusion[true_label][pred] += 1
if pred == true_label:
correct += 1
per_label_ok[true_label] += 1
acc = correct / total if total else 0.0
print(f"dataset={cfg.dataset_id} split={cfg.split} samples={total}")
print(f"accuracy={acc:.4f} ({correct}/{total})")
print("\nper-label accuracy:")
for label in sorted(per_label.keys()):
denom = per_label[label]
num = per_label_ok[label]
print(f"- {label}: {num}/{denom} = {num/denom:.4f}")
# print top confusions per label (lightweight)
print("\ncommon confusions (top-2 per true label):")
for true_label in sorted(confusion.keys()):
most = confusion[true_label].most_common(3)
# skip perfect-only rows
if len(most) == 1 and most[0][0] == true_label:
continue
top = ", ".join([f"{pred}:{cnt}" for pred, cnt in most])
print(f"- {true_label}: {top}")
return 0
if __name__ == "__main__":
raise SystemExit(main())
|