File size: 3,128 Bytes
7456e0b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | #!/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()
|