172 lines
8.7 KiB
Python
172 lines
8.7 KiB
Python
#!/usr/bin/env python3
|
|
"""Prepare a pinned Hugging Face multi-label 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/multilabel-toxic-comment-classifier"
|
|
DEFAULT_OUTPUT = Path("public/models/toxicity")
|
|
DEFAULT_MAX_LENGTH = 512
|
|
REQUIRED_LABELS = ("toxicity", "severe_toxicity", "obscene", "threat", "insult", "identity_attack", "sexual_explicit")
|
|
|
|
|
|
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("--max-length", type=int, default=DEFAULT_MAX_LENGTH)
|
|
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]) -> dict[str, int]:
|
|
indexed = {value.lower().replace("-", "_").replace(" ", "_"): int(index) for index, value in details["id2label"].items()}
|
|
if set(indexed) != set(REQUIRED_LABELS) or len(indexed) != len(REQUIRED_LABELS):
|
|
raise ValueError(f"Expected exactly these labels: {', '.join(REQUIRED_LABELS)}")
|
|
return {label: indexed[label] for label in REQUIRED_LABELS}
|
|
|
|
|
|
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"] != len(REQUIRED_LABELS): raise SystemExit(f"Expected {len(REQUIRED_LABELS)} labels, found {details['numLabels']}")
|
|
names = resolve_labels(details)
|
|
max_length = args.max_length
|
|
if max_length > (details["maxLength"] or max_length): raise SystemExit(f"Requested max length {max_length} exceeds model maximum {details['maxLength']}")
|
|
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", "LICENSE"):
|
|
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": {"names": names}, "scoring": "multi-label-sigmoid", "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" or manifest.get("scoring") != "multi-label-sigmoid": raise SystemExit("Unexpected model quantization, scoring, or license")
|
|
names = manifest.get("labels", {}).get("names", {})
|
|
if set(names) != set(REQUIRED_LABELS) or len(set(names.values())) != len(REQUIRED_LABELS): raise SystemExit("Unexpected model label mapping")
|
|
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()
|