v0.2: 4M, made from cbx-pii-extractor-17M v0.2 (16K vocab, compressed word table, 4 layers, 68M teacher)
c57104c verified | """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)) | |