Add role-aware multi-message inference helper

inference_mask_messages_onnx.py

@@ -0,0 +1,91 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+from common import decode_span_matrix, load_onnx_session, run_onnx_span, sigmoid_np
+
+
+def replacement(label: str) -> str:
+    return f"[PII:{label}]"
+
+
+def mask_text(text: str, spans: list[dict]) -> str:
+    out = text
+    for span in sorted(spans, key=lambda item: (item["start"], item["end"]), reverse=True):
+        out = out[: span["start"]] + replacement(span["label"]) + out[span["end"] :]
+    return out
+
+
+def infer_profile(role: str | None) -> str:
+    role_key = (role or "").strip().lower()
+    if role_key == "assistant":
+        return "assistant_public"
+    return "default"
+
+
+def predict(text: str, role: str | None, session, tokenizer, config, min_score: float):
+    encoded = tokenizer(text, return_offsets_mapping=True, return_tensors="np", truncation=True)
+    offsets = [tuple(item) for item in encoded["offset_mapping"][0].tolist()]
+    span_logits = run_onnx_span(session, encoded)
+    span_scores = sigmoid_np(span_logits[0])
+    profile = infer_profile(role)
+    spans = decode_span_matrix(text, offsets, span_scores, config, min_score, profile=profile)
+    for span in spans:
+        span["replacement"] = replacement(span["label"])
+    return profile, spans
+
+
+def load_messages(path: Path) -> list[dict]:
+    raw = path.read_text(encoding="utf-8")
+    if path.suffix.lower() == ".jsonl":
+        return [json.loads(line) for line in raw.splitlines() if line.strip()]
+    data = json.loads(raw)
+    if isinstance(data, dict):
+        messages = data.get("messages")
+        if isinstance(messages, list):
+            return messages
+    if isinstance(data, list):
+        return data
+    raise ValueError("Expected a JSON array, a JSON object with a `messages` array, or JSONL")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", required=True)
+    parser.add_argument("--input-file", required=True)
+    parser.add_argument("--min-score", type=float, default=0.5)
+    parser.add_argument("--json", action="store_true")
+    args = parser.parse_args()
+
+    session, tokenizer, config = load_onnx_session(args.model, onnx_file="model_quantized.onnx", onnx_subfolder="onnx")
+    messages = load_messages(Path(args.input_file))
+    output_messages = []
+    for message in messages:
+        role = message.get("role")
+        text = message.get("text", "")
+        profile, spans = predict(text, role, session, tokenizer, config, args.min_score)
+        output_messages.append(
+            {
+                **message,
+                "profile": profile,
+                "spans": spans,
+                "masked_text": mask_text(text, spans),
+            }
+        )
+    result = {
+        "model": args.model,
+        "backend": "onnx_global_pointer_q8",
+        "messages": output_messages,
+    }
+    if args.json:
+        print(json.dumps(result, indent=2, ensure_ascii=False))
+    else:
+        for message in output_messages:
+            print(f"[{message.get('role', 'unknown')}] {message['masked_text']}")
+
+
+if __name__ == "__main__":
+    main()