SpellSpeak Audience rc2.1
Browse files- MANIFEST.json +17 -9
- README.md +35 -10
- encoder_fp16.onnx +3 -0
- example.py +12 -6
- head_fp16.onnx +3 -0
- runtime/harness/addressee/classifier.py +33 -6
MANIFEST.json
CHANGED
|
@@ -1,14 +1,14 @@
|
|
| 1 |
{
|
| 2 |
"name": "SpellSpeak Audience",
|
| 3 |
-
"release": "rc2",
|
| 4 |
-
"package": "spellspeak-audience-rc2.
|
| 5 |
-
"created": "2026-10-
|
| 6 |
"model_run": "addr-r3-ettin-32m-tt-s0",
|
| 7 |
"inputs": "addressee-tt-0.6",
|
| 8 |
"files": {
|
| 9 |
"README.md": {
|
| 10 |
-
"bytes":
|
| 11 |
-
"sha256": "
|
| 12 |
},
|
| 13 |
"config.json": {
|
| 14 |
"bytes": 341,
|
|
@@ -18,14 +18,22 @@
|
|
| 18 |
"bytes": 127960962,
|
| 19 |
"sha256": "8db0b9abf290b82b12ae9e71066e607e91a60d310562d84fc67a8668a5020f0d"
|
| 20 |
},
|
|
|
|
|
|
|
|
|
|
|
|
|
| 21 |
"example.py": {
|
| 22 |
-
"bytes":
|
| 23 |
-
"sha256": "
|
| 24 |
},
|
| 25 |
"head.onnx": {
|
| 26 |
"bytes": 3692356,
|
| 27 |
"sha256": "67c08fae8fd4712ebf08e1955c5b40800467cd3b6efb9d425539d2afcca2eb8c"
|
| 28 |
},
|
|
|
|
|
|
|
|
|
|
|
|
|
| 29 |
"requirements.txt": {
|
| 30 |
"bytes": 270,
|
| 31 |
"sha256": "7ca7dcd5aa19a996d459ef0ab9f8e43f61d814931525369e57664044e66d72db"
|
|
@@ -59,8 +67,8 @@
|
|
| 59 |
"sha256": "ffbc520ee544ebc2fea6a25fb99b2d6edd3232a339a6d49b85a96a16de5db5d8"
|
| 60 |
},
|
| 61 |
"runtime/harness/addressee/classifier.py": {
|
| 62 |
-
"bytes":
|
| 63 |
-
"sha256": "
|
| 64 |
},
|
| 65 |
"runtime/harness/addressee/evidence.py": {
|
| 66 |
"bytes": 10185,
|
|
|
|
| 1 |
{
|
| 2 |
"name": "SpellSpeak Audience",
|
| 3 |
+
"release": "rc2.1",
|
| 4 |
+
"package": "spellspeak-audience-rc2.2",
|
| 5 |
+
"created": "2026-10-09T22:24:09Z",
|
| 6 |
"model_run": "addr-r3-ettin-32m-tt-s0",
|
| 7 |
"inputs": "addressee-tt-0.6",
|
| 8 |
"files": {
|
| 9 |
"README.md": {
|
| 10 |
+
"bytes": 16566,
|
| 11 |
+
"sha256": "c41e39a18c7a9d84d532065749abbde6d3b618eb35865a0b958d82881d513873"
|
| 12 |
},
|
| 13 |
"config.json": {
|
| 14 |
"bytes": 341,
|
|
|
|
| 18 |
"bytes": 127960962,
|
| 19 |
"sha256": "8db0b9abf290b82b12ae9e71066e607e91a60d310562d84fc67a8668a5020f0d"
|
| 20 |
},
|
| 21 |
+
"encoder_fp16.onnx": {
|
| 22 |
+
"bytes": 64103200,
|
| 23 |
+
"sha256": "9d71e550b8e497ed1c96dd6992f5133dc2f244fa708367cc5d024478a11ef7e2"
|
| 24 |
+
},
|
| 25 |
"example.py": {
|
| 26 |
+
"bytes": 4097,
|
| 27 |
+
"sha256": "676227bc4fb3c33c8fc0a99f2645dafacc56ac4586ee706233f0f1baac429338"
|
| 28 |
},
|
| 29 |
"head.onnx": {
|
| 30 |
"bytes": 3692356,
|
| 31 |
"sha256": "67c08fae8fd4712ebf08e1955c5b40800467cd3b6efb9d425539d2afcca2eb8c"
|
| 32 |
},
|
| 33 |
+
"head_fp16.onnx": {
|
| 34 |
+
"bytes": 1858520,
|
| 35 |
+
"sha256": "78361e61605406bb0d8c13de148847684e55724fe183f24a713df2c4cee18655"
|
| 36 |
+
},
|
| 37 |
"requirements.txt": {
|
| 38 |
"bytes": 270,
|
| 39 |
"sha256": "7ca7dcd5aa19a996d459ef0ab9f8e43f61d814931525369e57664044e66d72db"
|
|
|
|
| 67 |
"sha256": "ffbc520ee544ebc2fea6a25fb99b2d6edd3232a339a6d49b85a96a16de5db5d8"
|
| 68 |
},
|
| 69 |
"runtime/harness/addressee/classifier.py": {
|
| 70 |
+
"bytes": 9314,
|
| 71 |
+
"sha256": "3d3b7e8c36988721dcf2f257d09584d666065c0bd0f0fe1fe22f88954a258864"
|
| 72 |
},
|
| 73 |
"runtime/harness/addressee/evidence.py": {
|
| 74 |
"bytes": 10185,
|
README.md
CHANGED
|
@@ -19,7 +19,7 @@ tags:
|
|
| 19 |
- dialogue
|
| 20 |
---
|
| 21 |
|
| 22 |
-
# SpellSpeak Audience · RC2
|
| 23 |
|
| 24 |
When several AI characters listen to the same person, as in a game, an interactive story or any other multi-character voice experience, the application has to know who each line is for before any character answers or reacts. A name or a description often settles it, but much of what people say names nobody: "Hello.", "Who are you?", "You lied to me.". Once the characters talk among themselves, the application also needs to know who each character's line is for, to know who speaks next. The language model voicing each character could guess, but only by reading every character's context for every line. That adds latency, the characters can disagree, and a wrong guess is heard at once.
|
| 25 |
|
|
@@ -27,7 +27,18 @@ SpellSpeak Audience treats the question as classification by a small model. One
|
|
| 27 |
|
| 28 |
[Try it in your browser](https://huggingface.co/spaces/spellspeak/audience-demo) · [Examples and source](https://github.com/spellspeak/models/tree/main/audience)
|
| 29 |
|
| 30 |
-
## What is new in RC2
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 31 |
|
| 32 |
- **Questions to the room.** "Who knows where the mill key is?", "Anyone seen my horse?", "Does anybody here sell rope?" are for everyone who can hear them, whoever the speaker is looking at. RC1 sent most of them to the one person looked at. With the speaker facing one person, RC2 gets all 63 test lines right, with the facts and without them; RC1 gets 27% with the facts and 84% without.
|
| 33 |
- **Lines said by characters.** The speaker may be a character. The player is then one of the people the line can be said to, with a card of what the characters see of them. A character's line goes to the one it was answering, unless its words pick someone else out: "Wren, you were there, weren't you?" is for Wren, and "Ask Wren, she'd know." still answers the player.
|
|
@@ -35,24 +46,35 @@ SpellSpeak Audience treats the question as classification by a small model. One
|
|
| 35 |
|
| 36 |
| | |
|
| 37 |
|---|---|
|
| 38 |
-
| **Released** | 2026-10-09 · rc2 (package `spellspeak-audience-rc2.
|
| 39 |
| **What it does** | Works out who a line is for, the player's or a character's: one person, several named on purpose, the whole group, or unclear |
|
| 40 |
-
| **Size** | About 33M parameters
|
| 41 |
-
| **Speed** | p95 per line on 4 CPU threads (Apple M1 Max): 5.5 ms for 1 person, 5.7 for 2, 6.0 for 4, 6.6 for 8. On 1 thread: 11 to 12 ms. The first line of a scene also encodes everyone's card: 14 ms for 1 person to 61 ms for 8; cards are then cached |
|
| 42 |
-
| **Runtime** | ONNX Runtime and tokenizers, on CPU. A Python reference runtime ships with it, including the function that turns positions into the spatial facts and an instant rules baseline. It also runs RC1's model files unchanged |
|
| 43 |
-
| **Files** | `encoder.onnx`
|
| 44 |
| **Licence** | Apache-2.0 |
|
| 45 |
|
| 46 |
## Quick start
|
| 47 |
|
|
|
|
|
|
|
| 48 |
```bash
|
| 49 |
pip install -U huggingface_hub
|
| 50 |
-
hf download spellspeak/audience --revision rc2 --local-dir spellspeak-audience
|
| 51 |
cd spellspeak-audience
|
| 52 |
pip install -r requirements.txt
|
| 53 |
python example.py
|
| 54 |
```
|
| 55 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 56 |
The example puts a barkeep, an elf mercenary and an old sailor in a room, with the player facing the elf. It asks who four player lines are for, with and without the spatial facts. Then the barkeep speaks, after the player asked him where the mill key is:
|
| 57 |
|
| 58 |
```
|
|
@@ -68,6 +90,8 @@ Tomas: 'No idea, sorry.' text only player (1.00)
|
|
| 68 |
Tomas: 'Wren, were you there?' text only wren (1.00)
|
| 69 |
```
|
| 70 |
|
|
|
|
|
|
|
| 71 |
In your own code, with the downloaded folder as the working directory:
|
| 72 |
|
| 73 |
```python
|
|
@@ -78,7 +102,8 @@ from contracts.schemas.addressee import AddresseeRequest
|
|
| 78 |
from contracts.schemas.person_card import PersonCard
|
| 79 |
from harness.addressee.classifier import load_classifier
|
| 80 |
|
| 81 |
-
model = load_classifier(".") #
|
|
|
|
| 82 |
people = [
|
| 83 |
PersonCard(id="tomas", label="Tomas", aliases=["the barkeep"], features=[{"key": "headwear", "value": "red hat"}]),
|
| 84 |
PersonCard(id="wren", label="Wren", features=[{"key": "role", "value": "mercenary"}]),
|
|
@@ -151,4 +176,4 @@ Measured on lines as given only:
|
|
| 151 |
- **Spatial facts depend on the application's band edges.** Tune them (`BandConfig` in `bands.py`).
|
| 152 |
- **It has not met real users yet.** It was tested in game settings only, on generated and hand-written lines, not on logs from real people or from other multi-character voice experiences.
|
| 153 |
|
| 154 |
-
The numbers come from the training repository's reports for this release
|
|
|
|
| 19 |
- dialogue
|
| 20 |
---
|
| 21 |
|
| 22 |
+
# SpellSpeak Audience · RC2.1
|
| 23 |
|
| 24 |
When several AI characters listen to the same person, as in a game, an interactive story or any other multi-character voice experience, the application has to know who each line is for before any character answers or reacts. A name or a description often settles it, but much of what people say names nobody: "Hello.", "Who are you?", "You lied to me.". Once the characters talk among themselves, the application also needs to know who each character's line is for, to know who speaks next. The language model voicing each character could guess, but only by reading every character's context for every line. That adds latency, the characters can disagree, and a wrong guess is heard at once.
|
| 25 |
|
|
|
|
| 27 |
|
| 28 |
[Try it in your browser](https://huggingface.co/spaces/spellspeak/audience-demo) · [Examples and source](https://github.com/spellspeak/models/tree/main/audience)
|
| 29 |
|
| 30 |
+
## What is new in RC2.1
|
| 31 |
+
|
| 32 |
+
- **A 16-bit copy at half the size, with the same answers.** `encoder_fp16.onnx` and `head_fp16.onnx` hold the same weights stored in 16 bits: 66.0 MB against 131.7 MB. They are cast back to 32 bits when they load.
|
| 33 |
+
- On the test lines, 99.92% of their answers read exactly as full precision's, with the spatial facts and without them. The other two lines only change which of several tied people is on top.
|
| 34 |
+
- No scorecard item is worse. A difference counts when it is 0.1 points or more on a set of at least 30 lines.
|
| 35 |
+
- A wrong person scored 0.8 or more stays at 0.69% with facts and 0.45% without.
|
| 36 |
+
- **Memory and speed match full precision's**, because the weights become 32-bit when they load: 364 MB resident against 370 MB. The p95 per line is within a millisecond of full precision's, on one thread of an Apple M1 Max.
|
| 37 |
+
- **Full precision is unchanged and stays the default.** It is the same files as RC2, byte for byte.
|
| 38 |
+
|
| 39 |
+
Full precision is the reference. Every other number on this card is full precision's, and it is the file to keep for any further training.
|
| 40 |
+
|
| 41 |
+
## What RC2 added
|
| 42 |
|
| 43 |
- **Questions to the room.** "Who knows where the mill key is?", "Anyone seen my horse?", "Does anybody here sell rope?" are for everyone who can hear them, whoever the speaker is looking at. RC1 sent most of them to the one person looked at. With the speaker facing one person, RC2 gets all 63 test lines right, with the facts and without them; RC1 gets 27% with the facts and 84% without.
|
| 44 |
- **Lines said by characters.** The speaker may be a character. The player is then one of the people the line can be said to, with a card of what the characters see of them. A character's line goes to the one it was answering, unless its words pick someone else out: "Wren, you were there, weren't you?" is for Wren, and "Ask Wren, she'd know." still answers the player.
|
|
|
|
| 46 |
|
| 47 |
| | |
|
| 48 |
|---|---|
|
| 49 |
+
| **Released** | 2026-10-09 · rc2.1 (package `spellspeak-audience-rc2.2`, model run `addr-r3-ettin-32m-tt-s0`, inputs `addressee-tt-0.6`) |
|
| 50 |
| **What it does** | Works out who a line is for, the player's or a character's: one person, several named on purpose, the whole group, or unclear |
|
| 51 |
+
| **Size** | About 33M parameters. Full precision (fp32): 131.7 MB in two graphs, a 128.0 MB encoder and a 3.7 MB head. 16-bit: 66.0 MB, a 64.1 MB encoder and a 1.9 MB head. Both use the same 3.6 MB tokenizer |
|
| 52 |
+
| **Speed** | p95 per line on 4 CPU threads (Apple M1 Max): 5.5 ms for 1 person, 5.7 for 2, 6.0 for 4, 6.6 for 8. On 1 thread: 11 to 12 ms. The first line of a scene also encodes everyone's card: 14 ms for 1 person to 61 ms for 8; cards are then cached. The 16-bit files time the same: within a millisecond of full precision's on one thread, timed in the same session |
|
| 53 |
+
| **Runtime** | ONNX Runtime and tokenizers, on CPU. A Python reference runtime ships with it, including the function that turns positions into the spatial facts and an instant rules baseline. It reads full precision unless asked for the 16-bit files (`precision="fp16"`), and it also runs RC1's model files unchanged |
|
| 54 |
+
| **Files** | `encoder.onnx` and `head.onnx` (full precision, the default), `encoder_fp16.onnx` and `head_fp16.onnx` (16-bit), `tokenizer.json` and `config.json`, with the reference runtime (`runtime/`), `example.py` and `requirements.txt`. `MANIFEST.json` lists every file with its size and sha256 |
|
| 55 |
| **Licence** | Apache-2.0 |
|
| 56 |
|
| 57 |
## Quick start
|
| 58 |
|
| 59 |
+
Everything, full precision and the 16-bit files:
|
| 60 |
+
|
| 61 |
```bash
|
| 62 |
pip install -U huggingface_hub
|
| 63 |
+
hf download spellspeak/audience --revision rc2.1 --local-dir spellspeak-audience
|
| 64 |
cd spellspeak-audience
|
| 65 |
pip install -r requirements.txt
|
| 66 |
python example.py
|
| 67 |
```
|
| 68 |
|
| 69 |
+
Or the small download, the 16-bit files without the full-precision graphs:
|
| 70 |
+
|
| 71 |
+
```bash
|
| 72 |
+
hf download spellspeak/audience --revision rc2.1 --local-dir spellspeak-audience-fp16 --exclude encoder.onnx --exclude head.onnx
|
| 73 |
+
cd spellspeak-audience-fp16
|
| 74 |
+
pip install -r requirements.txt
|
| 75 |
+
python example.py --fp16
|
| 76 |
+
```
|
| 77 |
+
|
| 78 |
The example puts a barkeep, an elf mercenary and an old sailor in a room, with the player facing the elf. It asks who four player lines are for, with and without the spatial facts. Then the barkeep speaks, after the player asked him where the mill key is:
|
| 79 |
|
| 80 |
```
|
|
|
|
| 90 |
Tomas: 'Wren, were you there?' text only wren (1.00)
|
| 91 |
```
|
| 92 |
|
| 93 |
+
`python example.py --fp16` prints the same lines.
|
| 94 |
+
|
| 95 |
In your own code, with the downloaded folder as the working directory:
|
| 96 |
|
| 97 |
```python
|
|
|
|
| 102 |
from contracts.schemas.person_card import PersonCard
|
| 103 |
from harness.addressee.classifier import load_classifier
|
| 104 |
|
| 105 |
+
model = load_classifier(".") # full precision: encoder.onnx, head.onnx, tokenizer.json and config.json
|
| 106 |
+
# model = load_classifier(".", precision="fp16") # the 16-bit files instead: encoder_fp16.onnx and head_fp16.onnx
|
| 107 |
people = [
|
| 108 |
PersonCard(id="tomas", label="Tomas", aliases=["the barkeep"], features=[{"key": "headwear", "value": "red hat"}]),
|
| 109 |
PersonCard(id="wren", label="Wren", features=[{"key": "role", "value": "mercenary"}]),
|
|
|
|
| 176 |
- **Spatial facts depend on the application's band edges.** Tune them (`BandConfig` in `bands.py`).
|
| 177 |
- **It has not met real users yet.** It was tested in game settings only, on generated and hand-written lines, not on logs from real people or from other multi-character voice experiences.
|
| 178 |
|
| 179 |
+
The numbers come from the training repository's reports for this release: `2026-10-09-addr-r3.md` and `2026-10-09-addr-r3-setup.md` at commit `fce88bf`, and for the 16-bit files `2026-10-09-addr-precision.md` at commit `28e72a8`. The full-precision files are RC2's, trained and exported at git tag `spellspeak-audience-rc2` (commit `97f3320`), and this package was built at git tag `spellspeak-audience-rc2.2`. That repository is private.
|
encoder_fp16.onnx
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:9d71e550b8e497ed1c96dd6992f5133dc2f244fa708367cc5d024478a11ef7e2
|
| 3 |
+
size 64103200
|
example.py
CHANGED
|
@@ -1,11 +1,13 @@
|
|
| 1 |
"""Ask SpellSpeak Audience who a line is for.
|
| 2 |
|
| 3 |
pip install -r requirements.txt
|
| 4 |
-
python example.py
|
|
|
|
| 5 |
|
| 6 |
-
The model files (
|
| 7 |
-
|
| 8 |
"""
|
|
|
|
| 9 |
import pathlib
|
| 10 |
import sys
|
| 11 |
|
|
@@ -17,13 +19,17 @@ from contracts.schemas.person_card import PersonCard # noqa: E402
|
|
| 17 |
from harness.addressee.bands import Pose, bands # noqa: E402
|
| 18 |
from harness.addressee.classifier import load_classifier # noqa: E402
|
| 19 |
|
| 20 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
| 21 |
model_dir = HERE
|
| 22 |
if not all((HERE / f).exists() for f in MODEL_FILES):
|
| 23 |
from huggingface_hub import snapshot_download
|
| 24 |
|
| 25 |
-
model_dir = snapshot_download("spellspeak/audience", revision="rc2", allow_patterns=MODEL_FILES)
|
| 26 |
-
model = load_classifier(model_dir, threads=4)
|
| 27 |
|
| 28 |
# What the player can see of each person. Keys are open: a game uses its own.
|
| 29 |
people = [
|
|
|
|
| 1 |
"""Ask SpellSpeak Audience who a line is for.
|
| 2 |
|
| 3 |
pip install -r requirements.txt
|
| 4 |
+
python example.py # full precision: encoder.onnx and head.onnx
|
| 5 |
+
python example.py --fp16 # the 16-bit files, where the release has them: half the download, the same answers
|
| 6 |
|
| 7 |
+
The model files (the two graphs, tokenizer.json and config.json) are read from this folder. Without them, as in a git
|
| 8 |
+
checkout, the ones the chosen precision needs are downloaded once from Hugging Face (spellspeak/audience, release rc2.1).
|
| 9 |
"""
|
| 10 |
+
import argparse
|
| 11 |
import pathlib
|
| 12 |
import sys
|
| 13 |
|
|
|
|
| 19 |
from harness.addressee.bands import Pose, bands # noqa: E402
|
| 20 |
from harness.addressee.classifier import load_classifier # noqa: E402
|
| 21 |
|
| 22 |
+
parser = argparse.ArgumentParser(description="Ask SpellSpeak Audience who a line is for.")
|
| 23 |
+
parser.add_argument("--fp16", action="store_true", help="use the 16-bit files: half the download, the same answers")
|
| 24 |
+
PRECISION = "fp16" if parser.parse_args().fp16 else "fp32"
|
| 25 |
+
GRAPHS = ["encoder.onnx", "head.onnx"] if PRECISION == "fp32" else ["encoder_fp16.onnx", "head_fp16.onnx"]
|
| 26 |
+
MODEL_FILES = [*GRAPHS, "tokenizer.json", "config.json"]
|
| 27 |
model_dir = HERE
|
| 28 |
if not all((HERE / f).exists() for f in MODEL_FILES):
|
| 29 |
from huggingface_hub import snapshot_download
|
| 30 |
|
| 31 |
+
model_dir = snapshot_download("spellspeak/audience", revision="rc2.1", allow_patterns=MODEL_FILES)
|
| 32 |
+
model = load_classifier(model_dir, threads=4, precision=PRECISION)
|
| 33 |
|
| 34 |
# What the player can see of each person. Keys are open: a game uses its own.
|
| 35 |
people = [
|
head_fp16.onnx
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:78361e61605406bb0d8c13de148847684e55724fe183f24a713df2c4cee18655
|
| 3 |
+
size 1858520
|
runtime/harness/addressee/classifier.py
CHANGED
|
@@ -16,6 +16,25 @@ from contracts.schemas.addressee import AddresseeAnswer, AddresseeRequest
|
|
| 16 |
from harness.addressee.render import RENDER_VERSION, render, span_mask
|
| 17 |
|
| 18 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 19 |
def _sigmoid(x: float) -> float:
|
| 20 |
return 1.0 / (1.0 + math.exp(-x)) if x >= 0 else math.exp(x) / (1.0 + math.exp(x))
|
| 21 |
|
|
@@ -69,7 +88,7 @@ class CachedCardsClassifier:
|
|
| 69 |
version it was trained on (its config's `render_version`): rc1 reads addressee-tt-0.5, rc2
|
| 70 |
addressee-tt-0.6, which also takes lines spoken by characters."""
|
| 71 |
|
| 72 |
-
def __init__(self, model_dir: str | pathlib.Path, threads: int = 4):
|
| 73 |
import numpy as np
|
| 74 |
import onnxruntime as ort
|
| 75 |
from tokenizers import Tokenizer
|
|
@@ -78,6 +97,8 @@ class CachedCardsClassifier:
|
|
| 78 |
|
| 79 |
self.np = np
|
| 80 |
d = pathlib.Path(model_dir)
|
|
|
|
|
|
|
| 81 |
self.config = json.loads((d / "config.json").read_text())
|
| 82 |
self.version = self.config.get("render_version")
|
| 83 |
if self.version not in TT_VERSIONS:
|
|
@@ -91,8 +112,8 @@ class CachedCardsClassifier:
|
|
| 91 |
so = ort.SessionOptions()
|
| 92 |
so.intra_op_num_threads = threads
|
| 93 |
so.inter_op_num_threads = 1
|
| 94 |
-
self.enc = ort.InferenceSession(str(
|
| 95 |
-
self.head = ort.InferenceSession(str(
|
| 96 |
self.temps = self.config["temperatures"]
|
| 97 |
self.cache: dict[str, tuple] = {}
|
| 98 |
|
|
@@ -143,7 +164,13 @@ class CachedCardsClassifier:
|
|
| 143 |
return a, (time.perf_counter() - t0) * 1000
|
| 144 |
|
| 145 |
|
| 146 |
-
def load_classifier(model_dir: str | pathlib.Path, threads: int = 4):
|
| 147 |
-
"""The right runtime for an exported model folder: cached cards when it has a card encoder, else single pass.
|
|
|
|
|
|
|
| 148 |
d = pathlib.Path(model_dir)
|
| 149 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 16 |
from harness.addressee.render import RENDER_VERSION, render, span_mask
|
| 17 |
|
| 18 |
|
| 19 |
+
PRECISIONS = ("fp32", "fp16")
|
| 20 |
+
|
| 21 |
+
|
| 22 |
+
def graph_files(model_dir: str | pathlib.Path, precision: str = "fp32") -> tuple[pathlib.Path, pathlib.Path]:
|
| 23 |
+
"""The cached-cards model's two graphs for a precision: `encoder.onnx` and `head.onnx` (full precision, the default
|
| 24 |
+
and the reference), or `encoder_fp16.onnx` and `head_fp16.onnx` (16-bit weights, cast to fp32 when they load: half
|
| 25 |
+
the file, the same answers). A clear error when the chosen files are not in the folder."""
|
| 26 |
+
if precision not in PRECISIONS:
|
| 27 |
+
raise ValueError(f"precision must be one of {', '.join(PRECISIONS)}, not {precision!r}")
|
| 28 |
+
d = pathlib.Path(model_dir)
|
| 29 |
+
suffix = "" if precision == "fp32" else "_fp16"
|
| 30 |
+
enc, head = d / f"encoder{suffix}.onnx", d / f"head{suffix}.onnx"
|
| 31 |
+
missing = [p.name for p in (enc, head) if not p.exists()]
|
| 32 |
+
if missing:
|
| 33 |
+
raise FileNotFoundError(f"precision {precision!r} needs {' and '.join(missing)} in {d}; "
|
| 34 |
+
f"download them, or use precision='fp32' with encoder.onnx and head.onnx")
|
| 35 |
+
return enc, head
|
| 36 |
+
|
| 37 |
+
|
| 38 |
def _sigmoid(x: float) -> float:
|
| 39 |
return 1.0 / (1.0 + math.exp(-x)) if x >= 0 else math.exp(x) / (1.0 + math.exp(x))
|
| 40 |
|
|
|
|
| 88 |
version it was trained on (its config's `render_version`): rc1 reads addressee-tt-0.5, rc2
|
| 89 |
addressee-tt-0.6, which also takes lines spoken by characters."""
|
| 90 |
|
| 91 |
+
def __init__(self, model_dir: str | pathlib.Path, threads: int = 4, precision: str = "fp32"):
|
| 92 |
import numpy as np
|
| 93 |
import onnxruntime as ort
|
| 94 |
from tokenizers import Tokenizer
|
|
|
|
| 97 |
|
| 98 |
self.np = np
|
| 99 |
d = pathlib.Path(model_dir)
|
| 100 |
+
enc_file, head_file = graph_files(d, precision)
|
| 101 |
+
self.precision = precision
|
| 102 |
self.config = json.loads((d / "config.json").read_text())
|
| 103 |
self.version = self.config.get("render_version")
|
| 104 |
if self.version not in TT_VERSIONS:
|
|
|
|
| 112 |
so = ort.SessionOptions()
|
| 113 |
so.intra_op_num_threads = threads
|
| 114 |
so.inter_op_num_threads = 1
|
| 115 |
+
self.enc = ort.InferenceSession(str(enc_file), so, providers=["CPUExecutionProvider"])
|
| 116 |
+
self.head = ort.InferenceSession(str(head_file), so, providers=["CPUExecutionProvider"])
|
| 117 |
self.temps = self.config["temperatures"]
|
| 118 |
self.cache: dict[str, tuple] = {}
|
| 119 |
|
|
|
|
| 164 |
return a, (time.perf_counter() - t0) * 1000
|
| 165 |
|
| 166 |
|
| 167 |
+
def load_classifier(model_dir: str | pathlib.Path, threads: int = 4, precision: str = "fp32"):
|
| 168 |
+
"""The right runtime for an exported model folder: cached cards when it has a card encoder, else single pass.
|
| 169 |
+
`precision`: "fp32" (the default) reads encoder.onnx and head.onnx; "fp16" reads encoder_fp16.onnx and
|
| 170 |
+
head_fp16.onnx, in a folder that may hold only those (rc2.1 on)."""
|
| 171 |
d = pathlib.Path(model_dir)
|
| 172 |
+
if (d / "head.onnx").exists() or (d / "head_fp16.onnx").exists():
|
| 173 |
+
return CachedCardsClassifier(d, threads, precision)
|
| 174 |
+
if precision != "fp32":
|
| 175 |
+
raise ValueError(f"precision {precision!r} needs a cached-cards model (encoder and head); {d} holds a single-pass model.onnx")
|
| 176 |
+
return AddresseeClassifier(d, threads=threads)
|