feat(training): support dataset revision for enrichment branch

Ajoute --dataset-revision à train et eval pour charger
FlowRank/labeled_emails@enrichment (20 labels).

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

.gitignore

@@ -5,6 +5,7 @@ __pycache__/
 
 # training outputs (local only)
 outputs/
+outputs-enrichment/
 outputs_smoke/
 
 # HF caches

src/mailsort/eval.py

@@ -11,6 +11,7 @@ from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipe
 @dataclass(frozen=True)
 class Config:
     dataset_id: str
+    dataset_revision: str | None
     model_id_or_path: str
     subfolder: str | None
     split: str
@@ -28,6 +29,11 @@ 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-revision",
+        default=None,
+        help="Git revision of the dataset repo (branch, tag, or commit).",
+    )
     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("--split", default="test", help="Which split to evaluate (e.g. test).")
@@ -35,6 +41,7 @@ def _parse_args() -> Config:
     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,
@@ -45,7 +52,8 @@ def _parse_args() -> Config:
 def main() -> int:
     cfg = _parse_args()
 
-    ds = load_dataset(cfg.dataset_id)
+    kwargs = {"revision": cfg.dataset_revision} if cfg.dataset_revision else {}
+    ds = load_dataset(cfg.dataset_id, **kwargs)
     if cfg.split not in ds:
         raise SystemExit(f"Split '{cfg.split}' not found. Available: {list(ds.keys())}")
 

src/mailsort/train.py

@@ -18,6 +18,7 @@ from transformers import (
 @dataclass(frozen=True)
 class Config:
     dataset_id: str
+    dataset_revision: str | None
     model_name: str
     hub_model_id: str
     output_dir: str
@@ -41,6 +42,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-revision",
+        default=None,
+        help="Git revision of the dataset repo (branch, tag, or commit).",
+    )
     p.add_argument("--model-name", default="distilbert-base-uncased")
     p.add_argument("--hub-model-id", default="FlowRank/mailSort")
     p.add_argument("--output-dir", default="outputs")
@@ -54,6 +60,7 @@ def _parse_args() -> Config:
     a = p.parse_args()
     return Config(
         dataset_id=a.dataset_id,
+        dataset_revision=a.dataset_revision,
         model_name=a.model_name,
         hub_model_id=a.hub_model_id,
         output_dir=a.output_dir,
@@ -67,8 +74,9 @@ def _parse_args() -> Config:
     )
 
 
-def _load_ds(dataset_id: str, seed: int) -> DatasetDict:
-    ds = load_dataset(dataset_id)
+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
@@ -99,7 +107,7 @@ def _compute_metrics(eval_pred):
 def main() -> int:
     cfg = _parse_args()
 
-    ds = _load_ds(cfg.dataset_id, seed=cfg.seed)
+    ds = _load_ds(cfg.dataset_id, seed=cfg.seed, revision=cfg.dataset_revision)
     train_split = "train" if "train" in ds else list(ds.keys())[0]
     test_split = "test" if "test" in ds else ("validation" if "validation" in ds else None)