File size: 3,773 Bytes
85b0a00 | 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 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 | #!/usr/bin/env python3
from __future__ import annotations
import argparse
import json
import os
os.environ.setdefault("TRANSFORMERS_NO_TF", "1")
os.environ.setdefault("TRANSFORMERS_NO_FLAX", "1")
os.environ.setdefault("TRANSFORMERS_NO_TORCHVISION", "1")
os.environ["USE_TF"] = "0"
os.environ["USE_FLAX"] = "0"
os.environ["USE_TORCH"] = "1"
import torch
from transformers import AutoConfig
from common import (
boundary_label_thresholds_from_config,
decode_token_presence_segments,
label_max_span_tokens_from_config,
label_min_nonspace_chars_from_config,
label_names_from_config,
safe_auto_tokenizer,
token_extend_thresholds_from_config,
token_label_thresholds_from_config,
)
from multitask_model import IrishCoreTokenSpanModel
def replacement(label: str) -> str:
return f"[PII:{label}]"
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"]] + replacement(span["label"]) + out[span["end"] :]
return out
def predict(text: str, model, tokenizer, min_score: float):
encoded = tokenizer(text, return_offsets_mapping=True, return_tensors="pt", truncation=True)
offsets = [tuple(item) for item in encoded.pop("offset_mapping")[0].tolist()]
device = next(model.parameters()).device
encoded = {key: value.to(device) for key, value in encoded.items()}
with torch.no_grad():
output = model(**encoded)
token_scores = torch.sigmoid(output.token_logits[0]).cpu().numpy()
start_scores = torch.sigmoid(output.start_logits[0]).cpu().numpy()
end_scores = torch.sigmoid(output.end_logits[0]).cpu().numpy()
label_names = label_names_from_config(model.config)
thresholds = token_label_thresholds_from_config(model.config, min_score)
extend_thresholds = token_extend_thresholds_from_config(model.config)
max_span_tokens = label_max_span_tokens_from_config(model.config)
min_nonspace_chars = label_min_nonspace_chars_from_config(model.config)
boundary_thresholds = boundary_label_thresholds_from_config(model.config)
spans = decode_token_presence_segments(
text,
offsets,
token_scores,
label_names,
min_score,
thresholds,
extend_thresholds,
max_span_tokens,
min_nonspace_chars,
boundary_thresholds,
start_scores=start_scores,
end_scores=end_scores,
)
for span in spans:
span["replacement"] = replacement(span["label"])
return spans
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--model", required=True)
parser.add_argument("--text", required=True)
parser.add_argument("--min-score", type=float, default=0.5)
parser.add_argument("--device", choices=["auto", "cpu", "cuda"], default="auto")
parser.add_argument("--json", action="store_true")
args = parser.parse_args()
tokenizer = safe_auto_tokenizer(args.model)
config = AutoConfig.from_pretrained(args.model)
model = IrishCoreTokenSpanModel.from_pretrained(args.model, config=config)
if args.device == "auto":
device = "cuda" if torch.cuda.is_available() else "cpu"
else:
device = args.device
model.to(device)
model.eval()
spans = predict(args.text, model, tokenizer, args.min_score)
result = {
"model": args.model,
"backend": "transformers_token_span",
"min_score": args.min_score,
"spans": spans,
"masked_text": mask_text(args.text, spans),
}
if args.json:
print(json.dumps(result, indent=2, ensure_ascii=False))
else:
print(result["masked_text"])
if __name__ == "__main__":
main()
|