File size: 5,257 Bytes
818b1b2 | 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 | """Shared word tokenization, actual SentencePiece special IDs, complete windows."""
import re
TOKEN_RE = re.compile(r'\w+(?:-\w+)*|[^\w\s]', re.UNICODE)
MAX_LEN = 256
def tokenizer_contract(sp):
return dict(bos_token_id=sp.bos_id(), eos_token_id=sp.eos_id(),
pad_token_id=sp.pad_id(), unk_token_id=sp.unk_id())
def word_windows(tokens, sp, max_len=MAX_LEN, overlap=0):
"""Yield (start_word, end_word, IDs, first-subword positions) without dropping words."""
if max_len < 3 or overlap < 0:
raise ValueError('Invalid window settings')
pieces = [sp.encode(t, out_type=int) or [sp.unk_id()] for t in tokens]
start = 0
while start < len(tokens):
ids, positions, end = [sp.bos_id()], [], start
while end < len(tokens) and len(ids)+len(pieces[end])+1 <= max_len:
positions.append(len(ids))
ids.extend(pieces[end])
end += 1
if end == start:
raise ValueError(f'Word at index {start} cannot fit into a {max_len}-subtoken window')
ids.append(sp.eos_id())
yield start, end, ids, positions
if end == len(tokens):
break
start = max(start+1, end-overlap)
"""Body appended to the shared tokenization helpers when packaging the release."""
import hashlib
import json
from pathlib import Path
import numpy as np
import onnxruntime as ort
import sentencepiece as spm
class BashkirTokenLID:
"""Standalone CPU word-level LID with offsets and complete long-text windows."""
def __init__(self, model_dir=None, precision='int8'):
self.directory=Path(model_dir) if model_dir is not None else Path(__file__).resolve().parent
config=json.loads((self.directory/'config.json').read_text(encoding='utf-8'))
meta=json.loads((self.directory/'META.json').read_text(encoding='utf-8'))
if precision not in config['models']:
raise ValueError(f"precision must be one of {sorted(config['models'])}")
model_name=config['models'][precision];tokenizer_name=config['tokenizer_file']
for name in [model_name,tokenizer_name,'config.json']:
with (self.directory/name).open('rb') as stream:
digest=hashlib.sha256()
for block in iter(lambda:stream.read(1024*1024),b''): digest.update(block)
actual=digest.hexdigest()
if actual!=meta['artifacts'][name]['sha256']:
raise ValueError(f'Artifact checksum mismatch: {name}')
self.tags=tuple(config['labels'])
self.max_length=int(config['max_position_embeddings'])
self.overlap=int(config['word_window_overlap'])
self.sp=spm.SentencePieceProcessor(model_file=str(self.directory/tokenizer_name))
expected={k:int(config[k]) for k in ['bos_token_id','eos_token_id','pad_token_id','unk_token_id']}
if tokenizer_contract(self.sp)!=expected or self.sp.get_piece_size()!=config['vocab_size']:
raise ValueError('Tokenizer does not match the model contract')
options=ort.SessionOptions();options.intra_op_num_threads=4;options.inter_op_num_threads=1
self.session=ort.InferenceSession(str(self.directory/model_name),sess_options=options,providers=['CPUExecutionProvider'])
def predict(self,text):
"""Return token, label, model confidence, start and end for every token.
Offsets index Python Unicode strings. Confidence is an uncalibrated model
score. Empty text returns []; overlong individual words raise ValueError.
"""
if not isinstance(text,str): raise TypeError('text must be a string')
matches=list(TOKEN_RE.finditer(text));tokens=[m.group() for m in matches]
result=[None]*len(tokens);priorities=[-1]*len(tokens)
for start,end,ids,positions in word_windows(tokens,self.sp,max_len=self.max_length,overlap=self.overlap):
logits=self.session.run(['logits'],{'input_ids':np.asarray([ids],dtype=np.int64)})[0][0,positions]
probs=np.exp(logits-logits.max(-1,keepdims=True));probs/=probs.sum(-1,keepdims=True)
for local,index in enumerate(range(start,end)):
priority=min(local+1,end-start-local)
if priority>priorities[index]:
tag=int(probs[local].argmax())
result[index]=dict(token=tokens[index],label=self.tags[tag],confidence=float(probs[local,tag]),
start=matches[index].start(),end=matches[index].end())
priorities[index]=priority
if any(r is None for r in result): raise AssertionError('Missing token after window inference')
return result
if __name__=='__main__':
import argparse
import sys
if hasattr(sys.stdout,'reconfigure'): sys.stdout.reconfigure(encoding='utf-8')
parser=argparse.ArgumentParser(description='Word-level Bashkir / Russian / Tatar LID')
parser.add_argument('text');parser.add_argument('--precision',choices=['int8','fp32'],default='int8')
args=parser.parse_args()
print(json.dumps(BashkirTokenLID(precision=args.precision).predict(args.text),ensure_ascii=False,indent=2))
|