passageway 0.3.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
passageway/__init__.py ADDED
File without changes
passageway/chunk.py ADDED
@@ -0,0 +1,68 @@
1
+ """Turns an item into the pieces that are embedded and returned as passages."""
2
+
3
+ import re
4
+ from dataclasses import dataclass
5
+
6
+ from passageway.protocol import Item
7
+
8
+ MAX_CHARS = 1500 # a section over this is split; the embedding model cuts at a few hundred tokens anyway
9
+
10
+
11
+ @dataclass(frozen=True)
12
+ class Chunk:
13
+ heading: str
14
+ text: str
15
+ ordinal: int
16
+
17
+
18
+ def _cut_words(paragraph: str, cap: int) -> list[str]:
19
+ out, line = [], ""
20
+ for word in paragraph.split():
21
+ while len(word) > cap:
22
+ if line:
23
+ out.append(line)
24
+ line = ""
25
+ out.append(word[:cap])
26
+ word = word[cap:]
27
+ if line and len(line) + 1 + len(word) > cap:
28
+ out.append(line)
29
+ line = word
30
+ else:
31
+ line = f"{line} {word}" if line else word
32
+ if line:
33
+ out.append(line)
34
+ return out
35
+
36
+
37
+ def split_text(text: str, cap: int = MAX_CHARS) -> list[str]:
38
+ """Pieces of at most `cap` characters, cut at blank lines first, then between words."""
39
+ text = text.strip()
40
+ if len(text) <= cap:
41
+ return [text] if text else []
42
+ pieces, current = [], ""
43
+ for paragraph in re.split(r"\n\s*\n", text):
44
+ paragraph = paragraph.strip()
45
+ for part in [paragraph] if len(paragraph) <= cap else _cut_words(paragraph, cap):
46
+ if not part:
47
+ continue
48
+ if current and len(current) + 2 + len(part) > cap:
49
+ pieces.append(current)
50
+ current = part
51
+ else:
52
+ current = f"{current}\n\n{part}" if current else part
53
+ if current:
54
+ pieces.append(current)
55
+ return pieces
56
+
57
+
58
+ def chunks(item: Item) -> list[Chunk]:
59
+ out: list[Chunk] = []
60
+ for section in item.sections:
61
+ for piece in split_text(section.text):
62
+ out.append(Chunk(section.heading or "", piece, len(out)))
63
+ return out
64
+
65
+
66
+ def passage_text(item: Item, chunk: Chunk) -> str:
67
+ trail = " › ".join([*item.path, *([chunk.heading] if chunk.heading else [])])
68
+ return f"{item.title} › {trail}: {chunk.text}" if trail else f"{item.title}: {chunk.text}"
passageway/config.py ADDED
@@ -0,0 +1,315 @@
1
+ """Configuration: environment variables for the core, one TOML file for connectors (docs/install.md)."""
2
+
3
+ import math
4
+ import re
5
+ import sys
6
+ import threading
7
+ import tomllib
8
+ from collections.abc import Mapping
9
+ from dataclasses import dataclass, field
10
+ from pathlib import Path
11
+
12
+ from mcp.server.auth.routes import validate_issuer_url
13
+ from pydantic import AnyHttpUrl
14
+
15
+ from passageway.protocol import MAX_LINE_BYTES, MAX_OUTPUT_BYTES
16
+
17
+
18
+ class ConfigError(Exception):
19
+ pass
20
+
21
+
22
+ DEFAULT_TOOL_DESCRIPTION = (
23
+ "Search the user's indexed sources (documents, notes, decisions, how their systems and tools actually "
24
+ "behave) and get up to 3 short passages with links. Call it on your own before a task the sources may "
25
+ "cover, when you hit an unfamiliar error, host or service, and before making a decision they may already "
26
+ "have made. Returns nothing-relevant when there is no match. Open the link for the full document.")
27
+
28
+ DEFAULT_INSTRUCTIONS = (
29
+ "This server searches the user's indexed sources: their documents, notes, decisions and how their systems "
30
+ "and tools actually behave. Call `recall` before starting a task the sources may cover, when you hit an "
31
+ "unfamiliar error, host or service, and before making a decision they may already have made, even when you "
32
+ "think you know the answer: the sources hold what is true for this user. A hit is a pointer: open its link "
33
+ "for the full document before relying on it. If the context already holds passages headed \"Recall results\", "
34
+ "those came from this server; don't search for the same thing again.")
35
+
36
+ _NAME = re.compile(r"[a-z0-9][a-z0-9_-]*")
37
+
38
+
39
+ @dataclass(frozen=True)
40
+ class ConnectorConfig:
41
+ name: str
42
+ command: tuple[str, ...]
43
+ env: dict[str, str]
44
+ secrets: dict[str, str] # already resolved from env:NAME references
45
+ timeout: float = 300.0
46
+ max_line_bytes: int = MAX_LINE_BYTES
47
+ max_output_bytes: int = MAX_OUTPUT_BYTES
48
+ max_delete_fraction: float = 0.3 # a reconcile that would delete more than this share of a source is refused
49
+
50
+
51
+ EMBEDDER_KINDS = ("onnx", "command")
52
+ _ONNX_KEYS = {"kind", "min_similarity", "model", "model_dir", "device", "download"}
53
+ _COMMAND_KEYS = {"kind", "command", "env", "secrets", "min_similarity", "batch_size", "timeout"}
54
+
55
+
56
+ @dataclass(frozen=True)
57
+ class EmbedderConfig:
58
+ """How texts become vectors: the built-in ONNX model, or any program speaking the embedder protocol."""
59
+ kind: str = "onnx"
60
+ command: tuple[str, ...] = ()
61
+ env: dict[str, str] = field(default_factory=dict)
62
+ secrets: dict[str, str] = field(default_factory=dict) # already resolved from env:NAME references
63
+ min_similarity: float | None = None # the relevance gate of this model; None keeps the built-in default
64
+ batch_size: int = 32 # texts per request
65
+ timeout: float = 60.0 # seconds per request
66
+ model_dir: Path | None = None # onnx only; None: <data dir>/models/<model>
67
+ model: str = "multilingual-e5-small" # onnx only: a preset of passageway.models
68
+ device: str = "auto" # onnx only: auto | cpu | coreml
69
+ download: bool = True # onnx only: fetch the pinned, checksummed files when they are missing
70
+
71
+
72
+ @dataclass
73
+ class Config:
74
+ password: str # PASSAGEWAY_LOGIN_PASSWORD: static Bearer token and the OAuth login page's password
75
+ connectors: list[ConnectorConfig]
76
+ data_dir: Path
77
+ model_dir: Path
78
+ base_url: str | None = None # public URL; set, it turns on OAuth for /mcp
79
+ host: str = "127.0.0.1"
80
+ port: int = 8082
81
+ poll_seconds: int = 60
82
+ reconcile_seconds: int = 1800
83
+ log_level: str = "INFO"
84
+ tool_description: str = DEFAULT_TOOL_DESCRIPTION
85
+ instructions: str = DEFAULT_INSTRUCTIONS
86
+ instance: str | None = None # PASSAGEWAY_INSTANCE: one of several servers on this machine
87
+ models_root: Path | None = None # where the built-in models live; None: data_dir (instances share the default)
88
+ min_similarity: float | None = None # the relevance gate; None keeps the embedder preset's value
89
+ embedder: EmbedderConfig = field(default_factory=EmbedderConfig) # the one the server uses
90
+ embedders: dict[str, EmbedderConfig] = field(default_factory=dict) # candidates for `passageway eval compare`
91
+
92
+
93
+ _INSTANCE = r"[a-z0-9][a-z0-9-]*"
94
+
95
+
96
+ def instance_name(env: Mapping[str, str]) -> str | None:
97
+ """PASSAGEWAY_INSTANCE names one of several servers on this machine; unset, the one and only."""
98
+ name = env.get("PASSAGEWAY_INSTANCE", "").strip()
99
+ if name and not re.fullmatch(_INSTANCE, name):
100
+ raise ConfigError(f"PASSAGEWAY_INSTANCE must be lowercase letters, digits and hyphens, got {name!r}")
101
+ return name or None
102
+
103
+
104
+ def take_instance_argument(argv: list[str]) -> tuple[list[str], str | None]:
105
+ """`--instance NAME` or `--instance=NAME`, anywhere in the command line: the rest of it and the name."""
106
+ rest, name, args = [], None, iter(argv)
107
+ for arg in args:
108
+ if arg == "--instance":
109
+ name = next(args, "")
110
+ elif arg.startswith("--instance="):
111
+ name = arg.split("=", 1)[1]
112
+ else:
113
+ rest.append(arg)
114
+ continue
115
+ if not re.fullmatch(_INSTANCE, name):
116
+ raise ConfigError(f"--instance needs a name of lowercase letters, digits and hyphens, got {name!r}")
117
+ return rest, name
118
+
119
+
120
+ def existing_instances() -> list[str]:
121
+ return sorted(p.parent.name.removeprefix("passageway-") for p in (Path.home() / ".config").glob("passageway-*/passageway.toml"))
122
+
123
+
124
+ def instance_problem(env: Mapping[str, str]) -> str | None:
125
+ """Why this command would run against an instance that is not there, or None. A named instance must have
126
+ been created by init; with none named and no default config, the named ones are what the user meant."""
127
+ if env.get("PASSAGEWAY_CONFIG"):
128
+ return None
129
+ name, known = instance_name(env), existing_instances()
130
+ listed = ", ".join(known) or "none"
131
+ if name and not default_config_path(name).exists():
132
+ return (f"there is no passageway instance {name!r} (instances on this machine: {listed}); "
133
+ f"create it with: passageway --instance {name} init")
134
+ if not name and known and not default_config_path().exists():
135
+ return f"this machine has named passageway instances ({listed}): pick one with --instance NAME"
136
+ return None
137
+
138
+
139
+ def _suffix(instance: str | None) -> str:
140
+ return f"-{instance}" if instance else ""
141
+
142
+
143
+ def default_config_path(instance: str | None = None) -> Path:
144
+ return Path.home() / ".config" / f"passageway{_suffix(instance)}" / "passageway.toml"
145
+
146
+
147
+ def default_data_dir(instance: str | None = None) -> Path:
148
+ """The data folder; an instance gets its own next to the shared one, which keeps the built-in models."""
149
+ base = Path.home() / "Library" / "Application Support" if sys.platform == "darwin" else Path.home() / ".local" / "share"
150
+ return base / f"passageway{_suffix(instance)}"
151
+
152
+
153
+ def _int(env: Mapping[str, str], key: str, default: int) -> int:
154
+ raw = env.get(key)
155
+ if not raw:
156
+ return default
157
+ try:
158
+ return int(raw)
159
+ except ValueError:
160
+ raise ConfigError(f"{key} must be an integer, got {raw!r}") from None
161
+
162
+
163
+ def _resolve(connector: str, ref: object, env: Mapping[str, str], what: str = "connector") -> str:
164
+ if not isinstance(ref, str) or not ref.startswith("env:") or len(ref) == 4:
165
+ raise ConfigError(f"{what} {connector!r}: a secret must be a reference like \"env:NAME\" (in 'secrets')")
166
+ var = ref[4:]
167
+ if not env.get(var):
168
+ raise ConfigError(f"{what} {connector!r}: the environment variable {var} is not set")
169
+ return env[var]
170
+
171
+
172
+ def _table(connector: str, raw: dict, key: str, what: str = "connector") -> dict[str, str]:
173
+ table = raw.get(key, {})
174
+ if not isinstance(table, dict) or not all(isinstance(v, str) for v in table.values()):
175
+ raise ConfigError(f"{what} {connector!r}: '{key}' must be a table of strings")
176
+ return table
177
+
178
+
179
+ def _embedder(name: str, raw: object, env: Mapping[str, str]) -> EmbedderConfig:
180
+ where = f"embedder {name!r}"
181
+ if not isinstance(raw, dict):
182
+ raise ConfigError(f"{where}: must be a table")
183
+ kind = raw.get("kind", "onnx")
184
+ if kind not in EMBEDDER_KINDS:
185
+ raise ConfigError(f"{where}: 'kind' must be one of {', '.join(EMBEDDER_KINDS)}, got {kind!r}")
186
+ allowed = _ONNX_KEYS if kind == "onnx" else _COMMAND_KEYS
187
+ for key in raw:
188
+ if key not in allowed:
189
+ raise ConfigError(f"{where}: unknown key {key!r} for kind {kind!r} (known: {', '.join(sorted(allowed))})")
190
+ gate = raw.get("min_similarity")
191
+ if gate is not None:
192
+ if isinstance(gate, bool) or not isinstance(gate, (int, float)) or not 0 <= gate <= 1:
193
+ raise ConfigError(f"{where}: 'min_similarity' must be a number from 0 to 1")
194
+ gate = float(gate)
195
+ if kind == "onnx":
196
+ from passageway.models import DEVICES, PRESETS
197
+
198
+ model_dir = raw.get("model_dir")
199
+ if model_dir is not None and not isinstance(model_dir, str):
200
+ raise ConfigError(f"{where}: 'model_dir' must be a path")
201
+ model = raw.get("model", "multilingual-e5-small")
202
+ if not isinstance(model, str) or model not in PRESETS:
203
+ raise ConfigError(f"{where}: 'model' must be one of the presets {', '.join(PRESETS)}, got {model!r} "
204
+ "(another model runs as kind = \"command\")")
205
+ device = raw.get("device", "auto")
206
+ if device not in DEVICES:
207
+ raise ConfigError(f"{where}: 'device' must be one of {', '.join(DEVICES)}, got {device!r}")
208
+ auto = raw.get("download", True)
209
+ if not isinstance(auto, bool):
210
+ raise ConfigError(f"{where}: 'download' must be true or false")
211
+ return EmbedderConfig("onnx", min_similarity=gate, model_dir=Path(model_dir) if model_dir else None, model=model,
212
+ device=device, download=auto)
213
+ command = raw.get("command")
214
+ if not isinstance(command, list) or not command or not all(isinstance(c, str) for c in command):
215
+ raise ConfigError(f"{where}: 'command' must be a non-empty list of strings")
216
+ secrets = {k: _resolve(name, v, env, "embedder") for k, v in _table(name, raw, "secrets", "embedder").items()}
217
+ batch = raw.get("batch_size", 32)
218
+ if isinstance(batch, bool) or not isinstance(batch, int) or not 1 <= batch <= 1024:
219
+ raise ConfigError(f"{where}: 'batch_size' must be a whole number from 1 to 1024")
220
+ timeout = raw.get("timeout", 60.0)
221
+ if (isinstance(timeout, bool) or not isinstance(timeout, (int, float)) or not math.isfinite(timeout) or timeout <= 0
222
+ or timeout > threading.TIMEOUT_MAX):
223
+ raise ConfigError(f"{where}: 'timeout' must be a number of seconds greater than 0 (and not absurdly large)")
224
+ return EmbedderConfig("command", tuple(command), _table(name, raw, "env", "embedder"), secrets, gate, batch, float(timeout))
225
+
226
+
227
+ def _number(connector: str, raw: dict, key: str, default, *, kind=float, low=None, high=None):
228
+ value = raw.get(key, default)
229
+ ok = (isinstance(value, (int, float)) and not isinstance(value, bool) and (kind is float or isinstance(value, int))
230
+ and math.isfinite(value))
231
+ if not ok or (low is not None and value <= low) or (high is not None and value > high):
232
+ bound = f" (greater than {low}{'' if high is None else f' and at most {high}'})" if low is not None else ""
233
+ raise ConfigError(f"connector {connector!r}: '{key}' must be a number{bound}")
234
+ return kind(value)
235
+
236
+
237
+ def _load_toml(path: Path, env: Mapping[str, str]) -> tuple[list[ConnectorConfig], EmbedderConfig, dict[str, EmbedderConfig]]:
238
+ try:
239
+ data = tomllib.loads(path.read_text())
240
+ except FileNotFoundError:
241
+ raise ConfigError(f"{path} not found: list your connectors in a passageway.toml (see docs/install.md)") from None
242
+ except tomllib.TOMLDecodeError as e:
243
+ raise ConfigError(f"{path}: {e}") from None
244
+ connectors: list[ConnectorConfig] = []
245
+ entries = data.get("connector", [])
246
+ if not isinstance(entries, list) or not all(isinstance(e, dict) for e in entries):
247
+ raise ConfigError(f"{path}: connectors are written as [[connector]] tables, one per source")
248
+ for raw in entries:
249
+ name = raw.get("name", "")
250
+ if not isinstance(name, str) or not _NAME.fullmatch(name):
251
+ raise ConfigError(f"connector name {name!r} must match [a-z0-9][a-z0-9_-]* (it names the index file)")
252
+ if any(c.name == name for c in connectors):
253
+ raise ConfigError(f"duplicate connector name {name!r}")
254
+ command = raw.get("command")
255
+ if not isinstance(command, list) or not command or not all(isinstance(c, str) for c in command):
256
+ raise ConfigError(f"connector {name!r}: 'command' must be a non-empty list of strings")
257
+ secrets = {k: _resolve(name, v, env) for k, v in _table(name, raw, "secrets").items()}
258
+ connectors.append(ConnectorConfig(
259
+ name, tuple(command), _table(name, raw, "env"), secrets,
260
+ _number(name, raw, "timeout", 300.0, low=0, high=threading.TIMEOUT_MAX),
261
+ _number(name, raw, "max_line_bytes", MAX_LINE_BYTES, kind=int, low=0),
262
+ _number(name, raw, "max_output_bytes", MAX_OUTPUT_BYTES, kind=int, low=0),
263
+ _number(name, raw, "max_delete_fraction", 0.3, low=0, high=1)))
264
+ if not connectors:
265
+ raise ConfigError(f"no connectors in {path}: add at least one [[connector]]")
266
+ embedder = _embedder("default", data.get("embedder", {}), env)
267
+ candidates = data.get("embedders", {})
268
+ if not isinstance(candidates, dict):
269
+ raise ConfigError(f"{path}: 'embedders' holds the candidates for `passageway eval compare`, one [embedders.NAME] table each")
270
+ embedders: dict[str, EmbedderConfig] = {}
271
+ for name, raw in candidates.items():
272
+ if not _NAME.fullmatch(name):
273
+ raise ConfigError(f"embedder name {name!r} must match [a-z0-9][a-z0-9_-]*")
274
+ if name == "default":
275
+ raise ConfigError("the embedder name 'default' is taken: it stands for the [embedder] the server uses")
276
+ embedders[name] = _embedder(name, raw, env)
277
+ return connectors, embedder, embedders
278
+
279
+
280
+ def load_config(env: Mapping[str, str], *, require_password: bool = True) -> Config:
281
+ """require_password=False is for commands that do not serve (the evaluation): nobody logs in to them."""
282
+ password = env.get("PASSAGEWAY_LOGIN_PASSWORD", "").strip()
283
+ password_file = env.get("PASSAGEWAY_LOGIN_PASSWORD_FILE", "").strip()
284
+ if not password and password_file: # a file the service manager can point at (passageway init writes one)
285
+ try:
286
+ password = Path(password_file).read_text().strip()
287
+ except OSError as e:
288
+ raise ConfigError(f"PASSAGEWAY_LOGIN_PASSWORD_FILE={password_file}: {e.strerror or e}") from None
289
+ if not password:
290
+ raise ConfigError(f"PASSAGEWAY_LOGIN_PASSWORD_FILE={password_file} is empty")
291
+ if require_password and not password:
292
+ raise ConfigError("PASSAGEWAY_LOGIN_PASSWORD is not set (nor PASSAGEWAY_LOGIN_PASSWORD_FILE): clients authenticate "
293
+ "with it, so the engine refuses to start without it; `passageway init` generates one")
294
+ base_url = env.get("PASSAGEWAY_BASE_URL", "").strip().rstrip("/") or None
295
+ if base_url:
296
+ try:
297
+ validate_issuer_url(AnyHttpUrl(base_url))
298
+ except ValueError as e:
299
+ raise ConfigError(f"PASSAGEWAY_BASE_URL={base_url!r} is not a valid OAuth issuer "
300
+ f"(https, no query or fragment; plain http only for localhost): {e}") from None
301
+ instance = instance_name(env)
302
+ config_path = Path(env.get("PASSAGEWAY_CONFIG") or default_config_path(instance))
303
+ connectors, embedder, embedders = _load_toml(config_path, env)
304
+ data_dir = Path(env.get("PASSAGEWAY_DATA_DIR") or default_data_dir(instance))
305
+ models_root = Path(env.get("PASSAGEWAY_DATA_DIR") or default_data_dir())
306
+ return Config(
307
+ password=password, connectors=connectors, data_dir=data_dir, instance=instance, models_root=models_root,
308
+ model_dir=Path(env.get("PASSAGEWAY_MODEL_DIR") or models_root / "models" / "e5-small"),
309
+ base_url=base_url, host=env.get("PASSAGEWAY_HOST", "127.0.0.1"),
310
+ port=_int(env, "PASSAGEWAY_PORT", 8082), poll_seconds=_int(env, "PASSAGEWAY_POLL_SECONDS", 60),
311
+ reconcile_seconds=_int(env, "PASSAGEWAY_RECONCILE_SECONDS", 1800),
312
+ log_level=env.get("PASSAGEWAY_LOG_LEVEL", "INFO"),
313
+ tool_description=env.get("PASSAGEWAY_TOOL_DESCRIPTION", "").strip() or DEFAULT_TOOL_DESCRIPTION,
314
+ instructions=env.get("PASSAGEWAY_INSTRUCTIONS", "").strip() or DEFAULT_INSTRUCTIONS,
315
+ min_similarity=embedder.min_similarity, embedder=embedder, embedders=embedders)
passageway/embed.py ADDED
@@ -0,0 +1,213 @@
1
+ """Text embedders: the built-in ONNX model, any program speaking the embedder protocol, and a hashing stand-in for tests."""
2
+
3
+ import hashlib
4
+ import logging
5
+ import re
6
+ from pathlib import Path
7
+ from typing import Protocol
8
+
9
+ import numpy as np
10
+
11
+ from passageway.config import ConfigError, EmbedderConfig
12
+ from passageway.models import DEFAULT_PRESET, PRESETS, ModelError, download, files_missing, model_folder_for, verify
13
+
14
+ log = logging.getLogger(__name__)
15
+
16
+
17
+ class EmbedderError(Exception):
18
+ """An embedder failed: a program of its own died or answered badly, or the model raised."""
19
+
20
+
21
+ class EmbedderStartError(ConfigError):
22
+ """An embedder could not start. `public` is what an open endpoint may show; the message is for the log."""
23
+
24
+ def __init__(self, message: str, public: str):
25
+ super().__init__(message)
26
+ self.public = public
27
+
28
+
29
+ class Embedder(Protocol):
30
+ dim: int
31
+ name: str
32
+
33
+ def embed_queries(self, texts: list[str]) -> np.ndarray: ...
34
+
35
+ def embed_passages(self, texts: list[str]) -> np.ndarray: ...
36
+
37
+
38
+ def _normalise(v: np.ndarray) -> np.ndarray:
39
+ norms = np.linalg.norm(v, axis=1, keepdims=True)
40
+ norms[norms == 0] = 1
41
+ return (v / norms).astype(np.float32)
42
+
43
+
44
+ def token_batches(lengths: list[int], budget: int) -> list[list[int]]:
45
+ """Group indices so each batch's padded size (count × longest) stays within budget.
46
+
47
+ Memory at inference grows with the padded batch, so long passages go alone or in pairs
48
+ while short ones share a batch. Sorting by length keeps padding small.
49
+ """
50
+ batches: list[list[int]] = []
51
+ current: list[int] = []
52
+ for i in sorted(range(len(lengths)), key=lambda i: lengths[i]):
53
+ if current and (len(current) + 1) * lengths[i] > budget:
54
+ batches.append(current)
55
+ current = []
56
+ current.append(i)
57
+ if current:
58
+ batches.append(current)
59
+ return batches
60
+
61
+
62
+ class OnnxEmbedder:
63
+ """A built-in model: an ONNX encoder and its tokenizer from a preset, mean pooling over the attention mask, L2-normalised."""
64
+
65
+ max_tokens = 512
66
+ # Padded tokens per inference call: 1024 keeps peak RSS near 310 MB for e5-small (8 × 512 reached 550 MB).
67
+ token_budget = 1024
68
+
69
+ def __init__(self, model_dir: Path, preset=None, device: str = "auto", threads: int = 4):
70
+ import onnxruntime as ort
71
+
72
+ from passageway.models import PRESETS, choose_device
73
+
74
+ self.preset = preset or PRESETS["multilingual-e5-small"]
75
+ self.name, self.dim, self.min_similarity = self.preset.label, self.preset.dim, self.preset.min_similarity
76
+ self.max_tokens = self.preset.max_tokens
77
+ options = ort.SessionOptions()
78
+ options.intra_op_num_threads = threads
79
+ options.inter_op_num_threads = 1
80
+ # The arena would keep the largest batch's peak reserved for the life of the process.
81
+ options.enable_cpu_mem_arena = False
82
+ providers = choose_device(device, ort.get_available_providers())
83
+ self._session = ort.InferenceSession(str(Path(model_dir) / "model.onnx"), options, providers=providers)
84
+ self._inputs = {i.name for i in self._session.get_inputs()}
85
+ self._encode_batch = self._tokenizer(self.preset, Path(model_dir))
86
+
87
+ @staticmethod
88
+ def _tokenizer(preset, model_dir: Path):
89
+ """A function from texts to (ids, mask) arrays, padded to the longest text and cut at max_tokens."""
90
+ if preset.tokenizer == "sentencepiece":
91
+ import sentencepiece as spm
92
+
93
+ # The SentencePiece model gives the same ids as the HF tokenizer.json at ~60 MB instead of ~280 MB.
94
+ sp = spm.SentencePieceProcessor(model_file=str(model_dir / "sentencepiece.bpe.model"))
95
+
96
+ def encode(text: str) -> list[int]:
97
+ # XLM-R ids: fairseq offset of 1 over SentencePiece, its <unk> (0) is 3, wrapped in <s> (0) … </s> (2).
98
+ ids = [3 if i == 0 else i + 1 for i in sp.encode(text)][: preset.max_tokens - 2]
99
+ return [0, *ids, 2]
100
+
101
+ pad = 1
102
+ else:
103
+ from tokenizers import Tokenizer
104
+
105
+ tok = Tokenizer.from_file(str(model_dir / "tokenizer.json"))
106
+ tok.enable_truncation(preset.max_tokens)
107
+ tok.no_padding()
108
+ pad = next((i for i in (tok.token_to_id(t) for t in ("[PAD]", "<pad>")) if i is not None), 0)
109
+
110
+ def encode(text: str) -> list[int]:
111
+ return tok.encode(text).ids
112
+
113
+ def batch(texts: list[str]) -> tuple[np.ndarray, np.ndarray]:
114
+ encoded = [encode(t) for t in texts]
115
+ width = max(len(e) for e in encoded)
116
+ ids = np.full((len(encoded), width), pad, dtype=np.int64)
117
+ mask = np.zeros((len(encoded), width), dtype=np.int64)
118
+ for row, e in enumerate(encoded):
119
+ ids[row, : len(e)] = e
120
+ mask[row, : len(e)] = 1
121
+ return ids, mask
122
+
123
+ return batch
124
+
125
+ def _embed(self, texts: list[str]) -> np.ndarray:
126
+ if not texts:
127
+ return np.zeros((0, self.dim), dtype=np.float32)
128
+ lengths = [len(self._encode_batch([t])[0][0]) for t in texts]
129
+ out = np.zeros((len(texts), self.dim), dtype=np.float32)
130
+ for batch in token_batches(lengths, self.token_budget):
131
+ ids, mask = self._encode_batch([texts[i] for i in batch])
132
+ feed = {"input_ids": ids, "attention_mask": mask}
133
+ if "token_type_ids" in self._inputs:
134
+ feed["token_type_ids"] = np.zeros_like(ids)
135
+ hidden = self._session.run(None, feed)[0]
136
+ weights = mask[..., None].astype(np.float32)
137
+ out[batch] = (hidden * weights).sum(axis=1) / np.maximum(weights.sum(axis=1), 1e-9)
138
+ return _normalise(out)
139
+
140
+ def embed_queries(self, texts: list[str]) -> np.ndarray:
141
+ return self._embed([self.preset.query_prefix + t for t in texts])
142
+
143
+ def embed_passages(self, texts: list[str]) -> np.ndarray:
144
+ return self._embed([self.preset.passage_prefix + t for t in texts])
145
+
146
+
147
+ class HashEmbedder:
148
+ """Bag of lower-cased words hashed into buckets: overlap in words means similarity. Tests only."""
149
+
150
+ name = "hash-test"
151
+
152
+ def __init__(self, dim: int = 384):
153
+ self.dim = dim
154
+
155
+ def _embed(self, texts: list[str]) -> np.ndarray:
156
+ v = np.zeros((len(texts), self.dim), dtype=np.float32)
157
+ for row, text in enumerate(texts):
158
+ for word in re.findall(r"\w+", text.lower()):
159
+ bucket = int.from_bytes(hashlib.blake2b(word.encode(), digest_size=4).digest(), "big")
160
+ v[row, bucket % self.dim] += 1
161
+ return _normalise(v)
162
+
163
+ embed_queries = _embed
164
+ embed_passages = _embed
165
+
166
+
167
+ def model_folder(cfg: EmbedderConfig, config) -> Path:
168
+ """Where a preset's files live: one rule with `passageway model` (models.model_folder_for); config.model_dir is the
169
+ server's resolved PASSAGEWAY_MODEL_DIR for the default preset."""
170
+ if cfg.model_dir:
171
+ return cfg.model_dir
172
+ if cfg.model == DEFAULT_PRESET:
173
+ return config.model_dir
174
+ return model_folder_for(cfg.model, config.models_root or config.data_dir, None)
175
+
176
+
177
+ def build_embedder(cfg: EmbedderConfig, config, label: str = "default") -> Embedder:
178
+ """The embedder a configuration describes. A model or program that is not there is a ConfigError, so that a server
179
+ can start in its error mode and say what is wrong instead of crashing."""
180
+ if cfg.kind == "command":
181
+ from passageway.embed_command import CommandEmbedder
182
+
183
+ if cfg.min_similarity is None:
184
+ log.warning("embedder %r has no min_similarity: the relevance gate of the built-in model applies, which is "
185
+ "probably wrong for this one; find a value with `passageway eval --sweep`", label)
186
+ try:
187
+ return CommandEmbedder(label, cfg.command, env=cfg.env, secrets=cfg.secrets, batch_size=cfg.batch_size,
188
+ timeout=cfg.timeout)
189
+ except EmbedderError as e: # its output can hold hosts and paths: only the log gets it
190
+ raise EmbedderStartError(str(e), public=f"the embedder {label!r} could not start; see the server log") from None
191
+ preset = PRESETS[cfg.model]
192
+ folder = model_folder(cfg, config)
193
+ missing, wrong = files_missing(preset, folder), verify(preset, folder)
194
+ if missing or wrong:
195
+ if not cfg.download:
196
+ what = f"missing {', '.join(missing)}" if missing else f"a wrong checksum for {', '.join(wrong)}"
197
+ raise ConfigError(f"the model {cfg.model!r} in {folder} has {what}: run `passageway model download {cfg.model}` "
198
+ "(download = true in [embedder] does it at start)")
199
+ log.info("fetching the model %s (%d MB) from huggingface.co into %s, pinned and checksummed; the server answers "
200
+ "once it is there", cfg.model, preset.size_mb, folder)
201
+ try:
202
+ download(preset, folder, report=lambda m: log.info("%s", m))
203
+ except (ModelError, OSError) as e:
204
+ raise ConfigError(f"could not fetch the model {cfg.model!r}: {e}") from None
205
+ try:
206
+ embedder = OnnxEmbedder(folder, preset, cfg.device)
207
+ except ModelError as e:
208
+ raise ConfigError(str(e)) from None
209
+ if cfg.min_similarity is None and getattr(embedder, "min_similarity", None) is None:
210
+ log.warning("the model %r has no measured min_similarity yet, so the default gate of the built-in model (0.845) "
211
+ "applies, and on another similarity scale almost nothing will pass it: set min_similarity under "
212
+ "[embedder] from a run of `passageway eval --sweep`", cfg.model)
213
+ return embedder