feat: charge labeled_emails@main par défaut et finalise model card v2

- Ajoute mailsort.dataset (révision main = dernière version)
- eval utilise FlowRank/mailSort/model par défaut
- README model card HF (20 labels, accuracy 96%)

Co-authored-by: Hilarion Lefuneste <hilarionlefuneste@tutamail.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

README.md

@@ -6,118 +6,103 @@ tags:
   - text-classification
 language:
   - en
+datasets:
+  - FlowRank/labeled_emails
+base_model: distilbert-base-uncased
+license: mit
+metrics:
+  - accuracy
+model-index:
+  - name: mailSort
+    results:
+      - task:
+          type: text-classification
+        dataset:
+          name: FlowRank/labeled_emails
+          type: FlowRank/labeled_emails
+        metrics:
+          - type: accuracy
+            value: 0.9601
+widget:
+  - text: "Subject: Insurance claim\n\nBody: Hello, I need to update my policy urgently before the deadline."
 ---
 
 ## mailSort
 
-Repo minimal (Python) pour **entraîner, évaluer et publier** un modèle de classification multi-classes d’e-mails à partir du dataset Hugging Face [`FlowRank/labeled_emails`](https://huggingface.co/datasets/FlowRank/labeled_emails).
+Classifieur d’e-mails **20 classes** (catégorie + urgence), basé sur DistilBERT.
 
-**v2** (branche `main`, tag `v2.0`) : **20 classes** (catégorie + urgence), entraîné sur `FlowRank/labeled_emails@main`.  
-**v1** (tag `v1.0`) : 10 classes — `revision="v1.0"` pour le dataset et le modèle.
+Entraîné sur [`FlowRank/labeled_emails`](https://huggingface.co/datasets/FlowRank/labeled_emails) (révision `main`, ~19,8k exemples).  
+**Accuracy test : 96,01 %** (1 900 / 1 979).
 
-Le script principal est `mailsort.train` (Transformers `Trainer`).
+| Version | Labels | Révision dataset | Révision modèle |
+|---------|--------|------------------|-----------------|
+| **v2** (courante) | 20 | `main` ou `v2.0` | `main` ou `v2.0` |
+| v1 | 10 | `v1.0` | `v1.0` |
 
-### Prérequis
+### Labels (v2)
 
-- Python géré par `uv` (ce repo est prévu pour être lancé avec `uv run`)
-- (Optionnel) un GPU CUDA si tu veux accélérer l’entraînement
-- (Optionnel) un token HF si tu veux publier sur le Hub
+Chaque catégorie est déclinée en `_urgent` / `_not_urgent` : `family`, `finance`, `games`, `human resources`, `medical`, `pets`, `school`, `software engineering`, `sport`, `work/airbus`.
 
-### Installation
-
-`uv` installe/synchronise automatiquement les dépendances la première fois.
+## Inférence
 
-```bash
-uv sync
-```
-
-### Entraîner + évaluer (train + test)
+```python
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline
 
-Par défaut, le script charge le dataset **depuis le Hub** et utilise ses splits `train` et `test`.
-L’évaluation se fait à chaque epoch, puis une évaluation finale est exécutée à la fin.
-Les artefacts (modèle, tokenizer) sont sauvegardés dans `outputs/` (ou le dossier passé via `--output-dir`).
+tok = AutoTokenizer.from_pretrained("FlowRank/mailSort", subfolder="model")
+model = AutoModelForSequenceClassification.from_pretrained("FlowRank/mailSort", subfolder="model")
+clf = pipeline("text-classification", model=model, tokenizer=tok, truncation=True)
 
-```bash
-uv run python -m mailsort.train \
-  --dataset-id FlowRank/labeled_emails \
-  --model-name distilbert-base-uncased \
-  --hub-model-id FlowRank/mailSort \
-  --num-train-epochs 2
+text = "Subject: Insurance claim\n\nBody: I need to update my policy urgently."
+print(clf(text))
 ```
 
-Pour réentraîner sur une révision précise du dataset (ex. avant merge) :
+## Développement local
 
-```bash
-uv run python -m mailsort.train \
-  --dataset-id FlowRank/labeled_emails \
-  --dataset-revision enrichment \
-  --num-train-epochs 2
-```
+Repo Python minimal pour entraîner, évaluer et publier. Le script principal est `mailsort.train`.
 
-### Tester / évaluer uniquement
+Par défaut, le code charge la **dernière version du dataset** (`FlowRank/labeled_emails@main`).
 
-Le script n’a pas (encore) de mode “eval-only”.
-Le **minimum** pour faire uniquement une passe rapide est de mettre `--num-train-epochs 0` (ce qui évite l’entraînement) et de garder la phase `evaluate`.
+### Installation
 
 ```bash
-uv run python -m mailsort.train --num-train-epochs 0
+uv sync
 ```
 
-### Évaluer sur le split `test` du dataset (recommandé)
-
-Après ton entraînement dans `outputs/`, tu peux évaluer proprement sur le **split `test`** de `FlowRank/labeled_emails` :
+### Entraîner + évaluer
 
 ```bash
-uv run python -m mailsort.eval --model outputs --dataset-id FlowRank/labeled_emails --split test
+uv run python -m mailsort.train \
+  --model-name distilbert-base-uncased \
+  --hub-model-id FlowRank/mailSort \
+  --num-train-epochs 2
 ```
 
-(Optionnel) Pour un test rapide :
+Pour épingler une version précise du dataset :
 
 ```bash
-uv run python -m mailsort.eval --model outputs --split test --max-samples 200
+uv run python -m mailsort.train --dataset-revision v1.0
 ```
 
-### Publier sur le Hub (FlowRank/mailSort)
-
-Le push se fait automatiquement si la variable d’environnement `HF_TOKEN` (ou `HUGGINGFACE_HUB_TOKEN`) est définie.
+### Évaluer le modèle publié
 
 ```bash
-export HF_TOKEN="..."
-uv run python -m mailsort.train --hub-model-id FlowRank/mailSort
+uv run python -m mailsort.eval --split test
 ```
 
-### Publier via Git (README à la racine + artefacts dans `model/`)
+Évalue par défaut `FlowRank/mailSort` (sous-dossier `model/`) sur `FlowRank/labeled_emails@main`.
 
-Pour avoir un repo Hugging Face “complet” (doc + poids) tout en gardant une structure propre, on met :
-
-- `README.md` à la racine (documentation + model card)
-- les artefacts (config, poids, tokenizer) dans `model/`
-
-Hugging Face pourra charger le modèle via `subfolder="model"`.
-
-1) Préparer le dossier `model/` à partir de `outputs/` :
+### Publier via Git (README + `model/`)
 
 ```bash
 uv run python -m mailsort.prepare_model --outputs-dir outputs --model-dir model
-```
-
-2) Commit + push vers le repo Hugging Face `FlowRank/mailSort` :
-
-```bash
 git add README.md model
-git commit -m "Add model artifacts under model/ + docs"
+git commit -m "chore(model): update published artifacts"
 git push
 ```
 
-### Inférence (utiliser le modèle publié)
-
-```python
-from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline
+Le push Hub automatique via `Trainer` nécessite `HF_TOKEN` :
 
-tok = AutoTokenizer.from_pretrained("FlowRank/mailSort", subfolder="model")
-model = AutoModelForSequenceClassification.from_pretrained("FlowRank/mailSort", subfolder="model")
-clf = pipeline("text-classification", model=model, tokenizer=tok, truncation=True)
-text = "Subject: Insurance claim\n\nBody: Hello, I need to update my policy..."
-print(clf(text))
+```bash
+export HF_TOKEN="..."
+uv run python -m mailsort.train --hub-model-id FlowRank/mailSort
 ```
-

src/mailsort/dataset.py

@@ -0,0 +1,21 @@
+from __future__ import annotations
+
+from datasets import DatasetDict, load_dataset
+
+DEFAULT_DATASET_ID = "FlowRank/labeled_emails"
+DEFAULT_DATASET_REVISION = "main"
+
+
+def load_labeled_emails(
+    dataset_id: str = DEFAULT_DATASET_ID,
+    revision: str | None = None,
+    *,
+    seed: int = 42,
+) -> DatasetDict:
+    """Charge le dataset labeled_emails (révision `main` = dernière version publiée)."""
+    ds = load_dataset(dataset_id, revision=revision or DEFAULT_DATASET_REVISION)
+    if "train" in ds and "test" in ds:
+        return ds
+    if "train" in ds and "test" not in ds:
+        return ds["train"].train_test_split(test_size=0.1, seed=seed)
+    return ds

src/mailsort/eval.py

@@ -4,9 +4,10 @@ import argparse
 from collections import Counter, defaultdict
 from dataclasses import dataclass
 
-from datasets import load_dataset
 from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline
 
+from mailsort.dataset import DEFAULT_DATASET_ID, DEFAULT_DATASET_REVISION, load_labeled_emails
+
 
 @dataclass(frozen=True)
 class Config:
@@ -28,14 +29,18 @@ def _build_text(subject: str, body: str) -> str:
 
 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="FlowRank/labeled_emails")
+    p.add_argument("--dataset-id", default=DEFAULT_DATASET_ID)
     p.add_argument(
         "--dataset-revision",
         default=None,
-        help="Git revision of the dataset repo (branch, tag, or commit).",
+        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("--model", default="outputs", help="Local path OR Hugging Face repo id (e.g. FlowRank/mailSort).")
-    p.add_argument("--subfolder", default=None, help="Optional subfolder (e.g. model).")
+    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()
@@ -52,8 +57,7 @@ def _parse_args() -> Config:
 def main() -> int:
     cfg = _parse_args()
 
-    kwargs = {"revision": cfg.dataset_revision} if cfg.dataset_revision else {}
-    ds = load_dataset(cfg.dataset_id, **kwargs)
+    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())}")
 
