| """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": |
| |
| 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)) |
|
|
