| import json |
| import os |
| from pathlib import Path |
|
|
| import decode as D |
| import numpy as np |
| import rules |
|
|
|
|
| def windows(n, size, overlap): |
| if n <= size: |
| return [(0, n)] |
| out, start = [], 0 |
| while True: |
| end = min(n, start + size) |
| out.append((start, end)) |
| if end == n: |
| return out |
| start = end - overlap |
|
|
|
|
| class Detector: |
| def __init__(self, path, threads=None): |
| import onnxruntime as ort |
| from tokenizers import Tokenizer |
| path = Path(path) |
| self.cfg = json.loads((path / 'namecfg.json').read_text()) |
| self.tok = Tokenizer.from_file(str(path / 'tokenizer.json')) |
| self.tok.no_padding() |
| self.tok.no_truncation() |
| so = ort.SessionOptions() |
| so.intra_op_num_threads = threads or int(os.environ.get('NAME_THREADS', '4')) |
| so.inter_op_num_threads = 1 |
| so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL |
| self.sess = ort.InferenceSession(str(path / 'model.onnx'), so, providers=['CPUExecutionProvider']) |
| self.inputs = {i.name for i in self.sess.get_inputs()} |
| self.window = self.cfg['window'] |
| self.overlap = self.cfg['overlap'] |
| self.tau = self.cfg['tau'] |
| self.cls, self.sep = self.cfg['cls_id'], self.cfg['sep_id'] |
|
|
| def encode(self, text): |
| enc = self.tok.encode(text, add_special_tokens=False) |
| offsets = [] |
| for a, b in enc.offsets: |
| while a < b and text[a].isspace(): |
| a += 1 |
| while b > a and text[b - 1].isspace(): |
| b -= 1 |
| offsets.append((a, b)) |
| return enc.ids, offsets |
|
|
| def scores(self, text): |
| ids, offsets = self.encode(text) |
| n = len(ids) |
| probs = np.zeros(n, dtype=np.float32) |
| best = np.full(n, -1.0) |
| for a, b in windows(n, self.window, self.overlap): |
| x = np.array([[self.cls] + ids[a:b] + [self.sep]], dtype=np.int64) |
| feed = dict(input_ids=x, attention_mask=np.ones_like(x)) |
| if 'token_type_ids' in self.inputs: |
| feed['token_type_ids'] = np.zeros_like(x) |
| logits = self.sess.run(None, feed)[0][0, 1:-1] |
| e = np.exp(logits - logits.max(-1, keepdims=True)) |
| p = e[:, 1] / e.sum(-1) |
| j = np.arange(b - a, dtype=np.float64) |
| far = np.full(b - a, 1e9) |
| depth = np.minimum(j if a > 0 else far, (b - a - 1 - j) if b < n else far) |
| take = depth > best[a:b] |
| probs[a:b][take] = p[take] |
| best[a:b][take] = depth[take] |
| keep = [i for i, (a, b) in enumerate(offsets) if b > a] |
| return [list(offsets[i]) for i in keep], [float(probs[i]) for i in keep] |
|
|
| def findings(self, text, requester=None, authority=None): |
| offsets, probs = self.scores(text) |
| found = D.spans_for(text, offsets, probs, self.tau) |
| found = rules.apply(text, found, dict(offsets=offsets, probs=probs), requester, authority, self.tau) |
| cp = D.char_scores(text, offsets, probs) |
| return [dict(start=a, end=b, score=round(max(max(cp[a:b], default=0.0), self.tau), 4)) for a, b in found] |
|
|
