"""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  # noqa: E402
from safetensors import safe_open  # noqa: E402
from safetensors.torch import load_file, save_file  # noqa: E402

import common  # noqa: E402
from model import EncoderConfig, EncoderForClassification, count_parameters  # noqa: E402

LOG = logging.getLogger("export")

# Label contract (RUN_LOG entry 6): 0 legitimate, 1 phishing, 2 spam; the binary roll-up "unwanted" is
# the sum over every non-legitimate class. A 2-class head's class 1 IS the roll-up, hence its name.
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"   # Champa et al. mirror: Nazario.csv + Nigerian_Fraud.csv (phishing)
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))}


# ------------------------------------------------------------------------------ load
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)   # raises on an unsupported head size before anything is written
    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


# ----------------------------------------------------------------------------- export
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)
    # Verify by re-reading: same keys, same shapes, requested dtype, metadata intact.
    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   # None when the file has no evaluated three-way block
    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,
        # source only when a table is carried; otherwise say why (2-class head, or a split without phishing/spam ground truth)
        "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,   # informational: no phishing-vs-spam ground truth on TREC-06
        "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")},
    }


# --------------------------------------------------------------------------- model card
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 {}   # the checkpoint's own dials win
    pt = (cfg or {}).get("pretrain") or {}
    size_md = f"~{weight_bytes / 1e6:,.0f} MB"

    # Only reachable with --allow-metrics-mismatch: say on the card itself that the numbers were not produced on this checkpoint.
    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:
        # CPU-only evaluation (e.g. the GPU was busy): evaluate.py HAS run, so say what was measured
        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


# ------------------------------------------------------------------------------- app
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:  # defensive: never overwrite another owner's 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


# -------------------------------------------------------------------------------- main
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())