File size: 8,759 Bytes
77dfe62 da8de96 63a6829 da8de96 77dfe62 63a6829 77dfe62 da8de96 77dfe62 da8de96 77dfe62 da8de96 77dfe62 da8de96 77dfe62 da8de96 77dfe62 63a6829 77dfe62 63a6829 77dfe62 63a6829 77dfe62 63a6829 77dfe62 63a6829 77dfe62 63a6829 77dfe62 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 | """SpellSpeak Tone at run time: ONNX Runtime on the CPU, no torch.
A model folder holds `model.onnx` (encoder and heads in one graph; inputs input_ids, attention_mask and the
per-person marker masks; outputs the five heads' logits), `tokenizer.json` and `config.json` (class lists, the
hostility threshold tau, max length).
One encoder pass per line, whatever the number of people. The act toward each person is read at that person's
marker. The operating point is shared with the scorecard (harness/expression/decide.py): a pair is hostile exactly when
its hostility probability reaches tau (an invented grudge costs more than a miss).
r8 (line-tags-0.2): a model whose graph has an `exchange` output and whose config lists `exchanges` also returns the
line's exchange move (asks, answers, hedges, doesnt_know, withholds, closes, other) with its confidence, as
`LineTags02`. A model without it (RC1) returns `LineTags` exactly as before.
r10 (line-tags-0.3): a model whose graph also has a `sarcasm` output returns `sarcasm: {sarcastic, probability}` as
`LineTags03`: the calibrated chance the line is sarcastic, and the model's decision at its own threshold (the config's
`sarcasm.tau_raw` on the raw probability). Models without it return what they always did.
Precision: full precision ("fp32") is the default and the reference. From release rc2.1 a folder may also hold
`model_fp16.onnx`, the same graph with its weights stored in 16 bits (half the size, computed in 32 bits), picked with
`precision="fp16"`.
"""
from __future__ import annotations
import json
import pathlib
import time
from contracts.schemas.line_tags import ActTag, ExchangeTag, LineInput, LineTags, LineTags02, LineTags03, SarcasmTag, is_hostile
from harness.expression.calibrate import emotion_confidence, exchange_confidence, sarcasm_probability, tag_confidence
from harness.expression.decide import decide, summed_hostility
from harness.expression.render import render, span_mask
PRECISIONS = ("fp32", "fp16")
FP16_FILE = "model_fp16.onnx"
class LineClassifier:
def __init__(self, model_dir: str | pathlib.Path, onnx_file: str | None = None, threads: int = 4, precision: str = "fp32"):
"""`precision` picks the graph. "fp32" (the default): the transformer optimizer's graph when the export has one
(r3 on: same outputs, a little faster), else the plain graph. "fp16": `model_fp16.onnx`, the weights stored in
16 bits. `onnx_file` names a graph directly and overrides `precision`."""
import onnxruntime as ort
from tokenizers import Tokenizer
if precision not in PRECISIONS:
raise ValueError(f"precision must be one of {PRECISIONS}, not {precision!r}")
d = pathlib.Path(model_dir)
if onnx_file is None and precision == "fp16":
if not (d / FP16_FILE).exists():
raise FileNotFoundError(f"precision='fp16' needs {FP16_FILE} in {d}. It ships with release rc2.1 and later; "
"download it, or use precision='fp32'")
onnx_file = FP16_FILE
onnx_file = onnx_file or ("model.opt.onnx" if (d / "model.opt.onnx").exists() else "model.onnx")
self.onnx_file = onnx_file
self.precision = precision
self.config = json.loads((d / "config.json").read_text())
self.tok = Tokenizer.from_file(str(d / "tokenizer.json"))
self.tok.enable_truncation(self.config["max_len"])
self.tok.no_padding()
so = ort.SessionOptions()
so.intra_op_num_threads = threads
so.inter_op_num_threads = 1
self.sess = ort.InferenceSession(str(d / onnx_file), so, providers=["CPUExecutionProvider"])
self.acts = self.config["acts"]
self.calibration = self.config.get("calibration") # r7: reported confidences mean what they say
self.emotions = self.config["emotions"]
self.intensities = self.config["intensities"]
self.tau = float(self.config["tau"])
outputs = [o.name for o in self.sess.get_outputs()]
# r8: the exchange head, when the graph has it; RC1's graph does not
self.exchanges = self.config.get("exchanges") if "exchange" in outputs else None
self._exchange_at = outputs.index("exchange") if self.exchanges else None
# r10: the sarcasm head, when the graph has it; rc3 and earlier do not
self._sarcasm_at = outputs.index("sarcasm") if "sarcasm" in outputs else None
self.sarcasm_tau_raw = float((self.config.get("sarcasm") or {}).get("tau_raw", 0.5))
def logits(self, inp: LineInput):
"""Raw outputs for one line: emotion, emotion intensity, act per person, act intensity per person, backchannel."""
import numpy as np
text, spans = render(inp.model_dump())
enc = self.tok.encode(text)
ids = np.asarray([enc.ids], dtype=np.int64)
mask = np.asarray([enc.attention_mask], dtype=np.int64)
smask = np.asarray(span_mask(list(enc.offsets), spans), dtype=np.float32)
return self.sess.run(None, {"input_ids": ids, "attention_mask": mask, "span_mask": smask})
def tag(self, inp: LineInput) -> LineTags:
import numpy as np
def softmax(x):
e = np.exp(x - x.max(-1, keepdims=True))
return e / e.sum(-1, keepdims=True)
out = self.logits(inp)
le, lei, la, lai, lb = out[:5]
p_host = 1 / (1 + np.exp(-out[5])) if len(out) > 5 else None # the hostility head (every released model has one)
pe, pei, pa, pai = softmax(le[0]), softmax(lei[0]), softmax(la), softmax(lai)
emo = self.emotions[int(pe.argmax())]
if float(softmax(lb[0])[1]) >= 0.5: # a backchannel is not a turn: no act toward anyone
return self._with_sarcasm(self._with_exchange(LineTags(
emotion=emo, emotion_intensity="low" if emo == "neutral" else self.intensities[int(pei.argmax())],
emotion_confidence=round(emotion_confidence(float(pe.max()), self.calibration), 4),
acts={t: ActTag(act="none") for t in inp.targets}, backchannel=True), out), out)
acts = {}
for i, t in enumerate(inp.targets):
ph = float(p_host[i]) if p_host is not None else summed_hostility(pa[i])
act, inten = decide(pa[i], ph, self.tau, self.intensities[int(pai[i].argmax())])
# The confidence is the probability the decision rests on: for a hostile act, the hostility score (the
# act's own class can sit under 0.5 while the line is clearly hostile); otherwise the act's probability.
hostile = is_hostile(act, inten)
conf = ph if hostile else float(pa[i][self.acts.index(act)])
acts[t] = ActTag(act=act, intensity=inten, confidence=round(tag_confidence(conf, hostile, self.calibration), 4))
return self._with_sarcasm(self._with_exchange(LineTags(
emotion=emo, emotion_intensity="low" if emo == "neutral" else self.intensities[int(pei.argmax())],
emotion_confidence=round(emotion_confidence(float(pe.max()), self.calibration), 4), acts=acts, backchannel=False), out), out)
def _with_exchange(self, tags: LineTags, out) -> LineTags:
"""RC1 and any model without the exchange head: the tags unchanged. r8 on: the same tags plus the exchange move."""
if self.exchanges is None:
return tags
import numpy as np
z = np.asarray(out[self._exchange_at][0], dtype=np.float64)
p = np.exp(z - z.max())
p /= p.sum()
i = int(p.argmax())
conf = round(exchange_confidence(float(p[i]), self.calibration), 4)
return LineTags02(**{k: getattr(tags, k) for k in type(tags).model_fields},
exchange=ExchangeTag(label=self.exchanges[i], confidence=conf))
def _with_sarcasm(self, tags: LineTags, out) -> LineTags:
"""rc3 and earlier: the tags unchanged. r10 on: the same tags plus sarcasm, decided on the raw probability at the
model's threshold, reported as the calibrated probability."""
if self._sarcasm_at is None:
return tags
import numpy as np
z = np.asarray(out[self._sarcasm_at][0], dtype=np.float64)
p = np.exp(z - z.max())
raw = float(p[1] / p.sum())
return LineTags03(**{k: getattr(tags, k) for k in type(tags).model_fields},
sarcasm=SarcasmTag(sarcastic=raw >= self.sarcasm_tau_raw,
probability=round(sarcasm_probability(raw, self.calibration), 4)))
def timed(self, inp: LineInput) -> tuple[LineTags, float]:
t = time.perf_counter()
out = self.tag(inp)
return out, (time.perf_counter() - t) * 1000
|