61 lines
2.3 KiB
Python
61 lines
2.3 KiB
Python
"""Reusable vectors keyed by the exact input and model fingerprint."""
|
|
|
|
import hashlib
|
|
import json
|
|
import sqlite3
|
|
from contextlib import closing
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
|
|
|
|
def cache_key(profile: dict, text: str) -> str:
|
|
return hashlib.sha256(
|
|
json.dumps(
|
|
{"profile": profile["fingerprint"], "text": text}, ensure_ascii=False, sort_keys=True
|
|
).encode()
|
|
).hexdigest()
|
|
|
|
|
|
class EmbeddingCache:
|
|
def __init__(self, path: Path, *, readonly: bool = False):
|
|
self.path = Path(path)
|
|
self.readonly = readonly
|
|
if readonly:
|
|
return
|
|
self.path.parent.mkdir(parents=True, exist_ok=True)
|
|
with closing(sqlite3.connect(self.path)) as conn, conn:
|
|
conn.execute("""CREATE TABLE IF NOT EXISTS embeddings (
|
|
key TEXT PRIMARY KEY, dimensions INTEGER NOT NULL, vector BLOB NOT NULL,
|
|
vector_sha256 TEXT NOT NULL)""")
|
|
|
|
def load(self, keys: list[str]) -> dict[str, np.ndarray]:
|
|
found = {}
|
|
target = self.path.resolve().as_uri() + "?mode=ro" if self.readonly else str(self.path)
|
|
with closing(sqlite3.connect(target, uri=self.readonly)) as conn:
|
|
for key in set(keys):
|
|
row = conn.execute(
|
|
"SELECT dimensions,vector,vector_sha256 FROM embeddings WHERE key=?", (key,)
|
|
).fetchone()
|
|
if not row or hashlib.sha256(row[1]).hexdigest() != row[2]:
|
|
continue
|
|
vector = np.frombuffer(row[1], dtype=np.float32)
|
|
if (
|
|
len(vector) == row[0]
|
|
and np.isfinite(vector).all()
|
|
and np.linalg.norm(vector) > 0
|
|
):
|
|
found[key] = vector.copy()
|
|
return found
|
|
|
|
def save(self, vectors: dict[str, np.ndarray]):
|
|
if self.readonly:
|
|
raise RuntimeError("Cannot write to a read-only embedding cache")
|
|
with closing(sqlite3.connect(self.path)) as conn, conn:
|
|
for key, vector in vectors.items():
|
|
blob = np.asarray(vector, dtype=np.float32).tobytes()
|
|
conn.execute(
|
|
"INSERT OR REPLACE INTO embeddings VALUES(?,?,?,?)",
|
|
(key, len(vector), blob, hashlib.sha256(blob).hexdigest()),
|
|
)
|