wagesj45/vibeguard
wagesj45/vibeguard
Watch
1
0
Fork
You've already forked vibeguard
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Replace toxicity model with multi-label classifier

Browse source
This commit is contained in:
Jordan Wages 2026-08-26 17:00:49 -05:00
commit 1568c787f8
24 changed files with 2581129 additions and 119717 deletions
Download patch file Download diff file Expand all files Collapse all files

38
tools/convert_model.py
View file

@ -1,5 +1,5 @@
#!/usr/bin/env python3
"""Prepare a pinned Hugging Face binary classifier for VibeGuard."""
"""Prepare a pinned Hugging Face multi-label classifier for VibeGuard."""
from __future__ import annotations
import argparse
@ -11,8 +11,10 @@ import tempfile
from pathlib import Path
from typing import Any
DEFAULT_MODEL = "wagesj45/toxic-comment-classifier"
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:
@ -21,9 +23,7 @@ def make_parser() -> argparse.ArgumentParser:
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)
command.add_argument("--max-length", type=int, default=DEFAULT_MAX_LENGTH)
return command
@ -62,14 +62,11 @@ def inspect_source(source: Path) -> dict[str, Any]:
}
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 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:
@ -78,9 +75,10 @@ def prepare(args: argparse.Namespace) -> None:
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
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"
@ -88,12 +86,12 @@ def prepare(args: argparse.Namespace) -> None:
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"):
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": {"toxic": toxic, "nonToxic": non_toxic, "names": names}, "maxLength": max_length, "quantization": "int8-dynamic", "runtime": "onnxruntime-web-wasm", "license": "Apache-2.0"}
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))
@ -146,7 +144,9 @@ def validate(args: argparse.Namespace) -> None:
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")
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))