File size: 3,938 Bytes
16409df | 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 | """Standalone ONNX inference for Bashkir person-name recognition.
The graph returns logits at SentencePiece positions. Training assigned one BIO
label to the first piece of each original word, so this runtime does the same.
"""
from __future__ import annotations
import json
import re
from pathlib import Path
import numpy as np
import onnxruntime as ort
import sentencepiece as spm
TAGS = ("O", "B-PER", "I-PER")
WORD_RE = re.compile(r"[A-Za-zА-Яа-яЁёӘәҒғҘҙҠҡҢңӨөҪҫҮүҺһҚқІі][\w'\-]*", re.UNICODE)
class BashkirPersonNER:
def __init__(self, model_dir: str | Path = ".", precision: str = "int8",
providers: list[str] | None = None):
self.directory = Path(model_dir)
config = json.loads((self.directory / "config.json").read_text(encoding="utf-8"))
if precision not in config["models"]:
raise ValueError(f"precision must be one of {sorted(config['models'])}")
self.max_length = int(config["max_position_embeddings"])
self.bos_id = int(config["bos_token_id"])
self.eos_id = int(config["eos_token_id"])
self.sp = spm.SentencePieceProcessor(
model_file=str(self.directory / "spm_bashkir_bert_16k.model")
)
model = self.directory / config["models"][precision]
self.session = ort.InferenceSession(
str(model), providers=providers or ["CPUExecutionProvider"]
)
def predict_tokens(self, tokens: list[str]) -> tuple[list[str], bool]:
"""Return word-level BIO tags and whether input exceeded the context limit.
The caller supplies the same token boundaries used for gold evaluation.
On truncation, the returned tag list covers only its leading words.
"""
if not tokens:
return [], False
ids = [self.bos_id]
positions = []
for word in tokens:
pieces = list(self.sp.encode(word, out_type=int)) or [self.sp.unk_id()]
if len(ids) + len(pieces) + 1 > self.max_length:
break
positions.append(len(ids))
ids.extend(pieces)
if not positions:
return [], True
ids.append(self.eos_id)
logits = self.session.run(
["logits"], {"input_ids": np.asarray([ids], dtype=np.int64)}
)[0][0]
return [TAGS[int(logits[pos].argmax())] for pos in positions], len(positions) < len(tokens)
def predict(self, text: str) -> dict:
"""Return person spans with character offsets into the original text."""
if not isinstance(text, str):
raise TypeError("text must be a string")
matches = list(WORD_RE.finditer(text))
tags, truncated = self.predict_tokens([m.group() for m in matches])
entities = []
start = None
end = None
for match, tag in zip(matches, tags):
if tag == "B-PER":
if start is not None:
entities.append({"text": text[start:end], "start": start, "end": end})
start, end = match.span()
elif tag == "I-PER":
# Current benchmark convention: an orphan I-PER starts a span.
if start is None:
start = match.start()
end = match.end()
elif start is not None:
entities.append({"text": text[start:end], "start": start, "end": end})
start = end = None
if start is not None:
entities.append({"text": text[start:end], "start": start, "end": end})
return {"entities": entities, "truncated": truncated}
if __name__ == "__main__":
import sys
if hasattr(sys.stdout, "reconfigure"):
sys.stdout.reconfigure(encoding="utf-8")
ner = BashkirPersonNER(Path(__file__).resolve().parent)
for sample in sys.argv[1:]:
print(json.dumps(ner.predict(sample), ensure_ascii=False))
|