"""Run the model with ONNX Runtime instead of PyTorch (faster on CPU, especially for short texts).

    pip install onnxruntime
    python onnx_example.py

onnx/model.onnx takes `input_ids` [1, n] and returns `logits` [1, n, tags]. It uses ONNX Runtime's fused
MultiHeadAttention operator, so it runs with ONNX Runtime only. PIIExtractor still does the tokenisation and turns
tags into spans; only the model call changes.
"""
import os

import onnxruntime as ort
import torch

from pii import PIIExtractor


class OnnxModel:
    def __init__(self, file, threads=None):
        so = ort.SessionOptions()
        if threads:
            so.intra_op_num_threads = threads
        self.sess = ort.InferenceSession(file, so, providers=["CPUExecutionProvider"])

    def __call__(self, input_ids):
        return torch.from_numpy(self.sess.run(None, {"input_ids": input_ids.cpu().numpy()})[0])


pii = PIIExtractor(device="cpu")
pii.model = OnnxModel(os.path.join(os.path.dirname(os.path.abspath(__file__)), "onnx", "model.onnx"))

text = "Hi team, I'm Priya Raman (priya.raman@northwind.com). Please call me on +1-860-555-0185."
print(pii.extract(text))
print(pii.redact(text))