phi-redact

pip install onnxruntime tokenizers huggingface_hub numpy

import json
import numpy as np
import onnxruntime as ort
from huggingface_hub import hf_hub_download
from tokenizers import Tokenizer

repo = "CDT5058/phi-redact"
session = ort.InferenceSession(hf_hub_download(repo, "phi-redact-v7-int8.onnx"))
tokenizer = Tokenizer.from_file(hf_hub_download(repo, "tokenizer.json"))
tags = json.load(open(hf_hub_download(repo, "tags.json")))

text = "Patient Maria Garcia, DOB 1985-03-22, MRN 4471829, called 415-555-0189."
enc = tokenizer.encode(text)  # adds [CLS]/[SEP], truncates at 256 tokens
feed = {
    "input_ids": np.array([enc.ids], dtype=np.int64),
    "attention_mask": np.array([enc.attention_mask], dtype=np.int64),
    "token_type_ids": np.array([enc.type_ids], dtype=np.int64),
}
pred = session.run(["logits"], feed)[0][0].argmax(-1)

spans, start = [], None
for (s, e), tag in zip(enc.offsets, (tags[i] for i in pred)):
    if tag == "O" or s == e:
        start = None
        continue
    pos, label = tag[0], tag[2:]
    if pos in "BS":
        start = s
    if pos in "ES" and start is not None:
        spans.append((label, text[start:e]))
        start = None

print(spans)
# [('NAME_PATIENT', 'Maria Garcia'), ('DATE', '1985-03-22'),
#  ('MRN', 'MRN 4471829'), ('PHONE', '415-555-0189')]

from transformers import AutoTokenizer, AutoModelForTokenClassification
import torch

tokenizer = AutoTokenizer.from_pretrained("CDT5058/phi-redact")
model = AutoModelForTokenClassification.from_pretrained("CDT5058/phi-redact").eval()

text = "Patient Maria Garcia, DOB 1985-03-22, MRN 4471829, called 415-555-0189."
enc = tokenizer(text, return_tensors="pt", truncation=True, max_length=256,
                return_offsets_mapping=True)
offsets = enc.pop("offset_mapping")[0].tolist()

with torch.no_grad():
    pred = model(**enc).logits.argmax(-1)[0].tolist()

tags = [model.config.id2label[i] for i in pred]
# Decode BIOES tags to character spans exactly as in the ONNX example.