| """ |
| CPU-optimized handler for BERT-OJA-SkillLess (ONNX). |
| Uses ONNX Runtime with all available CPU threads, small internal batches |
| for cache efficiency, and numpy for fast softmax. |
| """ |
| from typing import Dict, List, Any |
| import os |
| import numpy as np |
| from transformers import AutoTokenizer |
| import onnxruntime as ort |
|
|
|
|
| class EndpointHandler: |
| def __init__(self, path=""): |
| self.tokenizer = AutoTokenizer.from_pretrained(path) |
|
|
| n_threads = int(os.environ.get("OMP_NUM_THREADS", os.cpu_count() or 4)) |
|
|
| opts = ort.SessionOptions() |
| opts.intra_op_num_threads = n_threads |
| opts.inter_op_num_threads = max(1, n_threads // 2) |
| opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL |
| opts.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL |
|
|
| onnx_path = os.path.join(path, "model.onnx") |
| self.session = ort.InferenceSession( |
| onnx_path, sess_options=opts, providers=["CPUExecutionProvider"], |
| ) |
| self.input_names = [i.name for i in self.session.get_inputs()] |
| self.batch_size = 128 |
| print(f"[handler] ONNX CPU ready — threads={n_threads}, batch={self.batch_size}") |
|
|
| def __call__(self, data: Dict[str, Any]) -> List[List[Dict[str, float]]]: |
| inputs = data.get("inputs", data.get("input", "")) |
| if isinstance(inputs, str): |
| inputs = [inputs] |
|
|
| all_results = [] |
| for i in range(0, len(inputs), self.batch_size): |
| batch = inputs[i : i + self.batch_size] |
| encoded = self.tokenizer( |
| batch, padding=True, truncation=True, |
| max_length=128, return_tensors="np", |
| ) |
| feed = {k: encoded[k] for k in self.input_names if k in encoded} |
| logits = self.session.run(None, feed)[0] |
|
|
| exp = np.exp(logits - logits.max(axis=-1, keepdims=True)) |
| probs = exp / exp.sum(axis=-1, keepdims=True) |
|
|
| for j in range(len(batch)): |
| all_results.append([ |
| {"label": "LABEL_0", "score": round(float(probs[j][0]), 6)}, |
| {"label": "LABEL_1", "score": round(float(probs[j][1]), 6)}, |
| ]) |
|
|
| return all_results |
|
|
