File size: 3,183 Bytes
77daf58 | 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 | 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]
|