| """Export a fine-tuned classifier for the Hugging Face model repo and the Gradio Space. |
| |
| Subcommand ``run --checkpoint <finetune best dir>``: |
| |
| 1. load ``EncoderForClassification`` strictly (``num_labels`` from ``config.json["model"]``), |
| cast to ``--dtype`` (bf16 by default: ~337M params x 2 B = ~675 MB, half the fp32 size, |
| and the Space runs the encoder under bf16 anyway; the classifier has no tied weights, |
| so a plain ``state_dict`` is safetensors-safe) and write ``model.safetensors`` with a |
| text-only metadata header (format, dtype, architecture, config JSON) so a loader can |
| rebuild the model from the file alone. The file is re-read and checked key-for-key. |
| 2. ``tokenizer.json`` (copied) and ``config.json`` = model config + label names |
| (``{"0": "Legitimate", "1": "Phishing", "2": "Spam"}`` for a 3-class head, binary names |
| otherwise) + ``binary_rollup`` (which class indices sum to "unwanted") + max_seq_len + |
| export dtype + source checkpoint + the evaluation metrics (reports/metrics.json for the |
| TREC-06 binary table and reports/metrics_test3.json for the three-way table, when |
| present) + the fine-tune dials and init descriptor + the pretraining summary (the |
| pretrain run's final/config.json reached through the checkpoint's ``pretrained:<path>`` |
| descriptor or ``--pretrain-config``; always null for a scratch-initialised checkpoint, |
| never guessed from the newest run). |
| 3. ``README.md`` model card (huggingface_hub ModelCard / ModelCardData front matter) with |
| BOTH metric tables, stating prominently that no pretrained weights were used and that |
| the phishing / spam labels follow corpus provenance (Nazario + Nigerian fraud + synthetic |
| lures = phishing; SpamAssassin + Enron-Spam spam folders = spam; TREC / CEAS positives |
| were trained as "unwanted" only). |
| 4. copy ``src/model.py`` -> ``<app-dir>/model.py`` byte-for-byte (sha1 asserted) and the |
| three weight files into ``<app-dir>/`` so ``app.py`` can load locally when they exist |
| and fall back to ``hf_hub_download`` otherwise. ``app.py`` and ``requirements.txt`` belong |
| to another owner and are never written here. |
| |
| Nothing is uploaded; ``upload_to_hf.py`` owns that step. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import logging |
| import shutil |
| import sys |
| from datetime import datetime |
| from pathlib import Path |
| from typing import Any, Sequence |
|
|
| sys.path.insert(0, str(Path(__file__).parent)) |
|
|
| import torch |
| from safetensors import safe_open |
| from safetensors.torch import load_file, save_file |
|
|
| import common |
| from model import EncoderConfig, EncoderForClassification, count_parameters |
|
|
| LOG = logging.getLogger("export") |
|
|
| |
| |
| UNWANTED_NAME = "Unwanted (phishing or spam)" |
| LABELS_3: dict[str, str] = {"0": "Legitimate", "1": "Phishing", "2": "Spam"} |
| LABELS_2: dict[str, str] = {"0": "Legitimate", "1": UNWANTED_NAME} |
| LABEL_PROVENANCE = ("Phishing = Nazario phishing corpus + Nigerian-fraud (advance-fee) corpus + synthetic BITS-style credential/BEC lures; " |
| "Spam = SpamAssassin and Enron-Spam spam folders; the TREC-05 / TREC-07 / CEAS-08 positives never separated spam " |
| "from phishing and were trained as 'unwanted' only (probability mass of phishing + spam), never as either class.") |
| DTYPES: dict[str, torch.dtype] = {"bf16": torch.bfloat16, "fp32": torch.float32} |
| WEIGHT_FILES: tuple[str, ...] = ("model.safetensors", "tokenizer.json", "config.json") |
| PROTECTED_APP_FILES: frozenset[str] = frozenset({"app.py", "requirements.txt"}) |
| DATASET_IDS: tuple[str, str] = ("puyang2025/seven-phishing-email-datasets", "corbt/enron-emails") |
| AURA_DATASET_ID = "kudzaiprichard/aura-phishing-email-corpus" |
| CARD_TAGS: tuple[str, ...] = ("phishing", "email-security", "text-classification", "from-scratch") |
|
|
|
|
| def labels_for(num_labels: int) -> dict[str, str]: |
| if num_labels == 3: |
| return dict(LABELS_3) |
| if num_labels == 2: |
| return dict(LABELS_2) |
| raise ValueError(f"num_labels={num_labels}: expected 2 or 3") |
|
|
|
|
| def binary_rollup_for(num_labels: int) -> dict[str, list[int]]: |
| """Class indices whose probabilities sum to P(unwanted): [1, 2] for the 3-class head, [1] for binary.""" |
| return {"unwanted": list(range(1, num_labels))} |
|
|
|
|
| |
| def load_classifier(ckpt_dir: Path) -> tuple[EncoderForClassification, EncoderConfig, dict]: |
| cfg_json = common.read_json(ckpt_dir / "config.json") |
| mcfg = EncoderConfig.from_dict(cfg_json["model"]) |
| mcfg.grad_checkpointing = False |
| labels_for(mcfg.num_labels) |
| model = EncoderForClassification(mcfg) |
| sd = load_file(str(ckpt_dir / "model.safetensors"), device="cpu") |
| model.load_state_dict(sd, strict=True) |
| model.eval() |
| LOG.info("loaded %s: %d tensors, %s params, num_labels=%d", ckpt_dir, len(sd), f"{count_parameters(model)['total']:,}", mcfg.num_labels) |
| return model, mcfg, cfg_json |
|
|
|
|
| def finetune_init(ckpt_cfg: dict) -> str | None: |
| """finetune.py's init descriptor as stored in the checkpoint config: ``"scratch"`` or |
| ``"pretrained:<path>"`` (top-level ``init`` or ``finetune.init``); None when absent.""" |
| v = ckpt_cfg.get("init") or (ckpt_cfg.get("finetune") or {}).get("init") |
| return str(v) if v else None |
|
|
|
|
| def is_scratch_init(ckpt_cfg: dict) -> bool: |
| return (finetune_init(ckpt_cfg) or "").strip().lower() == "scratch" |
|
|
|
|
| def find_pretrain_summary(ckpt_cfg: dict, explicit: str | None) -> tuple[dict | None, str | None]: |
| """``training`` block of the pretraining run's final/config.json. Provenance has to be |
| exact, so only two sources are accepted: an explicit ``--pretrain-config`` path, or a |
| ``pretrained:<path>`` descriptor written by finetune.py into the checkpoint config. |
| There is deliberately no "newest run" fallback: a classifier fine-tuned from random |
| init (``init: scratch``) has no pretraining run, and a pretrained-init checkpoint |
| whose path no longer resolves gets ``None`` (the card then says so) rather than |
| another run's step counts. Passing ``--pretrain-config`` for a scratch checkpoint is |
| an error, not a hint.""" |
| if is_scratch_init(ckpt_cfg): |
| if explicit: |
| raise ValueError(f"--pretrain-config given but the checkpoint was fine-tuned from scratch (init={finetune_init(ckpt_cfg)!r})") |
| return None, None |
| cands: list[Path] = [] |
| if explicit: |
| cands.append(Path(explicit)) |
|
|
| def walk(o: Any) -> None: |
| if isinstance(o, dict): |
| for v in o.values(): |
| walk(v) |
| elif isinstance(o, list): |
| for v in o: |
| walk(v) |
| elif isinstance(o, str) and o.startswith("pretrained:"): |
| p = Path(o[len("pretrained:"):]) |
| cands.append(p if p.is_absolute() else common.project_root() / p) |
|
|
| walk(ckpt_cfg) |
| for c in cands: |
| for p in (c, c / "config.json", c.parent / "config.json", c / "final" / "config.json"): |
| if p.is_file() and p.name == "config.json": |
| try: |
| j = common.read_json(p) |
| except (OSError, json.JSONDecodeError): |
| continue |
| if "training" in j: |
| return j["training"], str(p) |
| return None, None |
|
|
|
|
| |
| def export_weights(model: EncoderForClassification, mcfg: EncoderConfig, dtype: str, path: Path) -> dict[str, Any]: |
| model.to(DTYPES[dtype]) |
| sd = {k: v.contiguous() for k, v in model.state_dict().items()} |
| meta = {"format": "pt", "dtype": dtype, "architecture": "EncoderForClassification", |
| "config": json.dumps(mcfg.to_dict(), sort_keys=True), "labels": json.dumps(labels_for(mcfg.num_labels)), |
| "exported": datetime.now().isoformat(timespec="seconds")} |
| path.parent.mkdir(parents=True, exist_ok=True) |
| save_file(sd, str(path), metadata=meta) |
| |
| back = load_file(str(path), device="cpu") |
| if back.keys() != sd.keys(): |
| raise RuntimeError("exported safetensors keys differ from the state_dict") |
| bad = [k for k in sd if back[k].shape != sd[k].shape or back[k].dtype != DTYPES[dtype]] |
| if bad: |
| raise RuntimeError(f"exported tensors mismatch: {bad[:5]}") |
| head = back["classifier.out.weight"].shape[0] |
| if head != mcfg.num_labels: |
| raise RuntimeError(f"classifier head has {head} outputs but config num_labels={mcfg.num_labels}") |
| with safe_open(str(path), framework="pt") as f: |
| if f.metadata().get("architecture") != "EncoderForClassification": |
| raise RuntimeError("safetensors metadata missing") |
| size = path.stat().st_size |
| LOG.info("wrote %s: %d tensors, %s, %.1f MB", path, len(sd), dtype, size / 1e6) |
| return {"tensors": len(sd), "bytes": size, "dtype": dtype, "params": count_parameters(model)["total"]} |
|
|
|
|
| def build_config(mcfg: EncoderConfig, ckpt: Path, ckpt_cfg: dict, dtype: str, metrics: dict | None, metrics_path: str | None, |
| metrics3: dict | None, metrics3_path: str | None, training: dict | None, training_path: str | None, |
| metrics_mismatch: Sequence[str] = (), metrics3_mismatch: Sequence[str] = ()) -> dict[str, Any]: |
| scratch = is_scratch_init(ckpt_cfg) |
| m3s = _metrics3_summary(metrics3) if metrics3 else None |
| if scratch and training is not None: |
| raise AssertionError("scratch-initialised checkpoint must not carry a pretraining summary") |
| finetune = {k: v for k, v in ckpt_cfg.items() if k != "model"} |
| k = mcfg.num_labels |
| return { |
| "architecture": "EncoderForClassification", |
| "model": mcfg.to_dict(), |
| "num_labels": k, |
| "labels": labels_for(k), |
| "binary_rollup": binary_rollup_for(k), |
| "unwanted_name": UNWANTED_NAME, |
| "label_provenance": LABEL_PROVENANCE if k == 3 else "class 1 = unwanted mail (phishing or spam; the public corpora do not separate them)", |
| "max_seq_len": mcfg.max_seq_len, |
| "export_dtype": dtype, |
| "source_checkpoint": str(ckpt), |
| "exported": datetime.now().isoformat(timespec="seconds"), |
| "tokenizer": "tokenizer.json", |
| "model_input": 'subject + "\\n\\n" + body', |
| "finetune_checkpoint": finetune or None, |
| "finetune_init": finetune_init(ckpt_cfg), |
| "mlm_pretraining": "none (classifier trained from random init)" if scratch else ("summary below" if training else "yes, summary unavailable at export time"), |
| "metrics": _metrics_summary(metrics) if metrics else None, |
| "metrics_source": metrics_path, |
| "metrics_checkpoint_matches": (not metrics_mismatch) if metrics else None, |
| "metrics_mismatch": list(metrics_mismatch) or None, |
| "metrics_threeway": m3s, |
| |
| "metrics_threeway_source": metrics3_path if m3s is not None else None, |
| "metrics_threeway_reason": None if m3s is not None else (((metrics3 or {}).get("threeway") or {}).get("reason") |
| or ("2-class checkpoint: no phishing-vs-spam head" if k == 2 else None)), |
| "metrics_threeway_checkpoint_matches": (not metrics3_mismatch) if m3s is not None else None, |
| "metrics_threeway_mismatch": list(metrics3_mismatch) or None, |
| "training": training, |
| "training_source": training_path, |
| "from_scratch": True, |
| "external_pretrained_weights_used": "none", |
| } |
|
|
|
|
| def read_splits_manifest(data_processed: Path) -> dict | None: |
| """data_prep's splits_manifest.json, the single source for the dataset row counts quoted |
| in the model card (never hard-coded, so a re-run of data_prep cannot stale the card).""" |
| p = data_processed / "splits_manifest.json" |
| if not p.is_file(): |
| return None |
| try: |
| return common.read_json(p) |
| except (OSError, json.JSONDecodeError) as e: |
| LOG.warning("could not read %s: %s", p, e) |
| return None |
|
|
|
|
| def _metrics_summary(m: dict) -> dict[str, Any]: |
| """Binary roll-up summary of an evaluate.py metrics.json (the TREC-06 table).""" |
| at, tf, lat = m.get("at_threshold", {}), m.get("threshold_free", {}), m.get("latency", {}) |
| r1, r01 = tf.get("recall_at_fpr_1pct") or {}, tf.get("recall_at_fpr_0_1pct") or {} |
| lure = (m.get("credential_lure_slice") or {}).get("metrics") or {} |
| g, c = lat.get("gpu_bf16") or {}, lat.get("cpu_fp32") or {} |
| tw = m.get("threeway") or {} |
| bd = (tw.get("predicted_kind_breakdown") or {}).get("true_positives_flagged") |
| return { |
| "split": m.get("split"), "split_file": m.get("split_file"), "n_rows": (m.get("split_sizes") or {}).get("rows"), |
| "dev_run": m.get("dev_run"), "threshold": m.get("threshold"), "timestamp": m.get("timestamp"), "num_labels": m.get("num_labels"), |
| "positive_class": UNWANTED_NAME, |
| "accuracy": at.get("accuracy"), "precision": at.get("precision"), "recall": at.get("recall"), "f1": at.get("f1"), |
| "specificity": at.get("specificity"), "confusion_matrix": at.get("confusion_matrix"), |
| "roc_auc": tf.get("roc_auc"), "pr_auc": tf.get("pr_auc"), |
| "recall_at_fpr_1pct": r1.get("recall"), "recall_at_fpr_0_1pct": r01.get("recall"), |
| "best_f1_threshold": (tf.get("best_f1") or {}).get("threshold"), |
| "credential_lure_recall": lure.get("recall"), "credential_lure_precision": lure.get("precision"), "credential_lure_n": lure.get("n"), |
| "predicted_kind_of_flagged_true_positives": bd, |
| "latency_ms": {"gpu_bf16_avg": g.get("avg_ms"), "gpu_bf16_p50": g.get("p50_ms"), "cpu_fp32_avg": c.get("avg_ms"), |
| "cpu_fp32_p50": c.get("p50_ms"), "cpu_threads": c.get("threads"), |
| "instructor_benchmark_ms": lat.get("instructor_benchmark_ms")}, |
| } |
|
|
|
|
| def _metrics3_summary(m: dict) -> dict[str, Any] | None: |
| """Three-way summary of an evaluate.py metrics_test3.json; None when the file carries no |
| evaluated three-way block (2-class checkpoint, or a split without label3 ground truth).""" |
| tw = m.get("threeway") or {} |
| if not tw.get("evaluated"): |
| return None |
| at = m.get("at_threshold", {}) |
| pvs = tw.get("phish_vs_spam_on_true_unwanted") or {} |
| ko = pvs.get("kind_only") or {} |
| return { |
| "split": m.get("split"), "split_file": m.get("split_file"), "n_rows": (m.get("split_sizes") or {}).get("rows"), |
| "dev_run": m.get("dev_run"), "threshold": m.get("threshold"), "timestamp": m.get("timestamp"), |
| "rows_with_full_label": tw.get("rows_with_full_label"), "label3_counts": tw.get("label3_counts"), |
| "decision_rule": tw.get("decision_rule"), "accuracy": tw.get("accuracy"), "macro_f1": tw.get("macro_f1"), |
| "per_class": tw.get("per_class"), "confusion_matrix": tw.get("confusion_matrix"), |
| "kind_accuracy_on_true_unwanted": ko.get("accuracy"), "phishing_as_spam": ko.get("phishing_as_spam"), "spam_as_phishing": ko.get("spam_as_phishing"), |
| "n_true_unwanted": pvs.get("n_true_unwanted"), |
| "per_source": {s: {"n": v.get("n"), "counts": v.get("counts"), "accuracy": v.get("accuracy")} for s, v in (tw.get("per_source") or {}).items()}, |
| "binary_rollup_on_this_split": {"recall": at.get("recall"), "precision": at.get("precision"), "f1": at.get("f1"), "accuracy": at.get("accuracy")}, |
| } |
|
|
|
|
| |
| def _f(v: Any, nd: int = 4) -> str: |
| return "n/a" if v is None else f"{v:.{nd}f}" |
|
|
|
|
| def _n(v: Any) -> str: |
| """Thousands-separated integer or 'n/a' (manifest fields may be missing).""" |
| return "n/a" if v is None else f"{int(v):,}" |
|
|
|
|
| def _ms(v: Any) -> str: |
| """Latency cell: 'n/a' when that device was not measured (e.g. a CPU-only evaluation).""" |
| return "n/a" if v is None else f"{v} ms" |
|
|
|
|
| def _binary_table_md(ms: dict[str, Any], metrics: dict) -> str: |
| cm = ms.get("confusion_matrix") or {} |
| lat = ms.get("latency_ms") or {} |
| lines = [ |
| f"Evaluated by `src/evaluate.py` on the **{ms['split']}** split (`{Path(ms['split_file'] or '').name}`, {ms['n_rows']:,} rows, " |
| f"threshold {ms['threshold']}){' — DEV RUN on a row subset, not the final number' if ms.get('dev_run') else ''}. " |
| f"Positive class = **{UNWANTED_NAME}**" |
| + (", P(unwanted) = P(phishing) + P(spam); these definitions are identical to the binary-era runs, so the numbers are comparable." if ms.get("num_labels") == 3 else "."), "", |
| "| metric | value |", "|---|---:|", |
| f"| **Recall (unwanted)** | **{_f(ms['recall'])}** ({cm.get('fn', 'n/a')} missed of {(cm.get('fn') or 0) + (cm.get('tp') or 0):,}) |", |
| f"| Precision (unwanted) | {_f(ms['precision'])} |", f"| F1 (unwanted) | {_f(ms['f1'])} |", |
| f"| Accuracy | {_f(ms['accuracy'])} |", f"| Specificity | {_f(ms['specificity'])} |", |
| f"| ROC-AUC | {_f(ms['roc_auc'])} |", f"| PR-AUC | {_f(ms['pr_auc'])} |", |
| f"| Recall at 1% FPR | {_f(ms['recall_at_fpr_1pct'])} |", f"| Recall at 0.1% FPR | {_f(ms['recall_at_fpr_0_1pct'])} |", |
| f"| Credential-lure slice recall / precision (n={ms['credential_lure_n']}) | {_f(ms['credential_lure_recall'])} / {_f(ms['credential_lure_precision'])} |", |
| f"| Latency, batch 1, GPU bf16 (avg / p50) | {_ms(lat.get('gpu_bf16_avg'))} / {_ms(lat.get('gpu_bf16_p50'))} |", |
| f"| Latency, batch 1, CPU fp32, {lat.get('cpu_threads') or 'n/a'} threads (avg / p50) | {_ms(lat.get('cpu_fp32_avg'))} / {_ms(lat.get('cpu_fp32_p50'))} |", |
| f"| Instructor benchmark | {lat.get('instructor_benchmark_ms')} ms per sample |", "", |
| "Recall on the unwanted class is the headline number: in a defensive triage aid a missed phish is the " |
| "costly error, while a false positive costs a human one glance.", |
| ] |
| bd = ms.get("predicted_kind_of_flagged_true_positives") |
| if bd: |
| lines += ["", f"TREC-06 carries no phishing-vs-spam ground truth, so the three-way head can only be *described* here: of the " |
| f"{bd['n']:,} true positives it flagged, the model called {bd['predicted_phishing']:,} phishing and {bd['predicted_spam']:,} spam " |
| "(informational, not a metric)."] |
| return "\n".join(lines) |
|
|
|
|
| def _threeway_table_md(m3: dict[str, Any]) -> str: |
| pc = m3.get("per_class") or {} |
| cmx = m3.get("confusion_matrix") or {} |
| labs = cmx.get("labels") or list(pc) |
| lines = [ |
| f"Evaluated by `src/evaluate.py --split {m3['split']}` on `{Path(m3['split_file'] or '').name}` ({_n(m3['rows_with_full_label'])} rows with a " |
| f"full three-way label, threshold {m3['threshold']}){' — DEV RUN on a row subset, not the final number' if m3.get('dev_run') else ''}. " |
| f"Decision rule `{m3['decision_rule']}`: a row is unwanted iff P(phishing) + P(spam) >= threshold, then the kind is " |
| "argmax(P(phishing), P(spam)), so this table collapses exactly onto the binary roll-up.", "", |
| "| metric | value |", "|---|---:|", |
| f"| **Macro-F1 (3 classes)** | **{_f(m3['macro_f1'])}** |", f"| Accuracy | {_f(m3['accuracy'])} |", |
| ] |
| for lab, v in pc.items(): |
| lines.append(f"| {lab}: precision / recall / F1 (n={v.get('support', 0):,}) | {_f(v.get('precision'))} / {_f(v.get('recall'))} / {_f(v.get('f1'))} |") |
| if m3.get("kind_accuracy_on_true_unwanted") is not None: |
| lines.append(f"| Phishing-vs-spam accuracy on the {_n(m3['n_true_unwanted'])} true unwanted rows (kind decision alone) | " |
| f"{_f(m3['kind_accuracy_on_true_unwanted'])} ({_n(m3['phishing_as_spam'])} phishing read as spam, {_n(m3['spam_as_phishing'])} spam read as phishing) |") |
| br = m3.get("binary_rollup_on_this_split") or {} |
| lines.append(f"| Binary roll-up on this split: recall / precision / F1 | {_f(br.get('recall'))} / {_f(br.get('precision'))} / {_f(br.get('f1'))} |") |
| if cmx.get("matrix"): |
| lines += ["", "Confusion matrix (rows = true label by corpus provenance, columns = predicted):", "", |
| "| true \\ predicted | " + " | ".join(labs) + " |", "|---|" + "---:|" * len(labs)] |
| for i, lab in enumerate(labs): |
| lines.append(f"| **{lab}** | " + " | ".join(f"{v:,}" for v in cmx["matrix"][i]) + " |") |
| ps = m3.get("per_source") or {} |
| if ps: |
| lines += ["", "| source | rows | Legitimate | Phishing | Spam | accuracy |", "|---|---:|---:|---:|---:|---:|"] |
| for s, v in ps.items(): |
| c = v.get("counts") or {} |
| lines.append(f"| {s} | {_n(v.get('n'))} | {_n(c.get('Legitimate'))} | {_n(c.get('Phishing'))} | {_n(c.get('Spam'))} | {_f(v.get('accuracy'))} |") |
| lines += ["", f"**Label provenance.** {LABEL_PROVENANCE} The three-way holdout is drawn from text the tokenizer never saw and the pretraining " |
| "corpus does not contain (Nazario / Nigerian-fraud phishing, SpamAssassin / Enron-Spam spam, CEAS-08 / TREC-07 / SpamAssassin legitimate " |
| "mail from the former validation split). Because the kind labels come from which corpus a message sits in, part of any phishing-vs-spam " |
| "confusion measures corpus style rather than lure content."] |
| return "\n".join(lines) |
|
|
|
|
| def model_card(mcfg: EncoderConfig, cfg: dict, ckpt_cfg: dict, metrics: dict | None, metrics3: dict | None, training: dict | None, dtype: str, |
| params: int, weight_bytes: int, model_repo: str, manifest: dict | None = None, |
| metrics_mismatch: Sequence[str] = (), metrics3_mismatch: Sequence[str] = ()) -> str: |
| from huggingface_hub import EvalResult, ModelCard, ModelCardData |
|
|
| model_name = model_repo.split("/")[-1] |
| k = mcfg.num_labels |
| labels = labels_for(k) |
| ms = _metrics_summary(metrics) if metrics else None |
| m3 = _metrics3_summary(metrics3) if metrics3 else None |
| scratch = is_scratch_init(ckpt_cfg) |
| if scratch and training is not None: |
| raise AssertionError("scratch-initialised checkpoint must not carry a pretraining summary") |
| man = manifest or {} |
| totals = (man.get("counts") or {}).get("totals") or {} |
| tr_rows, va_rows, te = totals.get("finetune_train") or {}, totals.get("finetune_val") or {}, totals.get("test_holdout") or {} |
| te3 = totals.get("test_holdout3") or {} |
| pt_man = man.get("pretrain") or {} |
| decon = (((man.get("drops") or {}).get("decontamination_vs_holdout") or {}).get("by_reason") or {}).get("total") |
| holdout_src = man.get("holdout_test_source") or (cfg or {}).get("data", {}).get("holdout_test_source") or "TREC-06" |
| te_rate = f"{100 * te['base_rate']:.1f}% positive" if te.get("base_rate") is not None else "positive rate per metrics.json" |
| eval_results = [] |
| if ms and ms.get("split") == "test" and not ms.get("dev_run"): |
| eval_results += [ |
| EvalResult(task_type="text-classification", dataset_type=DATASET_IDS[0], dataset_name="TREC-06 holdout (puyang2025), binary roll-up", |
| metric_type=mt, metric_name=mn, metric_value=round(float(ms[kk]), 4)) |
| for kk, mt, mn in (("recall", "recall", "Recall (unwanted)"), ("precision", "precision", "Precision (unwanted)"), |
| ("f1", "f1", "F1 (unwanted)"), ("accuracy", "accuracy", "Accuracy"), ("roc_auc", "roc_auc", "ROC-AUC")) |
| if ms.get(kk) is not None |
| ] |
| if m3 and m3.get("split") == "test3" and not m3.get("dev_run"): |
| eval_results += [ |
| EvalResult(task_type="text-classification", dataset_type=AURA_DATASET_ID, |
| dataset_name="three-way holdout (Nazario + Nigerian Fraud phishing, SpamAssassin / Enron-Spam spam, CEAS-08 / TREC-07 / SpamAssassin ham)", |
| metric_type=mt, metric_name=mn, metric_value=round(float(m3[kk]), 4)) |
| for kk, mt, mn in (("macro_f1", "f1", "Macro-F1 (legitimate / phishing / spam)"), ("accuracy", "accuracy", "Accuracy (3-way)")) |
| if m3.get(kk) is not None |
| ] |
| datasets = list(DATASET_IDS) + ([AURA_DATASET_ID] if k == 3 else []) |
| data = ModelCardData(language="en", license="mit", library_name="pytorch", pipeline_tag="text-classification", |
| tags=list(CARD_TAGS) + (["spam-detection"] if k == 3 else []), datasets=datasets, model_name=model_name, |
| eval_results=eval_results or None) |
| ft = ckpt_cfg.get("finetune") or (cfg or {}).get("finetune") or {} |
| pt = (cfg or {}).get("pretrain") or {} |
| size_md = f"~{weight_bytes / 1e6:,.0f} MB" |
|
|
| |
| def mismatch_md(problems: Sequence[str]) -> str: |
| return (f"> **Warning:** this table was exported with `--allow-metrics-mismatch`; the metrics file does not belong to the " |
| f"exported checkpoint ({'; '.join(problems)}).\n\n") if problems else "" |
|
|
| if ms: |
| metrics_md = mismatch_md(metrics_mismatch) + _binary_table_md(ms, metrics or {}) |
| else: |
| metrics_md = "Not yet evaluated: run `src/evaluate.py run --checkpoint ... --split test` and re-export to fill this table." |
| if k == 3: |
| if m3: |
| metrics3_md = mismatch_md(metrics3_mismatch) + _threeway_table_md(m3) |
| else: |
| metrics3_md = ("Not yet evaluated: run `src/evaluate.py run --checkpoint ... --split test3` (writes `reports/metrics_test3.json`) " |
| "and re-export to fill this table. " + LABEL_PROVENANCE) |
| metrics3_section = f"\n### Three-way holdout (legitimate / phishing / spam)\n\n{metrics3_md}\n" |
| classes_md = (f"to three classes: `0 = Legitimate`, `1 = Phishing`, `2 = Spam`. The binary verdict shown by the demo is the roll-up " |
| f"**{UNWANTED_NAME}** = P(phishing) + P(spam) (`config.json[\"binary_rollup\"]`). {LABEL_PROVENANCE}") |
| else: |
| metrics3_section = "" |
| classes_md = (f"to two classes: `0 = Legitimate`, `1 = {UNWANTED_NAME}` (phishing **or** spam; the public corpora do not separate the two).") |
|
|
| if scratch: |
| pretrain_md = (f"1. **No masked-language-model pretraining.** This checkpoint was fine-tuned from random initialisation " |
| f"(`finetune.init = scratch`); it is the from-scratch control in the project's experiments comparison, so the " |
| f"encoder has seen only the labelled fine-tuning data below.") |
| else: |
| if training: |
| tr_md = "\n".join(f" - {kk}: {v}" for kk, v in training.items()) |
| else: |
| tr_md = (f" - pretraining summary not available at export time: the checkpoint's init descriptor " |
| f"`{finetune_init(ckpt_cfg) or 'pretrained'}` did not resolve to a `final/config.json` and no `--pretrain-config` " |
| f"was given; see RUN_LOG.md for the run record.") |
| pretrain_md = "\n".join([ |
| f"1. **Masked-language-model pretraining from scratch** on the deduplicated Enron email corpus " |
| f"(`{DATASET_IDS[1]}`: {_n(pt_man.get('rows_in'))} rows -> {_n(pt_man.get('rows_out'))} documents after exact/near-duplicate " |
| f"removal and quote stripping), {int(100 * pt.get('mlm_probability', 0.15))}% dynamic masking " |
| f"({'/'.join(str(int(100 * x)) for x in pt.get('mask_replace_probs', (0.8, 0.1, 0.1)))}), sequence packing to " |
| f"{mcfg.max_seq_len} tokens, AdamW (betas {tuple(pt.get('betas', (0.9, 0.98)))}, weight decay {pt.get('weight_decay', 0.01)}), " |
| f"peak LR {pt.get('lr', '1.5e-4')}, bf16 autocast on one RTX 3090 Ti. Run summary (from the pretraining run's " |
| f"`final/config.json`):", |
| tr_md]) |
| ft_md = "\n".join(f" - {kk}: {v}" for kk, v in ft.items()) if ft else " - see config.yaml `finetune`" |
| if k == 3: |
| ft_sources = (f"TREC-05, TREC-07 and CEAS-08 (positives trained as *unwanted* only), Enron-Spam and SpamAssassin (spam folders = `Spam`), " |
| f"the Nazario and Nigerian-fraud corpora of `{AURA_DATASET_ID}` (= `Phishing`) and synthetic BITS-style lures (= `Phishing`)") |
| loss_md = ("Rows with a clean three-way label train with 3-class cross-entropy; TREC / CEAS positives train on the probability mass " |
| "of {phishing, spam} (partial-label loss: -log(P(phishing) + P(spam))), so they shape the legitimate-vs-unwanted decision " |
| "without asserting a kind they were never labelled with.") |
| te3_md = (f" A second holdout, **test_holdout3** ({_n(te3.get('rows'))} rows), holds clean-provenance three-way rows never seen by the " |
| f"tokenizer or the pretraining corpus for the 3x3 confusion matrix.") if te3 else ( |
| " A second holdout, **test_holdout3**, holds clean-provenance three-way rows never seen by the tokenizer or the pretraining " |
| "corpus for the 3x3 confusion matrix.") |
| else: |
| ft_sources = "TREC-05, TREC-07, CEAS-08, Enron-Spam and SpamAssassin rows" |
| loss_md = "" |
| te3_md = "" |
| lat_ms = (ms or {}).get("latency_ms") or {} |
| lat_dev = " — DEV RUN numbers." if (ms or {}).get("dev_run") else "." |
| lat_head = "Measured single-email latency (batch 1, tokenize + forward + softmax, see the metrics table): " |
| if lat_ms.get("gpu_bf16_avg") is not None: |
| latency_md = (lat_head + f"GPU bf16 {lat_ms['gpu_bf16_avg']} ms avg on " |
| f"{(metrics or {}).get('latency', {}).get('gpu_bf16', {}).get('device', 'the evaluation GPU')}" |
| + (f", CPU fp32 ({lat_ms.get('cpu_threads')} threads) {lat_ms['cpu_fp32_avg']} ms avg" if lat_ms.get("cpu_fp32_avg") is not None else "") |
| + lat_dev) |
| elif lat_ms.get("cpu_fp32_avg") is not None: |
| |
| latency_md = (lat_head + f"CPU fp32 ({lat_ms.get('cpu_threads')} threads) {lat_ms['cpu_fp32_avg']} ms avg; GPU latency was not " |
| "measured in this evaluation (CPU-only run)" + lat_dev) |
| else: |
| latency_md = "Single-email latency is reported in the metrics table once `src/evaluate.py` has run." |
| rollup_py = ("p_unwanted = sum(float(probs[i]) for i in cfg[\"binary_rollup\"][\"unwanted\"]) # P(phishing) + P(spam)\n" |
| "print(\"unwanted\" if p_unwanted >= 0.5 else \"legitimate\", {cfg[\"labels\"][str(i)]: round(float(p), 4) for i, p in enumerate(probs)})") |
| files_cfg = ("model config, num_labels, labels, binary_rollup, label provenance, max_seq_len, export dtype, source checkpoint, fine-tune init, " |
| "binary metrics (TREC-06), three-way metrics (test_holdout3), pretraining summary (null for scratch init)") |
|
|
| body = f""" |
| # {model_name}: {'legitimate / phishing / spam' if k == 3 else 'phishing'} email classifier trained from scratch |
| |
| > **No external pretrained weights were used.** The encoder, the tokenizer and every parameter in `model.safetensors` |
| > were trained from random initialisation inside this project, in plain PyTorch, without the `transformers` |
| > library. The only external inputs are the public email corpora listed below. |
| > Initialisation of this checkpoint: `{finetune_init(ckpt_cfg) or 'not recorded'}`. |
| |
| ## What it is |
| |
| A {params:,}-parameter BERT-style **pre-LayerNorm Transformer encoder** ({mcfg.n_layers} layers, d_model {mcfg.d_model}, |
| {mcfg.n_heads} heads, FFN {mcfg.d_ff}, byte-level BPE vocabulary of {mcfg.vocab_size:,}, {mcfg.max_seq_len} positions) with a |
| CLS-pooled classification head that maps an email (`subject + "\\n\\n" + body`) {classes_md} |
| The architecture is defined in `model.py`, which ships with the weights. Weights are stored in **{dtype}** safetensors |
| ({size_md}); the header carries the architecture config and the label names as JSON. |
| |
| ## How it was trained |
| |
| {pretrain_md} |
| 2. **Fine-tuning** on `{DATASET_IDS[0]}` with our own split: {ft_sources} form train ({_n(tr_rows.get('rows'))}) / val ({_n(va_rows.get('rows'))}), grouped by |
| near-duplicate cluster so no template straddles the two; **{holdout_src} ({_n(te.get('rows'))} rows, {te_rate}) is held out |
| entirely** as the unseen binary test set and never touched by the tokenizer, pretraining or fine-tuning ({_n(decon)} pool rows that |
| near-duplicate {holdout_src} templates were removed).{te3_md} Counts come from `data/processed/splits_manifest.json`. |
| {loss_md} |
| Fine-tune dials (from the checkpoint's config.json): |
| {ft_md} |
| |
| ## Metrics |
| |
| ### {holdout_src} binary holdout ({UNWANTED_NAME}) |
| |
| {metrics_md} |
| {metrics3_section} |
| ## Intended use |
| |
| A **defensive triage aid for small-business inboxes**: rank or flag incoming mail for a human to review, surface a |
| probability and the attention-highlighted tokens that drove it, and shorten the time between a phishing wave |
| landing and someone noticing. It is a classroom/final-project model and is not a substitute for a mail gateway. |
| |
| ## Limitations |
| |
| - **Era**: the labelled data is 1998-2008 spam-track mail (TREC, CEAS, SpamAssassin, Enron){', plus Nazario phishing (1995-2022) and Nigerian-fraud mail' if k == 3 else ''}. |
| Modern lures (OAuth consent phishing, QR codes, MFA fatigue, SaaS-brand impersonation) are under-represented; expect lower recall on them. |
| - **Labels follow corpus provenance**: {LABEL_PROVENANCE if k == 3 else 'the public corpora label "unwanted" mail, not phishing specifically; the credential-lure slice in the metrics is the closest proxy for phishing-proper recall.'} |
| {'A phishing-vs-spam confusion therefore partly measures corpus style, and the credential-lure slice remains the closest proxy for phishing-proper recall.' if k == 3 else ''} |
| - **Length confound**: legitimate emails in these corpora average ~2x the tokens of positives. Inputs are truncated to |
| {mcfg.max_seq_len} tokens, and the error analysis checks whether errors follow length, but a short legitimate note can |
| still look suspicious. |
| - **English only**; byte-level BPE will tokenize other scripts but the model has not seen them. |
| - **Headers are not used** (TREC rewrote them); the model sees subject and body text only. Sender reputation, |
| SPF/DKIM and URL reputation are out of scope and should be combined with this score in any real deployment. |
| - Attention highlights are a signal of what the model attended to, not an explanation of why it decided. |
| |
| ## How to load |
| |
| `model.py` from this repository is the only code dependency (plain PyTorch). |
| |
| ```python |
| import json, torch |
| from safetensors.torch import load_file |
| from tokenizers import Tokenizer |
| from model import EncoderConfig, EncoderForClassification # model.py, shipped with the weights |
| |
| cfg = json.load(open("config.json")) |
| model = EncoderForClassification(EncoderConfig.from_dict(cfg["model"])) # num_labels = {k} |
| model.load_state_dict(load_file("model.safetensors"), strict=True) # bf16 weights cast into the fp32 module |
| model.eval() |
| |
| tok = Tokenizer.from_file("tokenizer.json") |
| tok.enable_truncation(max_length=cfg["max_seq_len"]) |
| enc = tok.encode("Your account has been suspended\\n\\nClick here to verify your identity within 24 hours.") |
| ids = torch.tensor([enc.ids]) |
| mask = torch.ones_like(ids, dtype=torch.bool) |
| with torch.inference_mode(): |
| probs = model(ids, mask).float().softmax(-1)[0] |
| {rollup_py} |
| ``` |
| |
| On a GPU wrap the forward in `torch.autocast("cuda", dtype=torch.bfloat16)`. {latency_md} |
| |
| ## Files |
| |
| - `model.safetensors` — {dtype} weights, metadata `format=pt, dtype={dtype}, architecture=EncoderForClassification, config=<json>, labels=<json>` |
| - `config.json` — {files_cfg} |
| - `tokenizer.json` — byte-level BPE, NFKC, case preserved, specials `[PAD]=0 [UNK]=1 [CLS]=2 [SEP]=3 [MASK]=4` |
| - `README.md` — this card |
| |
| ## Provenance |
| |
| Built for the Skillspire AI & Machine Learning course final project (BITS, 2026). Data decisions, measured VRAM / |
| throughput / latency and every challenge hit along the way are recorded in the project's RUN_LOG.md. |
| """ |
| card = ModelCard(f"---\n{data.to_yaml()}\n---\n{body}") |
| return card.content |
|
|
|
|
| |
| def copy_to_app(out_dir: Path, app_dir: Path, src_model_py: Path) -> list[Path]: |
| """model.py byte-for-byte (hash asserted) plus the weight files. app.py / requirements.txt |
| belong to the app owner and are never touched.""" |
| app_dir.mkdir(parents=True, exist_ok=True) |
| written: list[Path] = [] |
| dst = app_dir / "model.py" |
| shutil.copyfile(src_model_py, dst) |
| h_src, h_dst = common.sha1_hex(src_model_py.read_bytes().decode("utf-8")), common.sha1_hex(dst.read_bytes().decode("utf-8")) |
| if h_src != h_dst or src_model_py.read_bytes() != dst.read_bytes(): |
| raise RuntimeError(f"{dst} is not byte-identical to {src_model_py}") |
| LOG.info("copied %s -> %s (sha1 %s)", src_model_py, dst, h_src) |
| written.append(dst) |
| for name in WEIGHT_FILES: |
| if name in PROTECTED_APP_FILES: |
| continue |
| shutil.copyfile(out_dir / name, app_dir / name) |
| written.append(app_dir / name) |
| return written |
|
|
|
|
| def tree(paths: Sequence[Path], root: Path) -> str: |
| lines = [f"{root}/"] |
| for p in sorted(paths): |
| size = p.stat().st_size |
| unit = f"{size / 1e6:,.1f} MB" if size >= 1e6 else f"{size / 1e3:,.1f} kB" |
| lines.append(f" {p.name:<20} {unit:>12}") |
| return "\n".join(lines) |
|
|
|
|
| def check_metrics_match(m: dict | None, path: Path, what: str, ckpt: Path, k: int, expect_split: str, allow: bool) -> list[str]: |
| """Refuse to quote another run's numbers on this checkpoint's card; returns the list of mismatches. |
| |
| evaluate.py writes every metrics file under a fixed name in reports/, so evaluating the |
| cross_entropy / weighted / focal runs in sequence leaves the LAST run's numbers there, and an |
| export of the best run with the default paths would put those numbers (and, since the card |
| wording follows num_labels, a contradictory head description) on this model's card with no |
| warning. A metrics file is accepted only when its ``checkpoint`` resolves to the exported one, |
| its ``num_labels`` equals the head size and its ``split`` is the one the table describes |
| (``test`` for the TREC-06 binary table, ``test3`` for the three-way table). A file without a |
| ``checkpoint`` field cannot be attributed and counts as a mismatch. With |
| ``--allow-metrics-mismatch`` the file is still used, the mismatch is logged, and |
| ``metrics_checkpoint_matches: false`` is recorded in config.json / export_manifest.json.""" |
| if m is None: |
| return [] |
| problems: list[str] = [] |
| mc = m.get("checkpoint") |
| if mc is None or Path(str(mc)).resolve() != ckpt: |
| problems.append(f"checkpoint {mc!r} != {ckpt}") |
| mk = m.get("num_labels") |
| if mk is None or int(mk) != k: |
| problems.append(f"num_labels {mk!r} != {k}") |
| if m.get("split") != expect_split: |
| problems.append(f"split {m.get('split')!r} != {expect_split!r}") |
| if problems: |
| msg = f"{what}: {path} does not belong to the exported checkpoint: " + "; ".join(problems) |
| if not allow: |
| raise ValueError(msg + " (re-run src/evaluate.py on this checkpoint, point --metrics-json/--metrics3-json at its " |
| "metrics files, or pass --allow-metrics-mismatch to export anyway with the mismatch recorded)") |
| LOG.warning("%s; used anyway (--allow-metrics-mismatch), recorded as metrics_checkpoint_matches=false", msg) |
| return problems |
|
|
|
|
| def _read_metrics(path: Path, what: str) -> dict | None: |
| if not path.exists(): |
| LOG.info("%s: %s absent", what, path) |
| return None |
| try: |
| m = common.read_json(path) |
| except (OSError, json.JSONDecodeError) as e: |
| LOG.warning("%s: could not read %s: %s", what, path, e) |
| return None |
| LOG.info("%s: %s (%s split, %s%s)", what, path, m.get("split"), m.get("timestamp"), ", DEV RUN" if m.get("dev_run") else "") |
| return m |
|
|
|
|
| |
| def run(args: argparse.Namespace) -> int: |
| cfg = common.load_config(args.config) |
| global LOG |
| LOG = common.setup_logging("export", cfg["paths"]["logs"]) |
| common.set_seed(int(cfg["seed"])) |
| ckpt = Path(args.checkpoint).resolve() |
| run_name = ckpt.parent.name if ckpt.name in ("best", "last") else ckpt.name |
| out_dir = Path(args.out_dir) if args.out_dir else common.project_root() / "model_export" / run_name |
| app_dir = Path(args.app_dir) if args.app_dir else Path(cfg["paths"]["app_dir"]) |
| out_dir.mkdir(parents=True, exist_ok=True) |
| LOG.info("checkpoint=%s run_name=%s out_dir=%s app_dir=%s dtype=%s", ckpt, run_name, out_dir, app_dir, args.dtype) |
|
|
| model, mcfg, ckpt_cfg = load_classifier(ckpt) |
| k = mcfg.num_labels |
| LOG.info("labels: %s; binary roll-up: %s", labels_for(k), binary_rollup_for(k)) |
| weights = export_weights(model, mcfg, args.dtype, out_dir / "model.safetensors") |
|
|
| tok_src = Path(cfg["paths"]["tokenizer"]) |
| shutil.copyfile(tok_src, out_dir / "tokenizer.json") |
| LOG.info("copied %s (%.1f kB)", tok_src, tok_src.stat().st_size / 1e3) |
|
|
| reports = Path(cfg["paths"]["reports"]) |
| metrics_path = Path(args.metrics_json) if args.metrics_json else reports / "metrics.json" |
| metrics = _read_metrics(metrics_path, "binary metrics (TREC-06)") |
| mismatch = check_metrics_match(metrics, metrics_path, "binary metrics (TREC-06)", ckpt, k, "test", args.allow_metrics_mismatch) |
| if args.metrics3_json: |
| metrics3_path = Path(args.metrics3_json) |
| else: |
| m3cfg = (cfg.get("eval") or {}).get("model_card_metrics3") |
| metrics3_path = (common.project_root() / m3cfg) if m3cfg else reports / "metrics_test3.json" |
| metrics3 = _read_metrics(metrics3_path, "three-way metrics (test_holdout3)") if k == 3 else None |
| mismatch3 = check_metrics_match(metrics3, metrics3_path, "three-way metrics (test_holdout3)", ckpt, k, "test3", args.allow_metrics_mismatch) |
| if k == 2: |
| LOG.info("2-class checkpoint: three-way metrics table skipped") |
| elif metrics3 and not (metrics3.get("threeway") or {}).get("evaluated"): |
| LOG.warning("%s has no evaluated three-way block (%s); card table left as 'not yet evaluated'", metrics3_path, |
| (metrics3.get("threeway") or {}).get("reason")) |
| training, training_path = find_pretrain_summary(ckpt_cfg, args.pretrain_config) |
| LOG.info("fine-tune init: %s; pretraining summary: %s", finetune_init(ckpt_cfg), training_path or ("none (scratch init)" if is_scratch_init(ckpt_cfg) else "absent")) |
| manifest = read_splits_manifest(Path(cfg["paths"]["data_processed"])) |
| LOG.info("splits manifest: %s", "loaded" if manifest else "absent (card quotes n/a for dataset counts)") |
|
|
| export_cfg = build_config(mcfg, ckpt, ckpt_cfg, args.dtype, metrics, str(metrics_path) if metrics else None, |
| metrics3, str(metrics3_path) if metrics3 else None, training, training_path, |
| metrics_mismatch=mismatch, metrics3_mismatch=mismatch3) |
| common.write_json(export_cfg, out_dir / "config.json") |
| card = model_card(mcfg, cfg, ckpt_cfg, metrics, metrics3, training, args.dtype, weights["params"], weights["bytes"], cfg["hf"]["model_repo"], manifest, |
| metrics_mismatch=mismatch, metrics3_mismatch=mismatch3) |
| (out_dir / "README.md").write_text(card, encoding="utf-8") |
| LOG.info("wrote %s (%d chars)", out_dir / "README.md", len(card)) |
|
|
| app_written = copy_to_app(out_dir, app_dir, Path(__file__).resolve().parent / "model.py") |
|
|
| summary = { |
| "timestamp": datetime.now().isoformat(timespec="seconds"), "run_name": run_name, "checkpoint": str(ckpt), |
| "out_dir": str(out_dir), "app_dir": str(app_dir), "dtype": args.dtype, "weights": weights, |
| "num_labels": k, "labels": labels_for(k), "binary_rollup": binary_rollup_for(k), |
| "files": {p.name: p.stat().st_size for p in out_dir.iterdir()}, |
| "app_files": {p.name: p.stat().st_size for p in app_written}, |
| "model_py_sha1": common.sha1_hex((Path(__file__).resolve().parent / "model.py").read_text(encoding="utf-8")), |
| "metrics_included": metrics is not None, "metrics_threeway_included": bool(metrics3 and export_cfg["metrics_threeway"]), |
| "metrics_checkpoint_matches": export_cfg["metrics_checkpoint_matches"], "metrics_mismatch": export_cfg["metrics_mismatch"], |
| "metrics_threeway_checkpoint_matches": export_cfg["metrics_threeway_checkpoint_matches"], "metrics_threeway_mismatch": export_cfg["metrics_threeway_mismatch"], |
| "training_included": training is not None, |
| "finetune_init": finetune_init(ckpt_cfg), "manifest_included": manifest is not None, |
| } |
| common.write_json(summary, out_dir / "export_manifest.json") |
| print("\n" + tree(list(out_dir.iterdir()), out_dir)) |
| print(tree(app_written, app_dir)) |
| print(f"\nmodel.py sha1 {summary['model_py_sha1']} (src == app, asserted); weights {weights['bytes'] / 1e6:,.1f} MB {args.dtype}, " |
| f"{weights['params']:,} params; num_labels {k} {labels_for(k)}; init {finetune_init(ckpt_cfg)}; " |
| f"binary metrics {'included' if metrics else 'absent'}; three-way metrics " |
| f"{'included' if summary['metrics_threeway_included'] else ('n/a (2-class)' if k == 2 else 'absent')}; " |
| f"pretraining summary {'included' if training else ('n/a (scratch)' if is_scratch_init(ckpt_cfg) else 'absent')}") |
| return 0 |
|
|
|
|
| def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: |
| ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) |
| ap.add_argument("--config", default="config.yaml") |
| sub = ap.add_subparsers(dest="cmd", required=True) |
| r = sub.add_parser("run", help="export a fine-tune best/ dir to model_export/<run>/ and app/") |
| r.add_argument("--checkpoint", required=True, help="fine-tune best/ dir with model.safetensors + config.json") |
| r.add_argument("--out-dir", default=None, help="default model_export/<run_name>") |
| r.add_argument("--app-dir", default=None, help="default config.paths.app_dir") |
| r.add_argument("--dtype", choices=sorted(DTYPES), default="bf16") |
| r.add_argument("--metrics-json", default=None, help="binary TREC-06 metrics; default <reports>/metrics.json") |
| r.add_argument("--metrics3-json", default=None, help="three-way holdout metrics; default config eval.model_card_metrics3 or <reports>/metrics_test3.json") |
| r.add_argument("--pretrain-config", default=None, |
| help="pretrain final/config.json (default: the checkpoint's pretrained:<path> descriptor; none for init=scratch)") |
| r.add_argument("--allow-metrics-mismatch", action="store_true", |
| help="export even if a metrics file names another checkpoint / num_labels / split (logged; recorded as " |
| "metrics_checkpoint_matches=false in config.json and export_manifest.json). Default: refuse.") |
| return ap.parse_args(argv) |
|
|
|
|
| def main(argv: Sequence[str] | None = None) -> int: |
| return run(parse_args(argv)) |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
