""" Convert the training checkpoints of the development repository into the Hugging Face release tree (``huggingface_release/``). Reads (read-only): ../saved_models/distilbert/{kaggle,goemotions,dair}/distilbert_best.pt ../saved_models/bilstm_model/bilstm_attention_best.pt ../saved_models/textcnn_model/textcnn_complete.pt ../saved_models/gcul_{pca,lda}/model_phase{1,3}_best.pt ../datasets/vectorized_data_bilstm/vocab.pkl (bilstm vocabulary) ../github_release/data/splits/label_mapping.csv (authoritative id -> name) Writes: distilbert//{config.json,model.safetensors,tokenizer*} (HF format) bilstm/{config.json,model.safetensors,word2idx.json} textcnn/{config.json,model.safetensors,word2idx.json} gcul_{pca,lda}/{config.json,phase1.safetensors,phase3.safetensors} Design notes ------------ * id2label / label2id always come from ``data/splits/label_mapping.csv`` (single source of truth); the numeric-string ``class_names`` in ``bert_cached_data/*/config.pkl`` are a historical artefact and are never used. * The sklearn ``LabelEncoder`` stored inside ``textcnn_complete.pt`` is **discarded**: its ``classes_`` are identical to the stored ``class_names``, pickles are version sensitive (the local load already warns about sklearn 1.7.0 vs 1.6.1) and the Hub discourages them. * Vocabularies are released as ``word2idx.json`` (string keys; convert with ``int(k)``). * Every conversion is verified: key set, per-tensor shape/dtype and bitwise ``torch.equal`` against the source state_dict; the three DistilBERT models additionally have to pass a ``from_pretrained`` forward pass with output shape ``(batch, num_labels)``. Usage: python _convert.py [--dry-run] [--only distilbert,bilstm,textcnn,gcul] """ from __future__ import annotations import argparse import csv import json import pathlib import pickle import shutil import sys import torch from safetensors.torch import load_file, save_file HERE = pathlib.Path(__file__).resolve().parent PROJ = HERE.parent SM = PROJ / "saved_models" DATA = PROJ / "datasets" LABEL_MAP = HERE.parent / "github_release" / "data" / "splits" / "label_mapping.csv" _LOCAL_BASE = PROJ / "saved_models" / "distilbert-base-uncased" #: prefer the locally cached base checkpoint (offline runs); the released config.json records #: the public Hub name regardless. BASE_MODEL = str(_LOCAL_BASE) if _LOCAL_BASE.exists() else "distilbert-base-uncased" HUB_BASE_NAME = "distilbert-base-uncased" DATASETS = ("kaggle", "goemotions", "dair") REPORT: list[tuple] = [] #: transformers >= 5 renames LayerNorm parameters when saving (``gamma``/``beta`` are the #: canonical Hub names). Loading accepts both spellings, so verification canonicalises them. _RENAMED = {"gamma": "weight", "beta": "bias"} def _canon(key: str) -> str: head, _, tail = key.rpartition(".") return f"{head}.{_RENAMED.get(tail, tail)}" if "LayerNorm" in head else key def pick_metric(metrics: dict, names, default=None): """Return the first available numeric metric, normalised to a 0..1 fraction.""" for n in names: v = metrics.get(n) if isinstance(v, (int, float)): v = float(v) return v / 100.0 if v > 1 else v return default # --------------------------------------------------------------------------- # # helpers # --------------------------------------------------------------------------- # def id2label(dataset: str) -> dict: """Authoritative id -> emotion mapping for one dataset.""" rows = list(csv.DictReader(open(LABEL_MAP, encoding="utf-8"))) m = {int(r["encoded_label"]): r["class_name"] for r in rows if r["dataset"] == dataset} return {i: m[i] for i in sorted(m)} def write_config(path: pathlib.Path, payload: dict) -> None: path.write_text(json.dumps(payload, indent=2, ensure_ascii=False) + "\n", encoding="utf-8") def save_and_verify(sd: dict, out: pathlib.Path) -> int: """Save a state_dict as safetensors and verify it survives bit for bit.""" tensors = {k: v.detach().contiguous().clone() for k, v in sd.items() if torch.is_tensor(v)} assert len(tensors) == len(sd), "non-tensor entries in the state_dict" out.parent.mkdir(parents=True, exist_ok=True) save_file(tensors, str(out), metadata={"format": "pt"}) back = load_file(str(out)) assert set(back) == set(tensors), f"key mismatch in {out.name}" for k, v in tensors.items(): assert back[k].shape == v.shape, f"shape mismatch {k}" assert back[k].dtype == v.dtype, f"dtype mismatch {k}" assert torch.equal(back[k], v), f"value mismatch {k}" REPORT.append((out.parent.name, out.name, len(tensors), ",".join(sorted({str(v.dtype).replace("torch.", "") for v in tensors.values()})), out.stat().st_size)) return len(tensors) def dump_vocab(vocab: dict, out: pathlib.Path) -> int: """Write a token -> id mapping as JSON (string keys).""" out.parent.mkdir(parents=True, exist_ok=True) out.write_text(json.dumps({str(k): int(v) for k, v in vocab.items()}) + "\n", encoding="utf-8") REPORT.append((out.parent.name, out.name, len(vocab), "json", out.stat().st_size)) return len(vocab) def per_dataset_metrics(ds: str) -> dict: """Metrics recorded next to each DistilBERT checkpoint (model-index payload).""" p = SM / "distilbert" / ds / "distilbert_results.pkl" if not p.exists(): return {} with open(p, "rb") as fh: raw = pickle.load(fh) return {k: (float(v) if isinstance(v, (int, float)) else v) for k, v in raw.items()} # --------------------------------------------------------------------------- # # converters # --------------------------------------------------------------------------- # def convert_distilbert(ds: str, dry: bool) -> None: from transformers import AutoTokenizer, DistilBertForSequenceClassification sd = torch.load(SM / "distilbert" / ds / "distilbert_best.pt", map_location="cpu", weights_only=True) num_labels = int(sd["classifier.weight"].shape[0]) mapping = id2label(ds) assert len(mapping) == num_labels, f"label_mapping has {len(mapping)} entries, model {num_labels}" out = HERE / "distilbert" / ds if dry: print(f" [dry-run] distilbert/{ds}: {len(sd)} tensors, num_labels={num_labels}") return # the local base checkpoint carries a classifier head of a different size; the head is # re-initialised here and completely overwritten by ``sd`` right below. model = DistilBertForSequenceClassification.from_pretrained( BASE_MODEL, num_labels=num_labels, id2label={i: n for i, n in mapping.items()}, label2id={n: i for i, n in mapping.items()}, ignore_mismatched_sizes=True) missing, unexpected = model.load_state_dict(sd, strict=False) assert not missing and not unexpected, f"strict load failed: {missing} / {unexpected}" model.eval() out.mkdir(parents=True, exist_ok=True) model.save_pretrained(out, safe_serialization=True) AutoTokenizer.from_pretrained(BASE_MODEL).save_pretrained(out) metrics = per_dataset_metrics(ds) cfg = json.loads((out / "config.json").read_text(encoding="utf-8")) cfg["model_type"] = "distilbert" cfg["_name_or_path"] = HUB_BASE_NAME cfg["architectures"] = ["DistilBertForSequenceClassification"] cfg["datasets"] = [ds] cfg["num_labels"] = num_labels # save_pretrained omits it; see docs/models.md if metrics: acc = pick_metric(metrics, ["test_accuracy", "test_acc", "accuracy", "best_val_acc"]) f1 = pick_metric(metrics, ["macro_f1", "f1", "macro_f1_score"]) mlist = [] if acc is not None: mlist.append({"type": "accuracy", "name": "Test Accuracy", "value": round(acc, 4)}) if f1 is not None: mlist.append({"type": "f1_macro", "name": "Macro F1", "value": round(f1, 4)}) cfg["model-index"] = [{ "name": f"{HUB_BASE_NAME}-gcul-{ds}", "results": [{ "task": {"type": "text-classification", "name": "Emotion classification"}, "dataset": {"type": ds, "name": ds}, "metrics": mlist, }], }] write_config(out / "config.json", cfg) # --- verification of the *released* artefact --------------------------- # reloaded = DistilBertForSequenceClassification.from_pretrained(out) reloaded.eval() ids = torch.randint(0, 1000, (2, 16)) with torch.no_grad(): logits = reloaded(input_ids=ids, attention_mask=torch.ones_like(ids)).logits assert tuple(logits.shape) == (2, num_labels), f"forward pass gave {tuple(logits.shape)}" fwd = {_canon(k): v for k, v in load_file(str(out / "model.safetensors")).items()} assert set(fwd) == set(sd), f"saved keys differ: {sorted(set(fwd) ^ set(sd))}" for k, v in sd.items(): assert torch.equal(fwd[k], v), f"value mismatch {k}" for f in sorted(out.iterdir()): REPORT.append((f"distilbert/{ds}", f.name, len(sd) if f.suffix == ".safetensors" else "-", "pt" if f.suffix == ".safetensors" else "-", f.stat().st_size)) print(f" distilbert/{ds}: num_labels={num_labels}, forward OK (2, {num_labels}), " f"metrics={metrics}") def convert_bilstm(dry: bool) -> None: out = HERE / "bilstm" sd = torch.load(SM / "bilstm_model" / "bilstm_attention_best.pt", map_location="cpu", weights_only=True) if dry: print(f" [dry-run] bilstm: {len(sd)} tensors") return save_and_verify(sd, out / "model.safetensors") vocab = pickle.load(open(DATA / "vectorized_data_bilstm" / "vocab.pkl", "rb")) dump_vocab(vocab, out / "word2idx.json") mapping = id2label("kaggle") write_config(out / "config.json", { "architectures": ["BiLSTMWithAttention"], "model_type": "bilstm_attention", "dataset": "kaggle", "num_labels": len(mapping), "id2label": {str(i): n for i, n in mapping.items()}, "label2id": {n: i for i, n in mapping.items()}, "model_params": {"embedding_dim": 300, "hidden_dim": 128, "num_layers": 2, "dropout": 0.5}, "vocab_size": len(vocab), "max_len": 100, "vocab_file": "word2idx.json", "vocab_note": "JSON object with string keys; load with {int(k): v for k, v in data.items()}", }) print(f" bilstm: {len(sd)} tensors, vocab {len(vocab)}") def convert_textcnn(dry: bool) -> None: out = HERE / "textcnn" ckpt = torch.load(SM / "textcnn_model" / "textcnn_complete.pt", map_location="cpu", weights_only=False) sd = ckpt["model_state_dict"] if dry: print(f" [dry-run] textcnn: {len(sd)} tensors, label_encoder discarded") return save_and_verify(sd, out / "model.safetensors") # ``vocab_info["word2idx"]`` is None in this checkpoint; the vocabulary lives in the # sibling word2idx.pkl written by the preparation notebook. vocab = ckpt["vocab_info"].get("word2idx") if vocab is None: with open(SM / "textcnn_model" / "word2idx.pkl", "rb") as fh: vocab = pickle.load(fh) dump_vocab(vocab, out / "word2idx.json") mapping = id2label("kaggle") assert list(mapping.values()) == list(ckpt["class_names"]), "class_names disagree with label_mapping" write_config(out / "config.json", { "architectures": ["TextCNN"], "model_type": "textcnn", "dataset": "kaggle", "num_labels": len(mapping), "id2label": {str(i): n for i, n in mapping.items()}, "label2id": {n: i for i, n in mapping.items()}, "class_names": ckpt["class_names"], "model_params": ckpt["model_config"], "training": ckpt.get("training_config", {}), "test_accuracy": ckpt.get("test_accuracy"), "vocab_size": ckpt["vocab_info"]["vocab_size"], "max_len": 100, "vocab_file": "word2idx.json", "vocab_note": "JSON object with string keys; load with {int(k): v for k, v in data.items()}", "label_encoder": "discarded - identical to class_names (see docs/models.md)", }) print(f" textcnn: {len(sd)} tensors, vocab {len(vocab)} (label_encoder discarded)") def convert_gcul(which: str, dry: bool) -> None: out = HERE / which sds = {ph: torch.load(SM / which / f"model_{ph}_best.pt", map_location="cpu", weights_only=True) for ph in ("phase1", "phase3")} if dry: print(f" [dry-run] {which}: " + ", ".join(f"{k}={len(v)}" for k, v in sds.items())) return for ph, sd in sds.items(): save_and_verify(sd, out / f"{ph}.safetensors") mapping = id2label("kaggle") write_config(out / "config.json", { "architectures": ["BiLSTMWithBERTInput"], "model_type": "gcul", "reduction": which.split("_")[1], "dataset": "kaggle", "num_labels": len(mapping), "id2label": {str(i): n for i, n in mapping.items()}, "label2id": {n: i for i, n in mapping.items()}, "uncertain_id": 11, "n_clusters": 12, "n_components": 10, "model_params": {"input_dim": 768, "hidden_dim": 256, "num_layers": 2, "dropout": 0.55}, "phase_files": ["phase1.safetensors", "phase3.safetensors"], "phase_note": "phase1 = K-class model, phase3 = (K+1)-class model incl. the uncertain label", }) print(f" {which}: " + ", ".join(f"{k}={len(v)}" for k, v in sds.items())) # --------------------------------------------------------------------------- # def main(argv=None) -> int: ap = argparse.ArgumentParser(description="Build the Hugging Face release tree.") ap.add_argument("--dry-run", action="store_true") ap.add_argument("--only", default=None, help="comma separated subset of distilbert,bilstm,textcnn,gcul") args = ap.parse_args(argv) todo = ("distilbert", "bilstm", "textcnn", "gcul") if not args.only else tuple( x.strip() for x in args.only.split(",")) if not LABEL_MAP.exists(): print(f"label mapping not found: {LABEL_MAP}", file=sys.stderr) return 2 for what in todo: if what == "distilbert": for ds in DATASETS: convert_distilbert(ds, args.dry_run) elif what == "bilstm": convert_bilstm(args.dry_run) elif what == "textcnn": convert_textcnn(args.dry_run) elif what == "gcul": for which in ("gcul_pca", "gcul_lda"): convert_gcul(which, args.dry_run) else: raise SystemExit(f"unknown target {what!r}") if args.dry_run: return 0 print("\n### verified artefacts (key/shape/dtype/bitwise checked)") print(f"{'dir':<22}{'file':<26}{'tensors':>8}{'dtypes':>18}{'KB':>10}") for d, f, n, dt, size in REPORT: print(f"{d:<22}{f:<26}{str(n):>8}{dt:>18}{size/1024:>10.1f}") # never ship bytecode caches to the Hub (upload_folder would include them) shutil.rmtree(HERE / "__pycache__", ignore_errors=True) total = sum(r[4] for r in REPORT if isinstance(r[4], int)) print(f"\nhuggingface_release/ total: {total/1024/1024:.1f} MB in {len(REPORT)} files") return 0 if __name__ == "__main__": raise SystemExit(main())