43 lines
1.8 KiB
Python
43 lines
1.8 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):
|
|
self.path = Path(path)
|
|
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 = {}
|
|
with closing(sqlite3.connect(self.path)) 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]):
|
|
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()))
|