dethrottled 0.1.1__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.
- dethrottled/__init__.py +2 -0
- dethrottled/__main__.py +4 -0
- dethrottled/_quiet.py +122 -0
- dethrottled/cache.py +123 -0
- dethrottled/corpus.py +436 -0
- dethrottled/documents.py +372 -0
- dethrottled/domains.py +275 -0
- dethrottled/extract.py +258 -0
- dethrottled/fetch.py +1042 -0
- dethrottled/health.py +250 -0
- dethrottled/media.py +142 -0
- dethrottled/ocr.py +250 -0
- dethrottled/paths.py +58 -0
- dethrottled/rank.py +334 -0
- dethrottled/search.py +614 -0
- dethrottled/server.py +447 -0
- dethrottled-0.1.1.dist-info/METADATA +267 -0
- dethrottled-0.1.1.dist-info/RECORD +22 -0
- dethrottled-0.1.1.dist-info/WHEEL +5 -0
- dethrottled-0.1.1.dist-info/entry_points.txt +3 -0
- dethrottled-0.1.1.dist-info/licenses/LICENSE +21 -0
- dethrottled-0.1.1.dist-info/top_level.txt +1 -0
dethrottled/__init__.py
ADDED
dethrottled/__main__.py
ADDED
dethrottled/_quiet.py
ADDED
|
@@ -0,0 +1,122 @@
|
|
|
1
|
+
"""Suppress third-party noise that we have checked and understood.
|
|
2
|
+
|
|
3
|
+
Two sources, one mechanism: both are written to file descriptor 2 by C or
|
|
4
|
+
Rust code that never passes through `sys.stderr`, so `contextlib.
|
|
5
|
+
redirect_stderr` cannot see them and no logging call can filter them.
|
|
6
|
+
|
|
7
|
+
## onnxruntime GPU discovery
|
|
8
|
+
|
|
9
|
+
onnxruntime 1.27 probes /sys/class/drm during MODULE IMPORT -- not session
|
|
10
|
+
creation -- and writes straight to file descriptor 2 when it finds no GPU
|
|
11
|
+
vendor file:
|
|
12
|
+
|
|
13
|
+
[W:onnxruntime:Default, device_discovery.cc:283 GetGpuDevices] Failed to
|
|
14
|
+
detect devices under "/sys/class/drm/card1"
|
|
15
|
+
|
|
16
|
+
On a CPU-only host the probe is correct, and falling back
|
|
17
|
+
to CPU is what we want. The message is noise printed on every invocation.
|
|
18
|
+
|
|
19
|
+
`SessionOptions.log_severity_level` cannot suppress it. That configures a
|
|
20
|
+
session, and this is emitted by a static logger built during import, before any
|
|
21
|
+
session exists -- which is why setting it in corpus.py did not work. There is no
|
|
22
|
+
environment variable for it in 1.27 either; the knob is not there.
|
|
23
|
+
|
|
24
|
+
So it is caught where it is written: at the file descriptor, around the first
|
|
25
|
+
import only. `contextlib.redirect_stderr` does not work here because the write
|
|
26
|
+
comes from C++ and never passes through `sys.stderr`.
|
|
27
|
+
|
|
28
|
+
This hides C-level stderr chatter during that one import. It does not hide
|
|
29
|
+
import failures -- those raise, and Python prints the traceback after the
|
|
30
|
+
descriptor is restored.
|
|
31
|
+
"""
|
|
32
|
+
import os
|
|
33
|
+
import sys
|
|
34
|
+
|
|
35
|
+
|
|
36
|
+
def load():
|
|
37
|
+
"""The onnxruntime module, imported quietly the first time."""
|
|
38
|
+
module = sys.modules.get("onnxruntime")
|
|
39
|
+
if module is not None:
|
|
40
|
+
return module # already imported; it has had its say
|
|
41
|
+
|
|
42
|
+
saved = None
|
|
43
|
+
try:
|
|
44
|
+
sys.stderr.flush()
|
|
45
|
+
saved = os.dup(2)
|
|
46
|
+
devnull = os.open(os.devnull, os.O_WRONLY)
|
|
47
|
+
try:
|
|
48
|
+
os.dup2(devnull, 2)
|
|
49
|
+
finally:
|
|
50
|
+
os.close(devnull)
|
|
51
|
+
except OSError:
|
|
52
|
+
# No usable stderr to swap out. Import anyway and accept the noise;
|
|
53
|
+
# a missing descriptor is not a reason to fail to load the runtime.
|
|
54
|
+
if saved is not None:
|
|
55
|
+
os.close(saved)
|
|
56
|
+
saved = None
|
|
57
|
+
|
|
58
|
+
try:
|
|
59
|
+
import onnxruntime
|
|
60
|
+
finally:
|
|
61
|
+
if saved is not None:
|
|
62
|
+
try:
|
|
63
|
+
os.dup2(saved, 2)
|
|
64
|
+
finally:
|
|
65
|
+
os.close(saved)
|
|
66
|
+
return onnxruntime
|
|
67
|
+
|
|
68
|
+
|
|
69
|
+
# ── curl_cffi transport chatter ──────────────────────────────────────────────
|
|
70
|
+
#
|
|
71
|
+
# curl_cffi's HTTP/2 layer prints this on hosts that close a connection without
|
|
72
|
+
# a TLS close_notify:
|
|
73
|
+
#
|
|
74
|
+
# h2 connection driver error: peer closed connection without sending TLS
|
|
75
|
+
# close_notify: https://docs.rs/rustls/...
|
|
76
|
+
#
|
|
77
|
+
# It is written from Rust straight to fd 2, and it is not an error we have any
|
|
78
|
+
# say in: an unclean shutdown is the SERVER's choice, the response has already
|
|
79
|
+
# been received in full by the time it appears, and the fetch succeeds. Left
|
|
80
|
+
# alone it prints on every affected fetch and buries anything worth reading.
|
|
81
|
+
#
|
|
82
|
+
# Suppressed at the descriptor, around the call only -- not process-wide, which
|
|
83
|
+
# would also hide the next real traceback.
|
|
84
|
+
|
|
85
|
+
|
|
86
|
+
class quiet_fd2:
|
|
87
|
+
"""Silence writes to file descriptor 2 for the duration of a block.
|
|
88
|
+
|
|
89
|
+
A context manager rather than a decorator because the noisy call is one
|
|
90
|
+
line inside a function that has other things to say.
|
|
91
|
+
|
|
92
|
+
Degrades to doing nothing if the descriptor cannot be swapped: no usable
|
|
93
|
+
stderr is not a reason to fail a fetch.
|
|
94
|
+
"""
|
|
95
|
+
|
|
96
|
+
def __enter__(self):
|
|
97
|
+
self._saved = None
|
|
98
|
+
try:
|
|
99
|
+
sys.stderr.flush()
|
|
100
|
+
except Exception:
|
|
101
|
+
pass
|
|
102
|
+
try:
|
|
103
|
+
self._saved = os.dup(2)
|
|
104
|
+
devnull = os.open(os.devnull, os.O_WRONLY)
|
|
105
|
+
try:
|
|
106
|
+
os.dup2(devnull, 2)
|
|
107
|
+
finally:
|
|
108
|
+
os.close(devnull)
|
|
109
|
+
except OSError:
|
|
110
|
+
if self._saved is not None:
|
|
111
|
+
os.close(self._saved)
|
|
112
|
+
self._saved = None
|
|
113
|
+
return self
|
|
114
|
+
|
|
115
|
+
def __exit__(self, *exc):
|
|
116
|
+
if self._saved is not None:
|
|
117
|
+
try:
|
|
118
|
+
os.dup2(self._saved, 2)
|
|
119
|
+
finally:
|
|
120
|
+
os.close(self._saved)
|
|
121
|
+
self._saved = None
|
|
122
|
+
return False
|
dethrottled/cache.py
ADDED
|
@@ -0,0 +1,123 @@
|
|
|
1
|
+
#!/usr/bin/env python3
|
|
2
|
+
"""SQLite cache for the free stack.
|
|
3
|
+
|
|
4
|
+
The free stack has no quota to protect, but it does have manners and a clock.
|
|
5
|
+
Every cached hit is a page we do not re-request from someone else's server and
|
|
6
|
+
a second we do not spend, so caching is both the polite thing and the fast
|
|
7
|
+
thing. Search results expire quickly; fetched article bodies barely change, so
|
|
8
|
+
they are kept far longer.
|
|
9
|
+
|
|
10
|
+
One file, no daemon, trivially deletable — the whole lab is meant to be
|
|
11
|
+
removable with `rm -rf`.
|
|
12
|
+
"""
|
|
13
|
+
|
|
14
|
+
from __future__ import annotations
|
|
15
|
+
|
|
16
|
+
import hashlib
|
|
17
|
+
import json
|
|
18
|
+
import sqlite3
|
|
19
|
+
import threading
|
|
20
|
+
import time
|
|
21
|
+
from pathlib import Path
|
|
22
|
+
|
|
23
|
+
DEFAULT_TTL = {
|
|
24
|
+
"search": 6 * 3600, # a search result set goes stale fast
|
|
25
|
+
# The fetcher stores bodies under "extract". This table said "fetch", so
|
|
26
|
+
# every body fell through to the unnamed hourly default and 742 of the 762
|
|
27
|
+
# in the cache were already dead -- meaning almost every page was refetched
|
|
28
|
+
# almost every time, at nearly seven seconds each. Both spellings are here
|
|
29
|
+
# now; the intent was always twenty-one days.
|
|
30
|
+
"extract": 21 * 24 * 3600,
|
|
31
|
+
"fetch": 21 * 24 * 3600, # an article body essentially never changes
|
|
32
|
+
"robots": 24 * 3600,
|
|
33
|
+
}
|
|
34
|
+
|
|
35
|
+
|
|
36
|
+
class Cache:
|
|
37
|
+
def __init__(self, path: Path):
|
|
38
|
+
self.path = Path(path)
|
|
39
|
+
self.path.parent.mkdir(parents=True, exist_ok=True)
|
|
40
|
+
# Fetches run four at a time now, and a sqlite connection refuses to
|
|
41
|
+
# be used from a thread other than the one that created it. WAL plus a
|
|
42
|
+
# single lock is ample here: writes are small and rare, and the
|
|
43
|
+
# alternative -- a connection per thread -- is four writers racing for
|
|
44
|
+
# one file.
|
|
45
|
+
self._db = sqlite3.connect(str(self.path), timeout=30,
|
|
46
|
+
check_same_thread=False)
|
|
47
|
+
self._lock = threading.Lock()
|
|
48
|
+
self._db.execute("PRAGMA journal_mode=WAL")
|
|
49
|
+
self._db.execute("""
|
|
50
|
+
CREATE TABLE IF NOT EXISTS entries (
|
|
51
|
+
key TEXT PRIMARY KEY,
|
|
52
|
+
kind TEXT NOT NULL,
|
|
53
|
+
value TEXT NOT NULL,
|
|
54
|
+
ts REAL NOT NULL
|
|
55
|
+
)
|
|
56
|
+
""")
|
|
57
|
+
self._db.execute("CREATE INDEX IF NOT EXISTS idx_kind_ts ON entries(kind, ts)")
|
|
58
|
+
self._db.commit()
|
|
59
|
+
self.hits = 0
|
|
60
|
+
self.misses = 0
|
|
61
|
+
|
|
62
|
+
@staticmethod
|
|
63
|
+
def key(kind: str, *parts) -> str:
|
|
64
|
+
raw = "|".join(str(p) for p in parts)
|
|
65
|
+
digest = hashlib.blake2b(raw.encode("utf-8"), digest_size=16)
|
|
66
|
+
return "%s:%s" % (kind, digest.hexdigest())
|
|
67
|
+
|
|
68
|
+
def get(self, kind: str, *parts, ttl: float | None = None):
|
|
69
|
+
ttl = DEFAULT_TTL.get(kind, 3600) if ttl is None else ttl
|
|
70
|
+
key = self.key(kind, *parts)
|
|
71
|
+
with self._lock:
|
|
72
|
+
row = self._db.execute(
|
|
73
|
+
"SELECT value, ts FROM entries WHERE key=?", (key,)).fetchone()
|
|
74
|
+
if not row:
|
|
75
|
+
self.misses += 1
|
|
76
|
+
return None
|
|
77
|
+
value, ts = row
|
|
78
|
+
if ttl >= 0 and (time.time() - ts) > ttl:
|
|
79
|
+
self.misses += 1
|
|
80
|
+
return None
|
|
81
|
+
self.hits += 1
|
|
82
|
+
try:
|
|
83
|
+
return json.loads(value)
|
|
84
|
+
except json.JSONDecodeError:
|
|
85
|
+
return None
|
|
86
|
+
|
|
87
|
+
def put(self, kind: str, *parts, value) -> None:
|
|
88
|
+
key = self.key(kind, *parts)
|
|
89
|
+
# The commit belongs inside the lock: it is part of the write, not a
|
|
90
|
+
# separate thing that happens afterwards.
|
|
91
|
+
with self._lock:
|
|
92
|
+
self._db.execute(
|
|
93
|
+
"INSERT OR REPLACE INTO entries (key, kind, value, ts) "
|
|
94
|
+
"VALUES (?,?,?,?)",
|
|
95
|
+
(key, kind, json.dumps(value, ensure_ascii=False), time.time()),
|
|
96
|
+
)
|
|
97
|
+
self._db.commit()
|
|
98
|
+
|
|
99
|
+
def stats(self) -> dict:
|
|
100
|
+
with self._lock:
|
|
101
|
+
rows = self._db.execute(
|
|
102
|
+
"SELECT kind, COUNT(*) FROM entries GROUP BY kind").fetchall()
|
|
103
|
+
total = self.hits + self.misses
|
|
104
|
+
return {
|
|
105
|
+
"hits": self.hits,
|
|
106
|
+
"misses": self.misses,
|
|
107
|
+
"hit_rate": round(self.hits / total, 3) if total else 0.0,
|
|
108
|
+
"stored": {kind: count for kind, count in rows},
|
|
109
|
+
"path": str(self.path),
|
|
110
|
+
}
|
|
111
|
+
|
|
112
|
+
def purge(self, older_than_days: float = 30) -> int:
|
|
113
|
+
cutoff = time.time() - older_than_days * 86400
|
|
114
|
+
with self._lock:
|
|
115
|
+
cursor = self._db.execute("DELETE FROM entries WHERE ts < ?", (cutoff,))
|
|
116
|
+
self._db.commit()
|
|
117
|
+
return cursor.rowcount
|
|
118
|
+
|
|
119
|
+
def close(self):
|
|
120
|
+
try:
|
|
121
|
+
self._db.close()
|
|
122
|
+
except Exception:
|
|
123
|
+
pass
|
dethrottled/corpus.py
ADDED
|
@@ -0,0 +1,436 @@
|
|
|
1
|
+
"""Semantic search over everything already fetched.
|
|
2
|
+
|
|
3
|
+
The cache became a corpus without anyone building one. Every page fetched is
|
|
4
|
+
kept for twenty-one days, so hundreds of real documents in several languages
|
|
5
|
+
are already on disk -- and the only way to find anything in them was to fetch
|
|
6
|
+
it again.
|
|
7
|
+
|
|
8
|
+
This indexes them. Pages are split into passages, each passage is embedded, and
|
|
9
|
+
a query is answered by cosine similarity over the lot.
|
|
10
|
+
|
|
11
|
+
No vector database. A few thousand passages at 384 dimensions is a matrix of a
|
|
12
|
+
few megabytes and a single numpy dot product; an index would be machinery with
|
|
13
|
+
nothing to do. Past roughly a hundred thousand passages that stops being true
|
|
14
|
+
and something like pgvector or FAISS earns its place -- but the interface below
|
|
15
|
+
need not change when it does, so that is a problem for whoever gets there.
|
|
16
|
+
|
|
17
|
+
One model: all-MiniLM-L6-v2, 87MB, ~0.3ms per passage, 384 dimensions.
|
|
18
|
+
|
|
19
|
+
multilingual-e5-small was carried alongside it and has been dropped. Measured
|
|
20
|
+
over ten documents on neighbouring subjects with ten questions whose right
|
|
21
|
+
answer was known, both scored 1.00 accuracy and 1.000 MRR -- but:
|
|
22
|
+
|
|
23
|
+
model acc@1 MRR score margin size query
|
|
24
|
+
minilm 1.00 1.000 0.242 87MB 1.4ms
|
|
25
|
+
e5 1.00 1.000 0.037 465MB 3.3ms
|
|
26
|
+
|
|
27
|
+
The margin is the gap between the right answer and the best wrong one, and it
|
|
28
|
+
is the number that decides whether a relevance floor can be set safely at all.
|
|
29
|
+
MiniLM separates by 0.242; e5 compresses everything upward and separates by
|
|
30
|
+
0.037. Equal ranking, six times the headroom, a fifth of the size.
|
|
31
|
+
|
|
32
|
+
What is given up is cross-language retrieval -- e5 could answer a French query
|
|
33
|
+
from English documents. That is a real capability and this is an English-first
|
|
34
|
+
tool, so it goes. The `model` column in the schema stays, so adding a second
|
|
35
|
+
index later needs no migration.
|
|
36
|
+
"""
|
|
37
|
+
from __future__ import annotations
|
|
38
|
+
|
|
39
|
+
import hashlib
|
|
40
|
+
import json
|
|
41
|
+
import os
|
|
42
|
+
import sqlite3
|
|
43
|
+
import threading
|
|
44
|
+
import time
|
|
45
|
+
from pathlib import Path
|
|
46
|
+
|
|
47
|
+
from . import paths as _paths
|
|
48
|
+
|
|
49
|
+
MODELS = _paths.model_dir()
|
|
50
|
+
|
|
51
|
+
ENGLISH_MODEL = str(MODELS / "emb-minilm")
|
|
52
|
+
|
|
53
|
+
|
|
54
|
+
# Below this a result is not an answer, it is the least-bad thing on file.
|
|
55
|
+
#
|
|
56
|
+
# Cosine similarity always returns a best match: asked something the corpus has
|
|
57
|
+
# nothing for, it will still rank something first. Measured on general content,
|
|
58
|
+
# real answers score 0.375 to 0.585 and questions with nothing to answer them
|
|
59
|
+
# score 0.030 to 0.032 -- an enormous empty band, so 0.22 sits comfortably in
|
|
60
|
+
# the middle of it.
|
|
61
|
+
#
|
|
62
|
+
# The inherited value was 0.40, tuned on one domain-specific corpus, and it was
|
|
63
|
+
# cutting straight through the real answers: it rejected half of them here.
|
|
64
|
+
FLOOR = float(os.environ.get("DETHROTTLED_CORPUS_FLOOR", "0.22"))
|
|
65
|
+
|
|
66
|
+
# Much longer than the fetch cache's 21 days, on purpose: outliving that cache
|
|
67
|
+
# is the point of the corpus, because a page worth retrieving next month is
|
|
68
|
+
# precisely the one the cache will have dropped. But "forever" is not a policy.
|
|
69
|
+
RETENTION_DAYS = int(os.environ.get("DETHROTTLED_CORPUS_RETENTION_DAYS", "180"))
|
|
70
|
+
|
|
71
|
+
# 50,000, not 200,000.
|
|
72
|
+
#
|
|
73
|
+
# The cap is really a memory budget wearing a different unit. Every passage is
|
|
74
|
+
# a 384-dimension float32 vector, and searching means holding all of them in
|
|
75
|
+
# one matrix:
|
|
76
|
+
#
|
|
77
|
+
# passages matrix RAM disk rebuild search
|
|
78
|
+
# 20,000 31 MB 42 MB 99ms 1.0ms
|
|
79
|
+
# 50,000 77 MB 105 MB 236ms 2.5ms
|
|
80
|
+
# 100,000 154 MB 210 MB 504ms 8.4ms
|
|
81
|
+
# 200,000 307 MB 419 MB 1009ms 13.8ms
|
|
82
|
+
#
|
|
83
|
+
# 200,000 asks for 307MB of resident memory before Python, the ONNX runtime and
|
|
84
|
+
# two embedding models have taken their share, which is most of a 2GB Pi. The
|
|
85
|
+
# search was never the problem -- it is 14ms even at the top of that table --
|
|
86
|
+
# the resident matrix is. 50,000 passages is about 12,000 pages, holds in
|
|
87
|
+
# 77MB, and rebuilds in a quarter of a second.
|
|
88
|
+
#
|
|
89
|
+
# Raise it if you have the memory; nothing here breaks at a larger number.
|
|
90
|
+
MAX_PASSAGES = int(os.environ.get("DETHROTTLED_CORPUS_MAX_PASSAGES", "50000"))
|
|
91
|
+
|
|
92
|
+
CHUNK_CHARS = int(os.environ.get("DETHROTTLED_CORPUS_CHUNK", "1200"))
|
|
93
|
+
CHUNK_OVERLAP = int(os.environ.get("DETHROTTLED_CORPUS_OVERLAP", "200"))
|
|
94
|
+
|
|
95
|
+
_LOADED = {}
|
|
96
|
+
_LOCK = threading.Lock()
|
|
97
|
+
|
|
98
|
+
|
|
99
|
+
def _model():
|
|
100
|
+
"""Load the model once per process. Returns tokenizer, session and inputs."""
|
|
101
|
+
with _LOCK:
|
|
102
|
+
if "en" not in _LOADED:
|
|
103
|
+
import numpy
|
|
104
|
+
from transformers import AutoTokenizer
|
|
105
|
+
|
|
106
|
+
from ._quiet import load as _load_onnxruntime
|
|
107
|
+
onnxruntime = _load_onnxruntime()
|
|
108
|
+
options = onnxruntime.SessionOptions()
|
|
109
|
+
options.intra_op_num_threads = int(
|
|
110
|
+
os.environ.get("DETHROTTLED_EMBED_THREADS", "4"))
|
|
111
|
+
# Quiets this session's own logging. The GPU-discovery warning is
|
|
112
|
+
# NOT this -- it fires during import, before any session exists,
|
|
113
|
+
# which is why setting it here never suppressed it. See _quiet.py.
|
|
114
|
+
options.log_severity_level = 3
|
|
115
|
+
session = onnxruntime.InferenceSession(
|
|
116
|
+
ENGLISH_MODEL + "/model.onnx", options,
|
|
117
|
+
providers=["CPUExecutionProvider"])
|
|
118
|
+
_LOADED["en"] = {
|
|
119
|
+
"np": numpy,
|
|
120
|
+
"tok": AutoTokenizer.from_pretrained(ENGLISH_MODEL,
|
|
121
|
+
trust_remote_code=False),
|
|
122
|
+
"sess": session,
|
|
123
|
+
"inputs": {i.name for i in session.get_inputs()},
|
|
124
|
+
}
|
|
125
|
+
return _LOADED["en"]
|
|
126
|
+
|
|
127
|
+
|
|
128
|
+
def embed(texts: list, *, batch: int = 32) -> list:
|
|
129
|
+
"""Unit-normalised vectors, mean-pooled over the tokens that are real.
|
|
130
|
+
|
|
131
|
+
Pooling ignores padding: averaging over pad tokens dilutes a short passage
|
|
132
|
+
towards whatever the padding embeds to, which is a slow way to make every
|
|
133
|
+
document look alike.
|
|
134
|
+
"""
|
|
135
|
+
if not texts:
|
|
136
|
+
return []
|
|
137
|
+
model = _model()
|
|
138
|
+
np = model["np"]
|
|
139
|
+
|
|
140
|
+
out = []
|
|
141
|
+
for start in range(0, len(texts), batch):
|
|
142
|
+
chunk = list(texts[start:start + batch])
|
|
143
|
+
encoded = model["tok"](chunk, padding=True, truncation=True,
|
|
144
|
+
max_length=256, return_tensors="np")
|
|
145
|
+
feed = {k: np.asarray(v).astype(np.int64)
|
|
146
|
+
for k, v in encoded.items() if k in model["inputs"]}
|
|
147
|
+
if "token_type_ids" in model["inputs"] and "token_type_ids" not in feed:
|
|
148
|
+
feed["token_type_ids"] = np.zeros_like(feed["input_ids"])
|
|
149
|
+
|
|
150
|
+
hidden = model["sess"].run(None, feed)[0]
|
|
151
|
+
mask = np.asarray(encoded["attention_mask"])[..., None].astype(np.float32)
|
|
152
|
+
pooled = (hidden * mask).sum(1) / np.maximum(mask.sum(1), 1e-9)
|
|
153
|
+
pooled /= np.maximum(np.linalg.norm(pooled, axis=1, keepdims=True), 1e-9)
|
|
154
|
+
out.extend(pooled.tolist())
|
|
155
|
+
return out
|
|
156
|
+
|
|
157
|
+
|
|
158
|
+
def passages(text: str, title: str = "") -> list:
|
|
159
|
+
"""Split a page into overlapping passages.
|
|
160
|
+
|
|
161
|
+
Overlapping because the sentence that answers a question has no obligation
|
|
162
|
+
to sit tidily inside one window, and a fact split across a boundary is a
|
|
163
|
+
fact neither passage can be retrieved for.
|
|
164
|
+
|
|
165
|
+
The title rides on every passage: a paragraph three screens into an article
|
|
166
|
+
rarely names its own subject.
|
|
167
|
+
"""
|
|
168
|
+
text = " ".join((text or "").split())
|
|
169
|
+
if not text:
|
|
170
|
+
return []
|
|
171
|
+
head = ("%s. " % title.strip()) if title.strip() else ""
|
|
172
|
+
step = max(CHUNK_CHARS - CHUNK_OVERLAP, 200)
|
|
173
|
+
return [head + text[i:i + CHUNK_CHARS]
|
|
174
|
+
for i in range(0, len(text), step)
|
|
175
|
+
if len(text[i:i + CHUNK_CHARS]) > 120]
|
|
176
|
+
|
|
177
|
+
|
|
178
|
+
def _pack(vector) -> bytes:
|
|
179
|
+
"""A vector as raw float32.
|
|
180
|
+
|
|
181
|
+
1,536 bytes against roughly 4,600 characters of JSON for the same 384
|
|
182
|
+
numbers -- and no parsing back into floats on the way out.
|
|
183
|
+
"""
|
|
184
|
+
import numpy as np
|
|
185
|
+
return np.asarray(vector, dtype="float32").tobytes()
|
|
186
|
+
|
|
187
|
+
|
|
188
|
+
def _unpack(blob):
|
|
189
|
+
"""Read either form. Rows written as JSON before this change still load."""
|
|
190
|
+
import numpy as np
|
|
191
|
+
if isinstance(blob, (bytes, bytearray, memoryview)):
|
|
192
|
+
return np.frombuffer(blob, dtype="float32")
|
|
193
|
+
return np.asarray(json.loads(blob), dtype="float32")
|
|
194
|
+
|
|
195
|
+
|
|
196
|
+
class Corpus:
|
|
197
|
+
"""Passages and their vectors, in one sqlite file beside the page cache."""
|
|
198
|
+
|
|
199
|
+
def __init__(self, path=None):
|
|
200
|
+
self.path = Path(path or (_paths.data_dir() / "corpus.sqlite"))
|
|
201
|
+
self.path.parent.mkdir(parents=True, exist_ok=True)
|
|
202
|
+
self._db = sqlite3.connect(str(self.path), timeout=30,
|
|
203
|
+
check_same_thread=False)
|
|
204
|
+
self._lock = threading.Lock()
|
|
205
|
+
self._cached = None # (matrix, meta), dropped when rows change
|
|
206
|
+
with self._lock:
|
|
207
|
+
self._db.execute("PRAGMA journal_mode=WAL")
|
|
208
|
+
self._db.execute("""
|
|
209
|
+
CREATE TABLE IF NOT EXISTS passages (
|
|
210
|
+
id TEXT PRIMARY KEY,
|
|
211
|
+
url TEXT NOT NULL,
|
|
212
|
+
title TEXT,
|
|
213
|
+
text TEXT NOT NULL,
|
|
214
|
+
model TEXT NOT NULL,
|
|
215
|
+
vec BLOB NOT NULL,
|
|
216
|
+
ts REAL NOT NULL
|
|
217
|
+
)""")
|
|
218
|
+
self._db.execute(
|
|
219
|
+
"CREATE INDEX IF NOT EXISTS idx_url ON passages(url)")
|
|
220
|
+
self._db.commit()
|
|
221
|
+
|
|
222
|
+
# Kept as a column and a constant rather than removed outright: the schema
|
|
223
|
+
# already separates indexes by model, so adding a second one back later
|
|
224
|
+
# needs no migration and no rewrite of what is already stored.
|
|
225
|
+
model_name = "minilm"
|
|
226
|
+
|
|
227
|
+
def add(self, url: str, title: str, text: str) -> int:
|
|
228
|
+
"""Index one page. Returns how many passages were new."""
|
|
229
|
+
chunks = passages(text, title)
|
|
230
|
+
if not chunks:
|
|
231
|
+
return 0
|
|
232
|
+
ids = [hashlib.blake2b(("%s|%s|%d" % (url, self.model_name, i))
|
|
233
|
+
.encode("utf-8"), digest_size=12).hexdigest()
|
|
234
|
+
for i in range(len(chunks))]
|
|
235
|
+
with self._lock:
|
|
236
|
+
known = {r[0] for r in self._db.execute(
|
|
237
|
+
"SELECT id FROM passages WHERE url=? AND model=?",
|
|
238
|
+
(url, self.model_name))}
|
|
239
|
+
fresh = [(i, c) for i, c in zip(ids, chunks) if i not in known]
|
|
240
|
+
if not fresh:
|
|
241
|
+
return 0
|
|
242
|
+
|
|
243
|
+
vectors = embed([c for _, c in fresh])
|
|
244
|
+
now = time.time()
|
|
245
|
+
with self._lock:
|
|
246
|
+
self._db.executemany(
|
|
247
|
+
"INSERT OR REPLACE INTO passages "
|
|
248
|
+
"(id, url, title, text, model, vec, ts) VALUES (?,?,?,?,?,?,?)",
|
|
249
|
+
[(i, url, title, c, self.model_name, _pack(v), now)
|
|
250
|
+
for (i, c), v in zip(fresh, vectors)])
|
|
251
|
+
self._db.commit()
|
|
252
|
+
# APPEND to the cached matrix rather than dropping it.
|
|
253
|
+
#
|
|
254
|
+
# Dropping it meant the next search re-read every row and rebuilt
|
|
255
|
+
# the whole array: a second of work at 200,000 passages, paid after
|
|
256
|
+
# every single write -- and the corpus is written on every fetch, so
|
|
257
|
+
# in practice almost every search paid it. The new rows are right
|
|
258
|
+
# here; adding them costs a copy proportional to what arrived, not
|
|
259
|
+
# to what was already stored.
|
|
260
|
+
if self._cached is not None:
|
|
261
|
+
import numpy as np
|
|
262
|
+
old_matrix, old_meta = self._cached
|
|
263
|
+
added = np.asarray(vectors, dtype="float32")
|
|
264
|
+
self._cached = (np.vstack([old_matrix, added]),
|
|
265
|
+
old_meta + [{"url": url, "title": title,
|
|
266
|
+
"text": c} for _, c in fresh])
|
|
267
|
+
return len(fresh)
|
|
268
|
+
|
|
269
|
+
def matrix(self):
|
|
270
|
+
"""Every vector for this model, as one array. Built once, then reused.
|
|
271
|
+
|
|
272
|
+
Cached because it was being rebuilt on every search: 0.30s at 2,790
|
|
273
|
+
passages and a projected 2.1s at twenty thousand, spent re-reading rows
|
|
274
|
+
that had not changed. Dropped whenever passages are added.
|
|
275
|
+
"""
|
|
276
|
+
if self._cached is not None:
|
|
277
|
+
return self._cached
|
|
278
|
+
|
|
279
|
+
import numpy as np
|
|
280
|
+
with self._lock:
|
|
281
|
+
rows = self._db.execute(
|
|
282
|
+
"SELECT url, title, text, vec FROM passages WHERE model=?",
|
|
283
|
+
(self.model_name,)).fetchall()
|
|
284
|
+
if not rows:
|
|
285
|
+
return None, []
|
|
286
|
+
meta = [{"url": r[0], "title": r[1], "text": r[2]} for r in rows]
|
|
287
|
+
self._cached = (np.asarray([_unpack(r[3]) for r in rows],
|
|
288
|
+
dtype="float32"), meta)
|
|
289
|
+
return self._cached
|
|
290
|
+
|
|
291
|
+
def prune(self) -> int:
|
|
292
|
+
"""Drop passages past the age limit, then past the cap. Oldest first.
|
|
293
|
+
|
|
294
|
+
Scoped to THIS model. It used to count and delete across the whole
|
|
295
|
+
table, which was right when one index existed and wrong the moment a
|
|
296
|
+
second appeared: the cap is sized in passages-per-index -- 200,000 is
|
|
297
|
+
about 20,000 pages -- and counting two indexes together silently halved
|
|
298
|
+
the history it was chosen to hold. Deleting oldest-first across both
|
|
299
|
+
would also have trimmed whichever index happened to be written first.
|
|
300
|
+
"""
|
|
301
|
+
cutoff = time.time() - RETENTION_DAYS * 24 * 3600
|
|
302
|
+
with self._lock:
|
|
303
|
+
removed = self._db.execute(
|
|
304
|
+
"DELETE FROM passages WHERE ts < ? AND model = ?",
|
|
305
|
+
(cutoff, self.model_name)).rowcount
|
|
306
|
+
total = self._db.execute(
|
|
307
|
+
"SELECT COUNT(*) FROM passages WHERE model = ?",
|
|
308
|
+
(self.model_name,)).fetchone()[0]
|
|
309
|
+
if total > MAX_PASSAGES:
|
|
310
|
+
removed += self._db.execute(
|
|
311
|
+
"DELETE FROM passages WHERE id IN ("
|
|
312
|
+
" SELECT id FROM passages WHERE model = ?"
|
|
313
|
+
" ORDER BY ts ASC LIMIT ?)",
|
|
314
|
+
(self.model_name, total - MAX_PASSAGES)).rowcount
|
|
315
|
+
self._db.commit()
|
|
316
|
+
if removed:
|
|
317
|
+
# A full drop, deliberately. Deletion takes rows out of the
|
|
318
|
+
# middle of the array and the bookkeeping to patch that is
|
|
319
|
+
# worse than the rebuild -- and pruning runs hourly at most,
|
|
320
|
+
# where writes run continuously.
|
|
321
|
+
self._cached = None
|
|
322
|
+
return removed
|
|
323
|
+
|
|
324
|
+
def search(self, query: str, *, limit: int = 10, per_url: int = 2,
|
|
325
|
+
floor: float | None = None, domains: list | None = None) -> list:
|
|
326
|
+
"""The passages closest to a query, at most `per_url` from any one page.
|
|
327
|
+
|
|
328
|
+
Capped per page because a long article about the right subject will
|
|
329
|
+
otherwise fill every slot with its own paragraphs, and a reader learns
|
|
330
|
+
nothing from the same source five times.
|
|
331
|
+
|
|
332
|
+
`floor` is what stops this answering a question it has nothing for.
|
|
333
|
+
Cosine always returns a best match, and asked in French about Moroccan
|
|
334
|
+
tenders -- of which this corpus holds none -- the best match was a
|
|
335
|
+
Google Drive login page at 0.35, against 0.65 for a real answer to a
|
|
336
|
+
real question. Returning nothing is the correct answer to "we have
|
|
337
|
+
nothing", and much better than returning the least-bad thing on file.
|
|
338
|
+
|
|
339
|
+
`domains` scopes the search to particular sites. The cache is shared by
|
|
340
|
+
every subject ever queried, so a corpus built from it is cross-subject
|
|
341
|
+
by construction: a question about solar capacity will otherwise happily
|
|
342
|
+
return pages about semiconductor foundries, which is useful about as
|
|
343
|
+
often as it is misleading.
|
|
344
|
+
"""
|
|
345
|
+
import numpy as np
|
|
346
|
+
if floor is None:
|
|
347
|
+
floor = FLOOR
|
|
348
|
+
vectors, meta = self.matrix()
|
|
349
|
+
if vectors is None:
|
|
350
|
+
return []
|
|
351
|
+
q = np.asarray(embed([query])[0], dtype="float32")
|
|
352
|
+
|
|
353
|
+
scores = vectors @ q
|
|
354
|
+
wanted = tuple(d.lower() for d in (domains or ()))
|
|
355
|
+
seen, out = {}, []
|
|
356
|
+
for i in np.argsort(-scores):
|
|
357
|
+
if scores[i] < floor:
|
|
358
|
+
break # sorted, so nothing below clears it
|
|
359
|
+
row = meta[i]
|
|
360
|
+
if wanted and not any(d in row["url"].lower() for d in wanted):
|
|
361
|
+
continue
|
|
362
|
+
if seen.get(row["url"], 0) >= per_url:
|
|
363
|
+
continue
|
|
364
|
+
seen[row["url"]] = seen.get(row["url"], 0) + 1
|
|
365
|
+
out.append(dict(row, score=float(scores[i])))
|
|
366
|
+
if len(out) >= limit:
|
|
367
|
+
break
|
|
368
|
+
return out
|
|
369
|
+
|
|
370
|
+
def stats(self) -> dict:
|
|
371
|
+
with self._lock:
|
|
372
|
+
rows = self._db.execute(
|
|
373
|
+
"SELECT model, COUNT(*), COUNT(DISTINCT url) FROM passages "
|
|
374
|
+
"GROUP BY model").fetchall()
|
|
375
|
+
return {r[0]: {"passages": r[1], "pages": r[2]} for r in rows}
|
|
376
|
+
|
|
377
|
+
|
|
378
|
+
# How much of a page the corpus takes, regardless of how much was fetched.
|
|
379
|
+
# These are two different jobs: a caller wants the whole article, the index
|
|
380
|
+
# wants the part of it that is about something. A page's tail is its footer,
|
|
381
|
+
# its comments and its related-articles, and indexing those makes them
|
|
382
|
+
# retrievable -- the corpus would start answering questions with boilerplate.
|
|
383
|
+
INDEX_CHARS = int(os.environ.get("DETHROTTLED_CORPUS_INDEX_CHARS", "4000"))
|
|
384
|
+
|
|
385
|
+
|
|
386
|
+
# Retention only means something if it runs. prune() had no callers at all --
|
|
387
|
+
# neither the 180-day limit nor the 200,000 cap was ever applied to anything --
|
|
388
|
+
# which mattered little while the corpus was frozen and matters now that it
|
|
389
|
+
# grows on every fetch. Once an hour is often enough for limits measured in
|
|
390
|
+
# months, and keeps a DELETE scan off the common path.
|
|
391
|
+
PRUNE_EVERY = int(os.environ.get("DETHROTTLED_CORPUS_PRUNE_SECONDS", "3600"))
|
|
392
|
+
_last_prune = 0.0
|
|
393
|
+
|
|
394
|
+
|
|
395
|
+
def index_fetched(fetched) -> int:
|
|
396
|
+
"""Add freshly fetched bodies to the corpus. Best effort, never fatal.
|
|
397
|
+
|
|
398
|
+
Lives here rather than in provider.py, where it used to, because provider
|
|
399
|
+
is only one of two callers and was the quieter one. The HTTP routes fetch
|
|
400
|
+
far more pages than any batch job does, and indexed none of them: the corpus
|
|
401
|
+
sat at 2790 passages across 617 urls while the extract cache grew past a
|
|
402
|
+
thousand entries. A retrieval index that the main interface does not feed
|
|
403
|
+
is a retrieval index that slowly stops being about anything.
|
|
404
|
+
|
|
405
|
+
Takes (url, result) pairs. `result` may be a fetcher result (ok/text) or an
|
|
406
|
+
API row (quality/content), because both callers have one of those and
|
|
407
|
+
neither should have to reshape it.
|
|
408
|
+
"""
|
|
409
|
+
global _last_prune
|
|
410
|
+
if os.environ.get("DETHROTTLED_CORPUS_AUTOINDEX", "1") != "1":
|
|
411
|
+
return 0
|
|
412
|
+
due = (time.time() - _last_prune) > PRUNE_EVERY
|
|
413
|
+
added = 0
|
|
414
|
+
try:
|
|
415
|
+
corpus = Corpus()
|
|
416
|
+
except Exception:
|
|
417
|
+
return 0
|
|
418
|
+
if due:
|
|
419
|
+
try:
|
|
420
|
+
corpus.prune()
|
|
421
|
+
except Exception:
|
|
422
|
+
pass
|
|
423
|
+
for url, result in fetched:
|
|
424
|
+
try:
|
|
425
|
+
ok = result.get("ok") or result.get("quality") == "ok"
|
|
426
|
+
if not ok:
|
|
427
|
+
continue
|
|
428
|
+
body = result.get("text") or result.get("content") or ""
|
|
429
|
+
if body:
|
|
430
|
+
corpus.add(url, result.get("title", ""), body[:INDEX_CHARS])
|
|
431
|
+
added += 1
|
|
432
|
+
except Exception:
|
|
433
|
+
continue
|
|
434
|
+
if due:
|
|
435
|
+
_last_prune = time.time()
|
|
436
|
+
return added
|