rpi: tg — кэш векторов (npz) и служба tg-embed: поиск ~6 с вместо ~40
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
gilfanovr
co-authored by
Claude Opus 5.5
Claude Opus 5.5
parent
48c129615f
commit
3bdb75d465
@@ -8,6 +8,7 @@
|
||||
tg report [--hours 24] [--dry] сводка по темам за период → владельцу в Telegram (--dry — только напечатать)
|
||||
tg stats сколько чатов/сообщений/кусков, когда обновлялось
|
||||
tg index (сервис tg-export, юзер tgexport) — дописать в индекс новые сообщения выгрузки
|
||||
tg serve (служба tg-embed, юзер astra) — держать модель в памяти, векторы запросов по unix-сокету
|
||||
|
||||
Индекс — /mnt/backup/tg/index.db (SQLite: FTS5 + эмбеддинги potion-multilingual-128M), пишет tgexport,
|
||||
читают все из группы astra. Кусок = подряд идущие сообщения одного чата с паузами < 20 мин (до ~800 знаков).
|
||||
@@ -38,13 +39,44 @@ TZ = ZoneInfo("Asia/Nicosia")
|
||||
GAP = 20 * 60 # пауза, после которой начинается новый кусок
|
||||
CHUNK_CHARS = int(os.environ.get("TG_CHUNK_CHARS", "800"))
|
||||
TG_MAX = 3800 # одно сообщение Telegram (лимит 4096) — сводку режем на части
|
||||
# векторы всех кусков одним файлом рядом с индексом: чтение 85 тыс. BLOB из SQLite на HDD — ~40 с на запрос,
|
||||
# а np.load(mmap) — доли секунды; пересобирается после каждой индексации
|
||||
CACHE = DB.with_suffix(".emb.npz")
|
||||
|
||||
_emb = None
|
||||
SOCK = os.environ.get("TG_EMBED_SOCK", "/run/tg-embed/embed.sock")
|
||||
|
||||
|
||||
def embed(texts):
|
||||
def _embed_remote(texts):
|
||||
"""Векторы от службы tg-embed (модель уже в памяти) — иначе None. Загрузка модели в процесс — ~11 с на Pi 4."""
|
||||
import socket
|
||||
import numpy as np
|
||||
if not os.path.exists(SOCK):
|
||||
return None
|
||||
try:
|
||||
with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as s:
|
||||
s.settimeout(30)
|
||||
s.connect(SOCK)
|
||||
s.sendall(json.dumps({"texts": texts}).encode() + b"\n")
|
||||
buf = b""
|
||||
while not buf.endswith(b"\n"):
|
||||
part = s.recv(1 << 20)
|
||||
if not part:
|
||||
break
|
||||
buf += part
|
||||
r = json.loads(buf)
|
||||
return np.frombuffer(bytes.fromhex(r["vecs"]), dtype=np.float32).reshape(r["n"], r["dim"])
|
||||
except (OSError, ValueError, KeyError):
|
||||
return None
|
||||
|
||||
|
||||
def embed(texts, remote=True):
|
||||
global _emb
|
||||
import numpy as np
|
||||
if remote and len(texts) <= 64:
|
||||
v = _embed_remote(list(texts))
|
||||
if v is not None:
|
||||
return v
|
||||
if _emb is None:
|
||||
from fastembed import TextEmbedding
|
||||
_emb = TextEmbedding(EMB_MODEL, cache_dir=MODELS, threads=4)
|
||||
@@ -162,9 +194,24 @@ def cmd_index(a):
|
||||
total += len(new)
|
||||
c.execute("INSERT OR REPLACE INTO meta VALUES ('indexed', ?)", (dt.datetime.now(TZ).isoformat(timespec="seconds"),))
|
||||
c.commit()
|
||||
if total or not CACHE.exists():
|
||||
build_cache(c)
|
||||
print(f"проиндексировано новых сообщений: {total}")
|
||||
|
||||
|
||||
def build_cache(c):
|
||||
import numpy as np
|
||||
rows = c.execute("SELECT id, chat_id, ts1, emb FROM chunk ORDER BY id").fetchall()
|
||||
if not rows:
|
||||
return
|
||||
tmp = CACHE.with_name(CACHE.name + ".tmp.npz")
|
||||
np.savez(tmp, ids=np.array([r["id"] for r in rows], dtype=np.int64),
|
||||
chat=np.array([r["chat_id"] for r in rows], dtype=np.int64),
|
||||
ts=np.array([r["ts1"] for r in rows], dtype=np.int64),
|
||||
emb=np.frombuffer(b"".join(r["emb"] for r in rows), dtype=np.float32).reshape(len(rows), -1))
|
||||
os.replace(tmp, CACHE)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------- поиск
|
||||
|
||||
def link(chat, msg_id):
|
||||
@@ -216,16 +263,25 @@ def search(q, chat=None, since=None, n=12, min_sim=None):
|
||||
ranks[row["id"]] = ranks.get(row["id"], 0) + 1 / (60 + r)
|
||||
by_words.add(row["id"])
|
||||
qv = embed([q])[0]
|
||||
ids, mats = [], []
|
||||
for row in c.execute(f"SELECT c.id, c.emb FROM chunk c WHERE {w}", args): # 384 float32 на кусок — в памяти ок
|
||||
ids.append(row["id"])
|
||||
mats.append(row["emb"])
|
||||
if ids:
|
||||
m = np.frombuffer(b"".join(mats), dtype=np.float32).reshape(len(ids), -1)
|
||||
s = m @ qv
|
||||
for r, i in enumerate(np.argsort(-s)[:60]):
|
||||
sims[ids[i]] = float(s[i])
|
||||
ranks[ids[i]] = ranks.get(ids[i], 0) + 1 / (60 + r)
|
||||
if CACHE.exists():
|
||||
z = np.load(CACHE)
|
||||
ids, emb, mask = z["ids"], z["emb"], np.ones(len(z["ids"]), dtype=bool)
|
||||
if since:
|
||||
mask &= z["ts"] >= since
|
||||
if chat:
|
||||
cids = [r[0] for r in c.execute("SELECT id FROM chat WHERE title LIKE ? OR username LIKE ?", (f"%{chat}%", f"%{chat}%"))]
|
||||
mask &= np.isin(z["chat"], cids)
|
||||
ids, emb = ids[mask], emb[mask]
|
||||
else: # кэша нет (старый индекс) — медленно, из базы
|
||||
rows = c.execute(f"SELECT c.id, c.emb FROM chunk c WHERE {w}", args).fetchall()
|
||||
ids = np.array([r["id"] for r in rows], dtype=np.int64)
|
||||
emb = np.frombuffer(b"".join(r["emb"] for r in rows), dtype=np.float32).reshape(len(rows), -1) if rows else None
|
||||
if len(ids):
|
||||
s = emb @ qv
|
||||
top_i = np.argpartition(-s, min(60, len(s) - 1))[:60]
|
||||
for r, i in enumerate(top_i[np.argsort(-s[top_i])]):
|
||||
sims[int(ids[i])] = float(s[i])
|
||||
ranks[int(ids[i])] = ranks.get(int(ids[i]), 0) + 1 / (60 + r)
|
||||
top = sorted(ranks, key=ranks.get, reverse=True)
|
||||
if min_sim is not None: # для сводки: найденное словами — всегда, по смыслу — только достаточно близкое
|
||||
top = [i for i in top if i in by_words or sims.get(i, 0) >= min_sim]
|
||||
@@ -361,6 +417,30 @@ def cmd_report(a):
|
||||
print(f"сводка поставлена в очередь: {len(parts)} сообщ.")
|
||||
|
||||
|
||||
def cmd_serve(a):
|
||||
"""Служба tg-embed: модель в памяти, векторы запросов по unix-сокету (одна строка JSON → одна строка JSON)."""
|
||||
import socketserver
|
||||
embed(["прогрев"], remote=False)
|
||||
|
||||
class H(socketserver.StreamRequestHandler):
|
||||
def handle(self):
|
||||
try:
|
||||
req = json.loads(self.rfile.readline())
|
||||
v = embed(req["texts"][:64], remote=False)
|
||||
out = {"n": v.shape[0], "dim": v.shape[1], "vecs": v.astype("float32").tobytes().hex()}
|
||||
except Exception as e: # плохой запрос не роняет службу
|
||||
out = {"error": str(e)}
|
||||
self.wfile.write(json.dumps(out).encode() + b"\n")
|
||||
|
||||
if os.path.exists(SOCK):
|
||||
os.unlink(SOCK)
|
||||
old = os.umask(0o077) # сокет — только своему пользователю (astra)
|
||||
srv = socketserver.UnixStreamServer(SOCK, H)
|
||||
os.umask(old)
|
||||
print(f"tg-embed: {EMB_MODEL} в памяти, сокет {SOCK}", flush=True)
|
||||
srv.serve_forever()
|
||||
|
||||
|
||||
def main():
|
||||
p = argparse.ArgumentParser(prog="tg", description=__doc__.split("\n")[0])
|
||||
sub = p.add_subparsers(dest="cmd", required=True)
|
||||
@@ -373,6 +453,7 @@ def main():
|
||||
s = sub.add_parser("report"); s.add_argument("--hours", type=int, default=24); s.add_argument("--dry", action="store_true")
|
||||
sub.add_parser("stats")
|
||||
s = sub.add_parser("index"); s.add_argument("--quiet", action="store_true")
|
||||
sub.add_parser("serve")
|
||||
a = p.parse_args()
|
||||
globals()["cmd_" + a.cmd](a)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user