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