#!/usr/bin/env python3 """Prepare a pinned Hugging Face binary classifier for VibeGuard.""" from __future__ import annotations import argparse import json import shutil import subprocess import sys import tempfile from pathlib import Path from typing import Any DEFAULT_MODEL = "wagesj45/toxic-comment-classifier" DEFAULT_OUTPUT = Path("public/models/toxicity") def make_parser() -> argparse.ArgumentParser: command = argparse.ArgumentParser(description=__doc__) command.add_argument("action", choices=("inspect", "prepare", "validate")) command.add_argument("--model", default=DEFAULT_MODEL) command.add_argument("--revision", help="Immutable HF commit; required by prepare") command.add_argument("--output", type=Path, default=DEFAULT_OUTPUT) command.add_argument("--toxic-index", type=int) command.add_argument("--non-toxic-index", type=int) command.add_argument("--max-length", type=int) return command def load_source(model: str, revision: str | None) -> tuple[Path, dict[str, Any]]: try: from huggingface_hub import model_info, snapshot_download except ImportError as error: raise SystemExit("Install tools with: python3 -m pip install -r tools/requirements-model.txt") from error kwargs: dict[str, Any] = {"repo_id": model} if revision: kwargs["revision"] = revision source = Path(snapshot_download(**kwargs)) info = None try: info = model_info(model, revision=revision) except Exception as error: # A pinned, complete local snapshot is still usable offline. print(f"Warning: Hugging Face metadata unavailable; using local snapshot metadata: {error}", file=sys.stderr) card = getattr(info, "cardData", None) or {} license_id = card.get("license") if not license_id and (source / "LICENSE").exists(): license_text = (source / "LICENSE").read_text(errors="replace").lower() if "apache license" in license_text and "version 2.0" in license_text: license_id = "apache-2.0" return source, {"source": model, "revision": revision or getattr(info, "sha", "unknown"), "license": license_id} def inspect_source(source: Path) -> dict[str, Any]: config = json.loads((source / "config.json").read_text()) labels = {str(key): str(value) for key, value in (config.get("id2label") or {}).items()} return { "architecture": (config.get("architectures") or [config.get("model_type", "unknown")])[0], "modelType": config.get("model_type"), "numLabels": config.get("num_labels") or len(labels), "id2label": labels, "maxLength": config.get("max_position_embeddings"), "files": sorted(path.name for path in source.iterdir()), } def resolve_labels(details: dict[str, Any], toxic_override: int | None, non_toxic_override: int | None) -> tuple[int, int, dict[str, int]]: names = {str(index): value.lower().replace("-", "_").replace(" ", "_") for index, value in details["id2label"].items()} indexed = {name: int(index) for index, name in names.items()} toxic = toxic_override if toxic_override is not None else next((index for name, index in indexed.items() if name in ("toxic", "toxicity")), None) non_toxic = non_toxic_override if non_toxic_override is not None else next((index for name, index in indexed.items() if name in ("non_toxic", "nontoxic", "clean", "not_toxic")), None) if toxic is None or non_toxic is None or toxic == non_toxic: raise ValueError("Could not resolve distinct labels; pass --toxic-index and --non-toxic-index") return toxic, non_toxic, indexed def prepare(args: argparse.Namespace) -> None: if not args.revision: raise SystemExit("prepare requires --revision for reproducible release artifacts") source, source_metadata = load_source(args.model, args.revision) license_id = str(source_metadata.get("license") or "").lower() if license_id != "apache-2.0": raise SystemExit(f"Expected Apache-2.0 model metadata, found license={source_metadata.get('license')!r}") details = inspect_source(source) if details["numLabels"] != 2: raise SystemExit(f"Expected a binary classifier, found {details['numLabels']}") toxic, non_toxic, names = resolve_labels(details, args.toxic_index, args.non_toxic_index) max_length = args.max_length or details["maxLength"] or 512 with tempfile.TemporaryDirectory(prefix="vibeguard-model-") as temporary: export_dir = Path(temporary) / "onnx" float_model = export_dir / "model.onnx" export_onnx(source, float_model, max_length) quantized = export_dir / "model_quantized.onnx" quantize(float_model, quantized) args.output.mkdir(parents=True, exist_ok=True) for name in ("tokenizer.json", "tokenizer_config.json", "special_tokens_map.json", "vocab.txt", "merges.txt"): if (source / name).exists(): shutil.copy2(source / name, args.output / name) shutil.copy2(source / "config.json", args.output / "config.json") (args.output / "onnx").mkdir(exist_ok=True) shutil.copy2(quantized, args.output / "onnx" / "model_quantized.onnx") manifest = {"source": args.model, "revision": source_metadata["revision"], "architecture": details["architecture"], "labels": {"toxic": toxic, "nonToxic": non_toxic, "names": names}, "maxLength": max_length, "quantization": "int8-dynamic", "runtime": "onnxruntime-web-wasm", "license": "Apache-2.0"} (args.output / "model-manifest.json").write_text(json.dumps(manifest, indent=2) + "\n") print(json.dumps(manifest, indent=2)) def quantize(source: Path, destination: Path) -> None: try: from onnxruntime.quantization import QuantType, quantize_dynamic except ImportError as error: raise SystemExit("Install tools with: python3 -m pip install -r tools/requirements-model.txt") from error quantize_dynamic(str(source), str(destination), weight_type=QuantType.QInt8, per_channel=True, reduce_range=True) def export_onnx(source: Path, destination: Path, max_length: int) -> None: try: import torch from transformers import AutoModelForSequenceClassification, AutoTokenizer except ImportError as error: raise SystemExit("Install model tooling with: python3 -m pip install -r tools/requirements-model.txt") from error tokenizer = AutoTokenizer.from_pretrained(str(source), local_files_only=True) model = AutoModelForSequenceClassification.from_pretrained(str(source), local_files_only=True) model.eval() encoded = tokenizer("VibeGuard export calibration text.", return_tensors="pt", truncation=True, max_length=min(max_length, 32), padding="max_length") class ClassifierWrapper(torch.nn.Module): def __init__(self, wrapped: torch.nn.Module) -> None: super().__init__() self.wrapped = wrapped def forward(self, input_ids: Any, attention_mask: Any) -> Any: return self.wrapped(input_ids=input_ids, attention_mask=attention_mask).logits destination.parent.mkdir(parents=True, exist_ok=True) with torch.no_grad(): torch.onnx.export( ClassifierWrapper(model), (encoded["input_ids"], encoded["attention_mask"]), str(destination), input_names=["input_ids", "attention_mask"], output_names=["logits"], dynamic_axes={"input_ids": {0: "batch", 1: "sequence"}, "attention_mask": {0: "batch", 1: "sequence"}, "logits": {0: "batch"}}, opset_version=17, do_constant_folding=True, dynamo=False, ) def validate(args: argparse.Namespace) -> None: manifest_path = args.output / "model-manifest.json" model_path = args.output / "onnx" / "model_quantized.onnx" if not manifest_path.is_file() or not model_path.is_file(): raise SystemExit(f"Missing generated model files under {args.output}") manifest = json.loads(manifest_path.read_text()) if manifest.get("quantization") != "int8-dynamic" or str(manifest.get("license")).lower() != "apache-2.0": raise SystemExit("Unexpected model quantization or license") try: import onnx onnx.checker.check_model(str(model_path)) except ImportError as error: raise SystemExit("Install tools with: python3 -m pip install -r tools/requirements-model.txt") from error print(f"Validated {model_path} ({model_path.stat().st_size / 1024 / 1024:.1f} MiB)") def run(command: list[str]) -> None: print("+", " ".join(command)) subprocess.run(command, check=True) def main() -> None: args = make_parser().parse_args() if args.action == "inspect": source, metadata = load_source(args.model, args.revision) print(json.dumps({**metadata, **inspect_source(source)}, indent=2)) elif args.action == "prepare": prepare(args) else: validate(args) if __name__ == "__main__": main()