File size: 2,606 Bytes
268e727 | 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 | """CPU inference using only NumPy, Tokenizers and ONNX Runtime.
The host supplies normalized, optional observations. It also owns admission,
async execution and cancellation; this module creates no background workers.
"""
import json
from pathlib import Path
import numpy as np
from .router import Router
class JsonTokenizer:
"""Read the same exported fast tokenizer without loading Transformers."""
def __init__(self, path):
from tokenizers import Tokenizer
path = Path(path)
self.tokenizer = Tokenizer.from_file(str(path / 'tokenizer.json'))
config = json.loads((path / 'tokenizer_config.json').read_text())
pad = config['pad_token']
if isinstance(pad, dict):
pad = pad['content']
pad_id = self.tokenizer.token_to_id(pad)
if pad_id is None or config.get('padding_side', 'right') != 'right':
raise ValueError('Expected an exported tokenizer with right padding')
self.tokenizer.enable_truncation(max_length=384, direction='right')
self.tokenizer.enable_padding(direction='right', pad_id=pad_id, pad_token=pad)
def __call__(self, texts, *, padding, truncation, max_length, return_tensors):
if (padding is not True or truncation is not True
or max_length != 384 or return_tensors != 'np'):
raise ValueError('Unsupported tokenization contract')
values = self.tokenizer.encode_batch(texts, add_special_tokens=True)
return {'input_ids': np.asarray([r.ids for r in values], dtype=np.int64),
'attention_mask': np.asarray([r.attention_mask for r in values], dtype=np.int64)}
class OnnxRouter(Router):
"""Drop-in synchronous typed router, with no PyTorch or Studio dependency."""
def __init__(self, path, *, admission=None):
super().__init__(path, backend='onnx', device='cpu', admission=admission)
def _load(self):
import onnxruntime as ort
if self.meta.get('format') != 'onnx':
raise ValueError('Expected an ONNX artifact')
self.tokenizer = JsonTokenizer(self.path / 'tokenizer')
options = ort.SessionOptions()
options.intra_op_num_threads, options.inter_op_num_threads = 2, 1
options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
options.add_session_config_entry('session.intra_op.allow_spinning', '0')
options.add_session_config_entry('session.inter_op.allow_spinning', '0')
self.model = ort.InferenceSession(str(self.path / 'model.onnx'),
sess_options=options, providers=['CPUExecutionProvider'])
|