| """Fine-tune the email-risk classifier (HF Pro). |
| |
| Trains distilbert-base-uncased (4 classes) on the blocklist-derived dataset. |
| The dataset is class-imbalanced (DISPOSABLE >> SPAM_TRAP), so cross-entropy is |
| class-weighted and per-class precision is reported - the promotion gate. |
| |
| Usage: |
| python -m app.ml.train --data app/ml/data/email_risk.csv \ |
| --hub-repo <hf-user>/email-risk-distilbert [--eval-only --model <repo>] |
| |
| Requires requirements-ml.txt (torch CPU, transformers, datasets, evaluate). |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import sys |
| from collections import Counter |
| from pathlib import Path |
|
|
| LABEL_NAMES = ["NORMAL", "DISPOSABLE", "ROLE", "SPAM_TRAP"] |
| PROMOTION_PRECISION = 0.98 |
|
|
|
|
| def main() -> int: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--data", default="app/ml/data/email_risk.csv") |
| parser.add_argument("--hub-repo", default="") |
| parser.add_argument("--model", default="distilbert-base-uncased") |
| parser.add_argument("--eval-only", action="store_true") |
| parser.add_argument("--epochs", type=int, default=3) |
| parser.add_argument("--out", default="app/ml/metrics.json") |
| args = parser.parse_args() |
|
|
| import datasets |
| import evaluate |
| import numpy as np |
| import torch |
| from sklearn.metrics import precision_recall_fscore_support |
| from transformers import ( |
| AutoModelForSequenceClassification, |
| AutoTokenizer, |
| Trainer, |
| TrainingArguments, |
| ) |
|
|
| raw = datasets.load_dataset("csv", data_files=args.data)["train"] |
| raw = raw.train_test_split(test_size=0.15, seed=42) |
|
|
| tokenizer = AutoTokenizer.from_pretrained(args.model) |
|
|
| def tok(batch): |
| return tokenizer(batch["text"], padding="max_length", truncation=True, |
| max_length=64) |
|
|
| ds = raw.map(tok, batched=True) |
| labels = ds["train"]["label"] |
| counts = Counter(labels) |
| total = sum(counts.values()) |
| weights = [total / counts.get(i, 1) for i in range(len(LABEL_NAMES))] |
| print(f"class counts: {dict(sorted(counts.items()))} -> weights {weights}") |
|
|
| model = AutoModelForSequenceClassification.from_pretrained( |
| args.model, num_labels=len(LABEL_NAMES), |
| id2label={i: n for i, n in enumerate(LABEL_NAMES)}, |
| label2id={n: i for i, n in enumerate(LABEL_NAMES)}, |
| ) |
|
|
| def compute_metrics(eval_pred): |
| logits, y = eval_pred |
| pred = np.argmax(logits, axis=-1) |
| p, r, f1, _ = precision_recall_fscore_support( |
| y, pred, labels=range(len(LABEL_NAMES)), zero_division=0) |
| return { |
| "accuracy": float((pred == y).mean()), |
| **{f"precision_{LABEL_NAMES[i]}": float(p[i]) for i in range(len(LABEL_NAMES))}, |
| **{f"recall_{LABEL_NAMES[i]}": float(r[i]) for i in range(len(LABEL_NAMES))}, |
| } |
|
|
| class WeightedTrainer(Trainer): |
| def compute_loss(self, model, inputs, return_outputs=False, **kw): |
| y = inputs.pop("labels") |
| outputs = model(**inputs) |
| loss = torch.nn.functional.cross_entropy( |
| outputs.logits, y, weight=torch.tensor(weights, dtype=outputs.logits.dtype, |
| device=outputs.logits.device)) |
| return (loss, outputs) if return_outputs else loss |
|
|
| targs = TrainingArguments( |
| output_dir="app/ml/runs", |
| learning_rate=2e-5, |
| per_device_train_batch_size=64, |
| per_device_eval_batch_size=128, |
| num_train_epochs=args.epochs, |
| weight_decay=0.01, |
| eval_strategy="epoch", |
| save_strategy="epoch", |
| load_best_model_at_end=True, |
| use_cpu=not torch.cuda.is_available(), |
| hub_model_id=args.hub_repo or None, |
| hub_strategy="end", |
| push_to_hub=bool(args.hub_repo), |
| ) |
|
|
| trainer = WeightedTrainer(model=model, args=targs, train_dataset=ds["train"], |
| eval_dataset=ds["test"], compute_metrics=compute_metrics) |
|
|
| if args.eval_only: |
| metrics = trainer.evaluate() |
| else: |
| trainer.train() |
| metrics = trainer.evaluate() |
| if args.hub_repo: |
| trainer.push_to_hub() |
|
|
| disposable_precision = metrics.get("eval_precision_DISPOSABLE", 0.0) |
| result = { |
| "model": args.hub_repo or args.model, |
| "metrics": {k: round(v, 4) for k, v in metrics.items() if k.startswith("eval_")}, |
| "class_weights": weights, |
| "promotion": { |
| "gate": f"precision_DISPOSABLE >= {PROMOTION_PRECISION}", |
| "precision_disposable": round(disposable_precision, 4), |
| "promotion_ready": disposable_precision >= PROMOTION_PRECISION, |
| }, |
| } |
| Path(args.out).write_text(json.dumps(result, indent=2), encoding="utf-8") |
| print(json.dumps(result["promotion"], indent=2)) |
| print(f"metrics -> {args.out}") |
| print("PROMOTION READY" if result["promotion"]["promotion_ready"] else |
| "STAY ADVISORY - precision gate not met") |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
