inference_mask_messages_onnx.py

#!/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()