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