inference_mask_onnx.py

#!/usr/bin/env python3
import argparse
import json

from irish_core_decoder import repair_irish_core_spans_onnx
from onnx_token_classifier import load_onnx_token_classifier


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"]] + f"[{span['label']}]" + out[span["end"]:]
    return out


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", default=".")
    parser.add_argument("--onnx-file", default="onnx/model_quantized.onnx")
    parser.add_argument("--text", required=True)
    parser.add_argument("--ppsn-min-score", type=float, default=0.55)
    parser.add_argument("--other-min-score", type=float, default=0.50)
    parser.add_argument("--json", action="store_true")
    args = parser.parse_args()

    session, tokenizer, config, onnx_path = load_onnx_token_classifier(args.model, onnx_file=args.onnx_file)
    spans = repair_irish_core_spans_onnx(
        args.text,
        session,
        tokenizer,
        config,
        other_min_score=args.other_min_score,
        ppsn_min_score=args.ppsn_min_score,
    )
    result = {
        "model": args.model,
        "onnx_file": str(onnx_path),
        "masked_text": mask_text(args.text, spans),
        "spans": spans,
        "ppsn_decoder": "word_aligned",
        "general_decoder": "irish_core_label_aware",
        "ppsn_min_score": args.ppsn_min_score,
        "other_min_score": args.other_min_score,
        "backend": "onnx",
    }
    if args.json:
        print(json.dumps(result, indent=2, ensure_ascii=False))
    else:
        print(result["masked_text"])


if __name__ == "__main__":
    main()