| """Tern-Decide-0.1 public API (inference only, CPU, numpy + safetensors). |
| |
| from tern_decide import TernDecide |
| td = TernDecide.from_pretrained("path/to/Tern-Decide-0.1") # or a Hub id, see from_pretrained() |
| td.classify("I still have not received my new card", head="banking77").label |
| td.nli("A man plays a guitar.", "A man makes music.").probabilities |
| td.decide({"ticket": "..."}, "Is the customer asking for a refund?").probabilities |
| |
| Components load lazily on first use (thread-safe); call ``td.load()`` to load everything up front. |
| All outputs are deterministic: the same input always gives the same result.""" |
| from __future__ import annotations |
|
|
| import json |
| import os |
| import threading |
| from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple, Union |
|
|
| import numpy as np |
|
|
| from ._io import load, load_json |
| from .results import Classification, Decision, NLIResult |
|
|
| KINDS = ('noul', 'choice', 'score') |
| HEADS = ('banking77', 'clinc150') |
| MAX_OPTIONS = 255 |
| MAX_TEXT_CHARS = 200_000 |
| HUB_FILES = ('config.json', 'weights/decider.safetensors', 'weights/embeddings.safetensors', |
| 'weights/intent_banking77.safetensors', 'weights/intent_clinc150.safetensors', |
| 'weights/nli_encoder.safetensors', 'weights/oji_tokenizer.safetensors', |
| 'weights/valence_vocab.json', 'weights/word_vocab.json') |
| DEFAULT_REPO = 'OkeyMetaLtd/Tern-Decide-0.1' |
|
|
| Options = Union[None, Sequence[Any], Dict[Any, Any]] |
|
|
|
|
| class TernDecideError(ValueError): |
| """Raised for invalid inputs to Tern-Decide.""" |
|
|
|
|
| def _check_text(x: Any, name: str) -> str: |
| if not isinstance(x, str): |
| raise TernDecideError(f'{name} must be a str, got {type(x).__name__}') |
| if not x.strip(): |
| raise TernDecideError(f'{name} must be a non-empty string') |
| if len(x) > MAX_TEXT_CHARS: |
| raise TernDecideError(f'{name} is too long ({len(x)} chars; limit {MAX_TEXT_CHARS})') |
| return x |
|
|
|
|
| def _resolve_path(path: Optional[str], revision: Optional[str] = None) -> str: |
| """A local directory, $TERN_DECIDE_HOME, the repo this package sits in, or a Hugging Face Hub repo id.""" |
| if path is None: |
| path = os.environ.get('TERN_DECIDE_HOME') |
| if path is None: |
| here = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) |
| if os.path.isfile(os.path.join(here, 'weights', 'decider.safetensors')): |
| return here |
| path = DEFAULT_REPO |
| path = os.path.expanduser(str(path)) |
| if os.path.isdir(path): |
| if not os.path.isdir(os.path.join(path, 'weights')): |
| raise FileNotFoundError(f'{path!r} has no weights/ folder; point from_pretrained() at the Tern-Decide-0.1 repo root') |
| return path |
| if path.count('/') == 1 and not os.path.exists(path): |
| try: |
| from huggingface_hub import hf_hub_download |
| except ImportError as e: |
| raise ImportError(f'{path!r} is not a local folder. To download it from the Hugging Face Hub install ' |
| f'huggingface_hub (pip install "tern-decide[hub]") or download the repo yourself.') from e |
| root = None |
| for f in HUB_FILES: |
| p = hf_hub_download(repo_id=path, filename=f, revision=revision) |
| root = os.path.dirname(os.path.dirname(p)) if f.startswith('weights/') else os.path.dirname(p) |
| return root |
| raise FileNotFoundError(f'Tern-Decide weights not found at {path!r}') |
|
|
|
|
| class TernDecide: |
| """CPU-only decision / routing model (numpy + safetensors). |
| |
| decide(state, question, options=None, kind=None) -> Decision (kind: 'noul' | 'choice' | 'score') |
| classify(text, head='banking77' | 'clinc150') -> Classification |
| nli(premise, hypothesis) -> NLIResult |
| decide_batch / classify_batch / nli_batch -> lists of the same objects |
| """ |
|
|
| def __init__(self, path: Optional[str] = None, heads: Sequence[str] = HEADS, nli: bool = True, |
| lazy: bool = True, revision: Optional[str] = None): |
| for h in heads: |
| if h not in HEADS: |
| raise TernDecideError(f'unknown head {h!r}; available: {HEADS}') |
| self.path = _resolve_path(path, revision) |
| self._W = os.path.join(self.path, 'weights') |
| self.config = load_json(os.path.join(self.path, 'config.json')) |
| self._allowed_heads = tuple(heads) |
| self._nli_enabled = bool(nli) |
| self._lock = threading.RLock() |
| self._oji = self._words = self._emb = None |
| self._decider = None |
| self._heads: Dict[str, Any] = {} |
| self._enc = None |
| if not lazy: |
| self.load() |
|
|
| @classmethod |
| def from_pretrained(cls, path: Optional[str] = None, **kw) -> 'TernDecide': |
| """path: local repo folder, a Hub repo id (needs huggingface_hub), or None ($TERN_DECIDE_HOME / this repo).""" |
| return cls(path, **kw) |
|
|
| |
| def load(self) -> 'TernDecide': |
| """Load every enabled component now (otherwise each loads on first use).""" |
| self._get_decider() |
| for h in self._allowed_heads: |
| self._get_head(h) |
| if self._nli_enabled: |
| self._get_enc() |
| return self |
|
|
| @property |
| def heads(self) -> Tuple[str, ...]: |
| return self._allowed_heads |
|
|
| def _base(self): |
| if self._words is None: |
| with self._lock: |
| if self._words is None: |
| from .oji import Oji |
| from .words import Words |
| t, _ = load(os.path.join(self._W, 'oji_tokenizer.safetensors')) |
| oji = Oji(t['merges'], t['char_flags']) |
| e, _ = load(os.path.join(self._W, 'embeddings.safetensors')) |
| words = Words(e['word_E'], e['word_rowsum'], load_json(os.path.join(self._W, 'word_vocab.json'))) |
| self._emb, self._oji, self._words = e, oji, words |
| return self._oji, self._words |
|
|
| def _get_decider(self): |
| if self._decider is None: |
| with self._lock: |
| if self._decider is None: |
| from .features import Feat |
| _, words = self._base() |
| d, m = load(os.path.join(self._W, 'decider.safetensors')) |
| vv = load_json(os.path.join(self._W, 'valence_vocab.json')) |
| feat = Feat(words, dict(zip(vv, d['valence'].tolist())), json.loads(m['valence_norm'])) |
| w = {k: d['w_' + k] for k in KINDS} |
| self._decider = (feat, w, json.loads(m['T']), float(m['noul_delta'])) |
| return self._decider |
|
|
| def _get_head(self, head: str): |
| if head not in self._allowed_heads: |
| raise TernDecideError(f'head {head!r} not available; loaded heads: {self._allowed_heads}') |
| h = self._heads.get(head) |
| if h is None: |
| with self._lock: |
| h = self._heads.get(head) |
| if h is None: |
| from .intent import IntentHead |
| oji, words = self._base() |
| t, m = load(os.path.join(self._W, f'intent_{head}.safetensors')) |
| h = IntentHead(t, m, oji, words) |
| self._heads[head] = h |
| return h |
|
|
| def _get_enc(self): |
| if not self._nli_enabled: |
| raise RuntimeError('this TernDecide was created with nli=False') |
| if self._enc is None: |
| with self._lock: |
| if self._enc is None: |
| from .nli import PairEncoder |
| oji, words = self._base() |
| t, m = load(os.path.join(self._W, 'nli_encoder.safetensors')) |
| self._enc = PairEncoder(words, self._emb['tok_E'], self._emb['oov_map'], oji, t, m) |
| return self._enc |
|
|
| |
| @staticmethod |
| def _options(question: str, options: Options, kind: str) -> Tuple[List[str], List[str]]: |
| if kind == 'noul': |
| if options is None: |
| return ['true', 'false'], ['Yes: ' + question, 'No: it is not the case that ' + question] |
| if isinstance(options, dict): |
| if set(map(str, options)) != {'true', 'false'}: |
| raise TernDecideError("noul options as a dict need exactly the keys 'true' and 'false'") |
| return ['true', 'false'], [_check_text(str(options['true'] if 'true' in options else options[True]), 'options[true]'), |
| _check_text(str(options['false'] if 'false' in options else options[False]), 'options[false]')] |
| if isinstance(options, (list, tuple)) and len(options) == 2: |
| return ['true', 'false'], [_check_text(str(options[0]), 'options[0]'), _check_text(str(options[1]), 'options[1]')] |
| raise TernDecideError('noul options must be None, a (true_text, false_text) pair, or {"true": .., "false": ..}') |
| if options is None: |
| raise TernDecideError(f'kind={kind!r} needs options (a list or a dict of key -> description)') |
| if isinstance(options, dict): |
| keys = [str(k) for k in options] |
| texts = [f'{str(k).replace("_", " ")}: {v if isinstance(v, str) else json.dumps(v, sort_keys=True)}' for k, v in options.items()] |
| elif isinstance(options, (list, tuple)): |
| keys = [str(i) for i in range(len(options))] if kind == 'score' else [str(o) for o in options] |
| texts = [str(o) for o in options] |
| else: |
| raise TernDecideError(f'options must be a list, tuple or dict, got {type(options).__name__}') |
| if len(texts) < 2: |
| raise TernDecideError(f'kind={kind!r} needs at least 2 options') |
| if len(texts) > MAX_OPTIONS: |
| raise TernDecideError(f'at most {MAX_OPTIONS} options are supported, got {len(texts)}') |
| if len(set(keys)) != len(keys): |
| raise TernDecideError('option keys must be unique') |
| if any(not t.strip() for t in texts): |
| raise TernDecideError('options must be non-empty strings') |
| return keys, texts |
|
|
| def prepare(self, state: Any): |
| """Pre-process a state once (str, or JSON-serialisable dict/list) when asking several questions about it.""" |
| if isinstance(state, str): |
| _check_text(state, 'state') |
| else: |
| try: |
| json.dumps(state) |
| except (TypeError, ValueError) as e: |
| raise TernDecideError(f'state must be a str or JSON-serialisable dict/list: {e}') from e |
| feat = self._get_decider()[0] |
| return feat.prep_state(state) |
|
|
| def decide(self, state: Any, question: str, options: Options = None, kind: Optional[str] = None, |
| prepared=None) -> Decision: |
| """Zero-shot decision. kind defaults to 'noul' (yes/no) without options and 'choice' with options. |
| |
| EXPERIMENTAL: see the model card. On the only workflow benchmark we have, this decider is below a trivial |
| constant baseline; treat its Noul output as a calibrated prior, not as a reading of the text.""" |
| _check_text(question, 'question') |
| kind = kind or ('noul' if options is None else 'choice') |
| if kind not in KINDS: |
| raise TernDecideError(f'kind must be one of {KINDS}, got {kind!r}') |
| feat, w, T, delta = self._get_decider() |
| S = prepared if prepared is not None else self.prepare(state) |
| keys, texts = self._options(question, options, kind) |
| s = feat.item(S, question, texts, kind).astype(np.float64) @ w[kind] |
| z = s / T[kind] |
| if kind == 'noul': |
| z[0] += delta |
| p = np.exp(z - z.max()) |
| p /= p.sum() |
| probs = {k: float(v) for k, v in zip(keys, p)} |
| expected = float((np.arange(len(p)) * p).sum()) if kind == 'score' else None |
| return Decision(label=keys[int(np.argmax(p))], probabilities=probs, kind=kind, expected=expected) |
|
|
| def decide_batch(self, items: Iterable[Any], state: Any = None) -> List[Decision]: |
| """Many decisions in one call. |
| |
| With ``state`` given, each item is a question string or a dict {question, options?, kind?}; the state is |
| prepared once. Without it, each item is a dict {state, question, options?, kind?} or a tuple |
| (state, question[, options[, kind]]). Results come back in input order.""" |
| items = list(items) |
| out: List[Decision] = [] |
| shared = self.prepare(state) if state is not None else None |
| cache: Dict[int, Any] = {} |
| for i, it in enumerate(items): |
| if isinstance(it, str): |
| if shared is None: |
| raise TernDecideError(f'item {i}: a bare question needs decide_batch(..., state=...)') |
| st, q, opt, kd = None, it, None, None |
| elif isinstance(it, dict): |
| q, opt, kd = it.get('question'), it.get('options'), it.get('kind') |
| st = it.get('state') |
| if 'question' not in it: |
| raise TernDecideError(f'item {i}: missing "question"') |
| elif isinstance(it, (list, tuple)) and 2 <= len(it) <= 4: |
| st, q = it[0], it[1] |
| opt = it[2] if len(it) > 2 else None |
| kd = it[3] if len(it) > 3 else None |
| else: |
| raise TernDecideError(f'item {i}: expected a str, dict or (state, question[, options[, kind]]) tuple') |
| if st is None: |
| if shared is None: |
| raise TernDecideError(f'item {i}: no state given') |
| prep = shared |
| else: |
| key = id(st) |
| if key not in cache: |
| cache[key] = (st, self.prepare(st)) |
| prep = cache[key][1] |
| out.append(self.decide(None, q, opt, kd, prepared=prep)) |
| return out |
|
|
| |
| def classify(self, text: str, head: str = 'banking77', top_k: int = 5, oos: str = 'train') -> Classification: |
| """Supervised intent routing. head: 'banking77' (77 intents) or 'clinc150' (150 intents + 'oos'). |
| oos (clinc150 only): 'train' uses the trained 'oos' class; 'threshold' also returns 'oos' when the best in-scope |
| probability is below the dev-calibrated threshold.""" |
| _check_text(text, 'text') |
| if not isinstance(top_k, int) or isinstance(top_k, bool) or top_k < 1: |
| raise TernDecideError('top_k must be a positive int') |
| if oos not in ('train', 'threshold'): |
| raise TernDecideError("oos must be 'train' or 'threshold'") |
| if head not in HEADS: |
| raise TernDecideError(f'unknown head {head!r}; available: {HEADS}') |
| h = self._get_head(head) |
| p = h.proba(text) |
| best = int(p.argmax()) |
| if oos == 'threshold' and h.oos_th is not None: |
| pin = p[:len(h.labels) - 1] |
| best = len(h.labels) - 1 if pin.max() < h.oos_th else int(pin.argmax()) |
| order = np.argsort(-p, kind='stable')[:top_k] |
| probs = {h.labels[j]: float(p[j]) for j in range(len(p))} |
| return Classification(label=h.labels[best], probabilities=probs, kind='intent', head=head, |
| top=[(h.labels[j], float(p[j])) for j in order]) |
|
|
| def classify_batch(self, texts: Iterable[str], head: str = 'banking77', top_k: int = 5, |
| oos: str = 'train') -> List[Classification]: |
| return [self.classify(t, head=head, top_k=top_k, oos=oos) for t in texts] |
|
|
| |
| def nli_scores(self, premise: str, hypothesis: str) -> np.ndarray: |
| _check_text(premise, 'premise'); _check_text(hypothesis, 'hypothesis') |
| return self._get_enc().scores([(premise, hypothesis)])[0] |
|
|
| def _nli_result(self, s: np.ndarray, labels: List[str]) -> NLIResult: |
| z = s.astype(np.float64) / float(self.config['nli']['temperature']) |
| p = np.exp(z - z.max()) |
| p /= p.sum() |
| probs = {k: float(v) for k, v in zip(labels, p)} |
| return NLIResult(label=labels[int(np.argmax(p))], probabilities=probs, kind='nli') |
|
|
| def nli(self, premise: str, hypothesis: str) -> NLIResult: |
| """Entailment / neutral / contradiction for one (premise, hypothesis) pair.""" |
| enc = self._get_enc() |
| return self._nli_result(self.nli_scores(premise, hypothesis), enc.labels) |
|
|
| def nli_batch(self, pairs: Iterable[Tuple[str, str]], batch_size: int = 64) -> List[NLIResult]: |
| """Batched NLI (the projection layer runs once per chunk of ``batch_size`` pairs).""" |
| pairs = list(pairs) |
| for i, pr in enumerate(pairs): |
| if not isinstance(pr, (list, tuple)) or len(pr) != 2: |
| raise TernDecideError(f'pair {i}: expected a (premise, hypothesis) tuple') |
| _check_text(pr[0], f'pairs[{i}][0]'); _check_text(pr[1], f'pairs[{i}][1]') |
| if not isinstance(batch_size, int) or batch_size < 1: |
| raise TernDecideError('batch_size must be a positive int') |
| enc = self._get_enc() |
| out: List[NLIResult] = [] |
| for a in range(0, len(pairs), batch_size): |
| S = enc.scores([tuple(p) for p in pairs[a:a + batch_size]]) |
| out.extend(self._nli_result(s, enc.labels) for s in S) |
| return out |
|
|
