File size: 1,754 Bytes
d1ddb26 | 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 | #!/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()
|