@@ -61,12 +65,12 @@ def main() -> int:
     if cfg.max_samples is not None:
         rows = rows.select(range(min(cfg.max_samples, len(rows))))
 
-    kwargs = {}
+    model_kwargs = {}
     if cfg.subfolder:
-        kwargs["subfolder"] = cfg.subfolder
+        model_kwargs["subfolder"] = cfg.subfolder
 
-    tokenizer = AutoTokenizer.from_pretrained(cfg.model_id_or_path, **kwargs)
-    model = AutoModelForSequenceClassification.from_pretrained(cfg.model_id_or_path, **kwargs)
+    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,

src/mailsort/train.py

@@ -5,7 +5,7 @@ import os
 from dataclasses import dataclass
 
 import numpy as np
-from datasets import DatasetDict, load_dataset
+from datasets import DatasetDict
 from transformers import (
     AutoModelForSequenceClassification,
     AutoTokenizer,
@@ -14,6 +14,8 @@ from transformers import (
     TrainingArguments,
 )
 
+from mailsort.dataset import DEFAULT_DATASET_ID, DEFAULT_DATASET_REVISION, load_labeled_emails
+
 
 @dataclass(frozen=True)
 class Config:
@@ -41,11 +43,11 @@ def _build_text(subject: str, body: str) -> str:
 
 def _parse_args() -> Config:
     p = argparse.ArgumentParser(description="Train & push email classifier to Hugging Face Hub.")
-    p.add_argument("--dataset-id", default="FlowRank/labeled_emails")
+    p.add_argument("--dataset-id", default=DEFAULT_DATASET_ID)
     p.add_argument(
         "--dataset-revision",
         default=None,
-        help="Git revision of the dataset repo (branch, tag, or commit).",
+        help=f"Révision Git du dataset (défaut: {DEFAULT_DATASET_REVISION} = dernière version).",
     )
     p.add_argument("--model-name", default="distilbert-base-uncased")
     p.add_argument("--hub-model-id", default="FlowRank/mailSort")
@@ -75,15 +77,7 @@ def _parse_args() -> Config:
 
 
 def _load_ds(dataset_id: str, seed: int, revision: str | None = None) -> DatasetDict:
-    kwargs = {"revision": revision} if revision else {}
-    ds = load_dataset(dataset_id, **kwargs)
-    if "train" in ds and "test" in ds:
-        return ds  # already split
-    # fallback: split if only a single split exists
-    if "train" in ds and "test" not in ds:
-        return ds["train"].train_test_split(test_size=0.1, seed=seed)
-    # if weird structure, just return as-is and let Trainer fail loudly
-    return ds
+    return load_labeled_emails(dataset_id=dataset_id, revision=revision, seed=seed)
 
 
 def _prepare(ds: DatasetDict, tokenizer: AutoTokenizer, label2id: dict[str, int], max_length: int) -> DatasetDict: