Replace toxicity model with multi-label classifier

tests/cache.test.ts

@@ -6,8 +6,7 @@ const result = (id: string): InferenceResult => ({
   requestId: id,
   id,
   textHash: id,
-  label: "toxic",
-  probability: .9,
+  scores: { toxicity: .9, severe_toxicity: 0, obscene: 0, threat: 0, insult: 0, identity_attack: 0, sexual_explicit: 0 },
   navigationId: "nav"
 });
 

tests/filter.test.ts

@@ -1,14 +1,16 @@
 import { afterEach, describe, expect, it } from "vitest";
-import { applyResult, removeOrphanedPlaceholders, restore } from "../src/content/filter";
+import { applyResult, removeOrphanedPlaceholders, restore, selectFilterLabel } from "../src/content/filter";
 import { DEFAULT_SETTINGS } from "../src/shared/types";
 
+const scores = { toxicity: .91, severe_toxicity: .1, obscene: .3, threat: .2, insult: .8, identity_attack: .1, sexual_explicit: .1 };
+
 describe("content filtering", () => {
   afterEach(() => document.body.replaceChildren());
 
   it("collapses toxic content and restores it", () => {
     const element = document.createElement("article");
     document.body.append(element);
-    applyResult(element, { requestId: "1", id: "1", textHash: "x", label: "toxic", probability: .91, navigationId: "n" }, DEFAULT_SETTINGS);
+    applyResult(element, { requestId: "1", id: "1", textHash: "x", scores, navigationId: "n" }, DEFAULT_SETTINGS);
     expect(element.style.display).toBe("none");
     expect(document.querySelector("[data-vibeguard-placeholder]")).not.toBeNull();
     restore(element);
@@ -19,13 +21,19 @@ describe("content filtering", () => {
   it("preserves the original display value when filtering is reconciled repeatedly", () => {
     const element = document.createElement("article");
     document.body.append(element);
-    const result = { requestId: "1", id: "1", textHash: "x", label: "toxic" as const, probability: .91, navigationId: "n" };
+    const result = { requestId: "1", id: "1", textHash: "x", scores, navigationId: "n" };
     applyResult(element, result, DEFAULT_SETTINGS);
     applyResult(element, result, DEFAULT_SETTINGS);
     restore(element);
     expect(element.style.display).toBe("");
   });
 
+  it("uses the highest enabled category above the threshold", () => {
+    const result = { requestId: "1", id: "1", textHash: "x", scores, navigationId: "n" };
+    expect(selectFilterLabel(result, { ...DEFAULT_SETTINGS, enabledLabels: ["insult", "threat"], threshold: .5 })).toEqual({ label: "insult", probability: .8 });
+    expect(selectFilterLabel(result, { ...DEFAULT_SETTINGS, enabledLabels: [], threshold: .5 })).toBeUndefined();
+  });
+
   it("removes placeholders no longer paired with the current post element", () => {
     const stale = document.createElement("div");
     const orphan = document.createElement("div");

tests/inference-worker.test.ts

@@ -23,7 +23,7 @@ class FakeWorker implements InferenceWorker {
   terminate(): void { this.terminated = true; }
 
   reply(): void {
-    const event = new MessageEvent("message", { data: { results: [{ requestId: request.requestId, id: request.id, textHash: "hash", label: "not_toxic", probability: .1, navigationId: request.navigationId }] } });
+    const event = new MessageEvent("message", { data: { results: [{ requestId: request.requestId, id: request.id, textHash: "hash", scores: { toxicity: .1, severe_toxicity: 0, obscene: 0, threat: 0, insult: 0, identity_attack: 0, sexual_explicit: 0 }, navigationId: request.navigationId }] } });
     this.listeners.get("message")?.forEach((listener) => listener(event));
   }
 }

tests/model-metadata.test.ts

@@ -3,25 +3,26 @@ import { validateManifest } from "../src/inference/model-metadata";
 import { resolveOutputIndex } from "../src/inference/classifier";
 
 const manifest = {
-  source: "wagesj45/toxic-comment-classifier",
+  source: "wagesj45/multilabel-toxic-comment-classifier",
   revision: "abc123",
-  architecture: "DistilBertForSequenceClassification",
-  labels: { toxic: 1, nonToxic: 0, names: { toxic: 1, non_toxic: 0 } },
+  architecture: "ModernBertForSequenceClassification",
+  labels: { names: { toxicity: 0, severe_toxicity: 1, obscene: 2, threat: 3, insult: 4, identity_attack: 5, sexual_explicit: 6 } },
+  scoring: "multi-label-sigmoid" as const,
   maxLength: 512,
   quantization: "int8-dynamic",
   runtime: "onnxruntime-web-wasm"
 };
 
 describe("model contract", () => {
-  it("accepts a binary manifest and resolves common output labels", () => {
+  it("accepts a multi-label manifest and resolves output labels", () => {
     expect(validateManifest(manifest)).toEqual(manifest);
     expect(resolveOutputIndex("LABEL_1", manifest)).toBe(1);
-    expect(resolveOutputIndex("not-toxic", manifest)).toBe(0);
-    expect(resolveOutputIndex("toxic", manifest)).toBe(1);
+    expect(resolveOutputIndex("identity-attack", manifest)).toBe(5);
+    expect(resolveOutputIndex("toxicity", manifest)).toBe(0);
   });
 
   it("rejects incomplete metadata", () => {
-    expect(() => validateManifest({ ...manifest, labels: { toxic: 1, nonToxic: 1, names: {} } })).toThrow();
+    expect(() => validateManifest({ ...manifest, labels: { names: { ...manifest.labels.names, threat: 1 } } })).toThrow();
     expect(() => validateManifest({ ...manifest, maxLength: 0 })).toThrow();
     expect(() => validateManifest({ ...manifest, labels: undefined as unknown as typeof manifest.labels })).toThrow();
   });

tests/queue.test.ts

@@ -11,11 +11,11 @@ describe("InferenceQueue", () => {
     let calls = 0;
     const queue = new InferenceQueue(async (requests) => {
       calls += 1;
-      return requests.map((item) => ({ requestId: item.requestId, id: item.id, textHash: "hash", label: "toxic" as const, probability: .9, navigationId: item.navigationId }));
+      return requests.map((item) => ({ requestId: item.requestId, id: item.id, textHash: "hash", scores: { toxicity: .9, severe_toxicity: 0, obscene: 0, threat: 0, insult: 0, identity_attack: 0, sexual_explicit: 0 }, navigationId: item.navigationId }));
     }, undefined, 10, 8, 0);
     const first = await queue.enqueue(request("a", 0));
     const second = await queue.enqueue({ ...request("b", 0), text: "text a" });
-    expect(first.probability).toBe(.9);
+    expect(first.scores.toxicity).toBe(.9);
     expect(second.requestId).toBe("b");
     expect(calls).toBe(1);
     expect(queue.cacheSize).toBe(1);

tests/settings.test.ts

@@ -24,4 +24,9 @@ describe("legacy settings", () => {
     expect(mergeSettings({ cpuThreads: 99 })).not.toHaveProperty("cpuThreads");
     expect(mergeSettings({ inferenceDevice: "webgpu" } as never)).not.toHaveProperty("inferenceDevice");
   });
+
+  it("defaults missing categories and removes invalid stored labels", () => {
+    expect(mergeSettings({}).enabledLabels).toEqual(DEFAULT_SETTINGS.enabledLabels);
+    expect(mergeSettings({ enabledLabels: ["threat", "threat", "unknown"] as never }).enabledLabels).toEqual(["threat"]);
+  });
 });