| 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) |
|
|
| 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("\ncommon confusions (top-2 per true label):") |
| for true_label in sorted(confusion.keys()): |
| most = confusion[true_label].most_common(3) |
| |
| 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()) |
|
|
|
|
