"""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): # looks like a Hub repo id 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: # per-file download (no snapshot) 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) # ------------------------------------------------------------------ loading (lazy, thread-safe) 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 # ------------------------------------------------------------------ decider @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)) # keep st alive so id() stays unique prep = cache[key][1] out.append(self.decide(None, q, opt, kd, prepared=prep)) return out # ------------------------------------------------------------------ intent heads 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] # ------------------------------------------------------------------ NLI 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