80 lines
3.2 KiB
Python
80 lines
3.2 KiB
Python
|
import os
|
||
|
|
import numpy as np
|
||
|
|
import onnxruntime as ort
|
||
|
|
from tokenizers import Tokenizer
|
||
|
|
|
||
|
|
import config
|
||
|
|
|
||
|
|
|
||
|
|
MODEL_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "models", "all-MiniLM-L6-v2")
|
||
|
|
MAX_TOKENS = 128
|
||
|
|
|
||
|
|
|
||
|
|
def _mean_pool(last_hidden: np.ndarray, attention_mask: np.ndarray) -> np.ndarray:
|
||
|
|
mask = attention_mask.astype(np.float32)[..., None]
|
||
|
|
summed = (last_hidden * mask).sum(axis=1)
|
||
|
|
counts = np.clip(mask.sum(axis=1), a_min=1e-9, a_max=None)
|
||
|
|
return summed / counts
|
||
|
|
|
||
|
|
|
||
|
|
def _l2_normalize(x: np.ndarray) -> np.ndarray:
|
||
|
|
norm = np.linalg.norm(x, axis=-1, keepdims=True)
|
||
|
|
return x / np.clip(norm, a_min=1e-12, a_max=None)
|
||
|
|
|
||
|
|
|
||
|
|
class IntentClassifier:
|
||
|
|
def __init__(self, model_dir: str = MODEL_DIR):
|
||
|
|
tok_path = os.path.join(model_dir, "tokenizer.json")
|
||
|
|
onnx_path = os.path.join(model_dir, "model.onnx")
|
||
|
|
self._tokenizer = Tokenizer.from_file(tok_path)
|
||
|
|
self._tokenizer.enable_truncation(max_length=MAX_TOKENS)
|
||
|
|
self._tokenizer.enable_padding(pad_id=0, pad_token="[PAD]")
|
||
|
|
so = ort.SessionOptions()
|
||
|
|
so.intra_op_num_threads = max(1, (os.cpu_count() or 2) // 2)
|
||
|
|
so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
|
||
|
|
self._session = ort.InferenceSession(onnx_path, sess_options=so, providers=["CPUExecutionProvider"])
|
||
|
|
self._input_names = {inp.name for inp in self._session.get_inputs()}
|
||
|
|
self._cache: dict[str, np.ndarray] = {}
|
||
|
|
|
||
|
|
def embed(self, texts: list[str]) -> np.ndarray:
|
||
|
|
if not texts:
|
||
|
|
return np.zeros((0, 384), dtype=np.float32)
|
||
|
|
encs = self._tokenizer.encode_batch(texts)
|
||
|
|
ids = np.array([e.ids for e in encs], dtype=np.int64)
|
||
|
|
mask = np.array([e.attention_mask for e in encs], dtype=np.int64)
|
||
|
|
feeds = {"input_ids": ids, "attention_mask": mask}
|
||
|
|
if "token_type_ids" in self._input_names:
|
||
|
|
feeds["token_type_ids"] = np.zeros_like(ids)
|
||
|
|
outputs = self._session.run(None, feeds)
|
||
|
|
last_hidden = outputs[0]
|
||
|
|
pooled = _mean_pool(last_hidden, mask)
|
||
|
|
return _l2_normalize(pooled).astype(np.float32)
|
||
|
|
|
||
|
|
def prime(self, candidates: dict[str, list[str]]) -> None:
|
||
|
|
self._cache.clear()
|
||
|
|
for cmd_id, phrases in candidates.items():
|
||
|
|
if not phrases:
|
||
|
|
continue
|
||
|
|
self._cache[cmd_id] = self.embed(list(phrases))
|
||
|
|
|
||
|
|
def match(self, utterance: str, candidates: dict[str, list[str]]) -> dict:
|
||
|
|
utterance = (utterance or "").strip()
|
||
|
|
if not utterance:
|
||
|
|
return {"cmd": "", "score": 0.0}
|
||
|
|
for cmd_id, phrases in candidates.items():
|
||
|
|
if cmd_id not in self._cache and phrases:
|
||
|
|
self._cache[cmd_id] = self.embed(list(phrases))
|
||
|
|
query = self.embed([utterance])[0]
|
||
|
|
best_cmd = ""
|
||
|
|
best_score = -1.0
|
||
|
|
for cmd_id, mat in self._cache.items():
|
||
|
|
sims = mat @ query
|
||
|
|
top = float(sims.max()) if sims.size else -1.0
|
||
|
|
if top > best_score:
|
||
|
|
best_score = top
|
||
|
|
best_cmd = cmd_id
|
||
|
|
threshold = getattr(config, "INTENT_SIMILARITY_THRESHOLD", 0.45)
|
||
|
|
if best_score < threshold:
|
||
|
|
return {"cmd": "", "score": max(0.0, best_score)}
|
||
|
|
return {"cmd": best_cmd, "score": max(0.0, best_score)}
|