Add initial VibeGuard browser extension

tools/convert_model.py

@@ -0,0 +1,172 @@
+#!/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()