File size: 9,252 Bytes
0d058cf 6a274a1 0d058cf 6a274a1 4a52986 0b3a005 0d058cf 0b3a005 0d058cf 4a52986 0d058cf 4a52986 0d058cf 4a52986 0d058cf 4a52986 0d058cf 6a274a1 0d058cf 6a274a1 0d058cf 6a274a1 0d058cf 6a274a1 0d058cf 4a52986 0d058cf 4a52986 0d058cf 6a274a1 | 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 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 | """Run the quantized ONNX models from this repository with ONNX Runtime.
Needs only `onnxruntime`, `tokenizers` and `numpy` (plus `huggingface_hub` to download by repo id).
from inference import load, maxsim
emb = load("nynefold/pplx-embed-v1-0.6b-ONNX") # recommended variant
q = emb.embed_queries(["how do I reset the router?"])
d = emb.embed_documents(["Hold the reset button for ten seconds.", "The router has four ports."])
print(q @ d.T) # cosine similarities
rr = load("nynefold/zerank-1-small-ONNX", variant="mix30-gptq-sq50-a4")
print(rr.rerank("how do I reset the router?", ["Hold the reset button ...", "The router has ..."]))
late = load("nynefold/pplx-embed-v2-late-0.6b-ONNX") # late interaction: one vector per token
docs = late.embed_documents(["Hold the reset button for ten seconds.", "The router has four ports."])
print(maxsim(late.embed_queries(["how do I reset the router?"])[0], docs))
`load` takes a repo id or a local folder. A repo id downloads the variant into
~/.cache/nynefold/<repo id> (or $NYNEFOLD_CACHE/<repo id>) as plain files. All model-specific details (prompts, truncation, output
transform) come from `config.json`, which is part of the contract: the graphs expect
exactly these inputs.
"""
import json
import os
from pathlib import Path
import numpy as np
import onnxruntime as ort
from tokenizers import Tokenizer
CONFIG = "config.json" # also the file the Hub counts downloads by
def _local(repo_or_dir, variant):
p = Path(repo_or_dir)
if p.is_dir():
return p
# Plain files in a local folder, not the Hub cache: its symlinks put model.onnx and model.onnx.data
# in different blob folders, and ONNX Runtime refuses external data outside the model's folder.
from huggingface_hub import hf_hub_download, snapshot_download
dst = Path(os.environ.get("NYNEFOLD_CACHE", Path.home() / ".cache" / "nynefold")) / repo_or_dir
cfg = json.loads(Path(hf_hub_download(repo_or_dir, CONFIG, local_dir=dst)).read_text())
v = variant or cfg["recommended_variant"]
return Path(snapshot_download(repo_or_dir, local_dir=dst,
allow_patterns=["*.json", "*.txt", "*.jinja", f"{v}/*"]))
def load(repo_or_dir, variant=None, threads=None):
"""A ready model: Embedder or Reranker. variant=None picks the repo's recommended variant."""
root = _local(repo_or_dir, variant)
cfg = json.loads((root / CONFIG).read_text())
variant = variant or cfg["recommended_variant"]
vdir = root / variant
tok_file = vdir / "tokenizer.json" if (vdir / "tokenizer.json").exists() else root / "tokenizer.json"
so = ort.SessionOptions()
if threads: # otherwise ORT's default: one thread per physical core
so.intra_op_num_threads = threads
so.inter_op_num_threads = 1
so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
sess = ort.InferenceSession(str(vdir / "model.onnx"), so, providers=["CPUExecutionProvider"])
cls = {"embed": Embedder, "rerank": Reranker, "late": LateEmbedder}[cfg["kind"]]
return cls(cfg, Tokenizer.from_file(str(tok_file)), sess)
class _Model:
def __init__(self, cfg, tok, sess):
self.cfg, self.tok, self.sess = cfg, tok, sess
self.max_len = cfg["max_length"]
self.pad_id = tok.token_to_id(cfg["pad_token"])
def _ids(self, text, special=True):
return self.tok.encode(text, add_special_tokens=special).ids
def _run(self, ids, max_tokens=8192, max_batch=64, per_token=False):
"""Right-padded, length-sorted batches through the graph; one output row per input
(per_token: a list with one [tokens, dim] array per input, padding removed)."""
out = [None] * len(ids)
order = sorted(range(len(ids)), key=lambda i: -len(ids[i]))
i = 0
while i < len(order):
T = len(ids[order[i]])
idx = order[i : i + max(1, min(max_batch, max_tokens // max(T, 1)))]
inp = np.full((len(idx), T), self.pad_id, dtype=np.int64)
msk = np.zeros((len(idx), T), dtype=np.int64)
for r, j in enumerate(idx):
inp[r, : len(ids[j])] = ids[j]
msk[r, : len(ids[j])] = 1
y = self.sess.run(None, {"input_ids": inp, "attention_mask": msk})[0]
for r, j in enumerate(idx):
out[j] = y[r, : len(ids[j])] if per_token else y[r]
i += len(idx)
return out if per_token else np.stack(out)
class Embedder(_Model):
"""Pooling is inside the graph. Returns L2-normalized float32 vectors; compare by dot product."""
def encode_ids(self, texts, role):
prefix = self.cfg["query_prefix"] if role == "query" else self.cfg["document_prefix"]
out = []
for t in texts:
x = self._ids(prefix + t, special=self.cfg["add_special_tokens"])
if len(x) > self.max_len:
# last-token pooling reads the final (end) token, so keep it
x = x[: self.max_len - 1] + x[-1:] if self.cfg["pooling"] == "last" else x[: self.max_len]
out.append(x)
return out
def _embed(self, texts, role):
e = self._run(self.encode_ids(list(texts), role)).astype(np.float32)
if self.cfg["output"] == "int8_tanh": # the base model's native int8 output: round(127 * tanh(x))
e = np.clip(np.round(127 * np.tanh(e)), -128, 127)
return e / np.linalg.norm(e, axis=1, keepdims=True)
def embed_queries(self, texts):
return self._embed(texts, "query")
def embed_documents(self, texts):
return self._embed(texts, "document")
class Reranker(_Model):
"""Scores (query, document) pairs; higher is more relevant."""
def encode_ids(self, pairs):
c = self.cfg
out = []
if c["format"] == "chat": # prompt from the template; if too long, only the document's tokens are cut
head, tail = c["template"].split("{document}")
suf = self._ids(tail, False)
for q, d in pairs:
q, d = q[: c["query_max_chars"]].strip(), d[: c["document_max_chars"]].strip()
x = self._ids(head.format(query=q) + d + tail, special=False)
if len(x) > self.max_len: # keep the end of the template: the score is read at its last token
pre = self._ids(head.format(query=q), False)
x = pre + self._ids(d, False)[: max(0, self.max_len - len(pre) - len(suf))] + suf
out.append(x)
else: # fixed prefix and suffix around a body; only the body is truncated
pre, suf = self._ids(c["prefix"], False), self._ids(c["suffix"], False)
room = self.max_len - len(pre) - len(suf)
for q, d in pairs:
body = self._ids(c["body"].format(instruction=c["instruction"], query=q, document=d), False)
out.append(pre + body[:room] + suf)
return out
def logits(self, query, documents):
return self._run(self.encode_ids([(query, d) for d in documents])).reshape(-1).astype(np.float32)
def rerank(self, query, documents):
"""Relevance scores in [0, 1], one per document, in input order."""
z = self.logits(query, documents) * self.cfg["score_scale"]
return 1.0 / (1.0 + np.exp(-z))
class LateEmbedder(_Model):
"""Late interaction (ColBERT-style, pplx-embed-v2-late): one L2-normalized 128-d vector per token.
Score a query against documents with `maxsim`. Models that share a token space (the 0.6B and 9B)
can be mixed: index documents with one, embed queries with the other."""
def __init__(self, cfg, tok, sess):
super().__init__(cfg, tok, sess)
self.skip = np.array(sorted({i for i in (tok.token_to_id(t) for t in cfg["document_skiplist"]) if i is not None}))
def encode_ids(self, texts, role):
c = self.cfg
prefix, n = ((c["query_prefix"], c["query_max_length"]) if role == "query"
else (c["document_prefix"], c["document_max_length"]))
return [self._ids(prefix + t, special=c["add_special_tokens"])[:n] for t in texts]
def _embed(self, texts, role):
ids = self.encode_ids(list(texts), role)
out = [v.astype(np.float32) for v in self._run(ids, per_token=True)]
if role == "document": # punctuation tokens don't take part in scoring
out = [v[~np.isin(x, self.skip)] for v, x in zip(out, ids)]
return out
def embed_queries(self, texts):
"""A list of [query tokens, 128] float32 arrays."""
return self._embed(texts, "query")
def embed_documents(self, texts):
"""A list of [document tokens, 128] float32 arrays (punctuation tokens removed): the index."""
return self._embed(texts, "document")
def maxsim(query, documents):
"""Late-interaction scores of one query ([tokens, dim]) against documents (a list of [tokens, dim]):
for each query token the best-matching document token's dot product, summed. Higher is more relevant."""
return np.array([float((query @ d.T).max(1).sum()) for d in documents], dtype=np.float32)
|