Spaces:
Sleeping
Sleeping
File size: 2,523 Bytes
526950e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 | """Embedding cache: (model, text_sha256) -> vector.
Append-only JSONL on disk; in-memory dict for fast lookup. Same input always
returns the same vector regardless of the doc that produced it, so the cache
is reusable across builds, branches, and tests.
Atomic flush via temp file + os.replace.
"""
from __future__ import annotations
import hashlib
import json
import os
from datetime import datetime, timezone
from pathlib import Path
def _hash(text: str) -> str:
return hashlib.sha256(text.encode("utf-8")).hexdigest()
class EmbeddingCache:
def __init__(self, path: Path) -> None:
self.path = path
self._mem: dict[tuple[str, str], list[float]] = {}
self._dirty: list[dict] = []
def load(self) -> None:
"""Read existing cache file (if present) into memory; last-wins on duplicates."""
if not self.path.is_file():
return
with open(self.path, encoding="utf-8") as f:
for line in f:
line = line.strip()
if not line:
continue
try:
rec = json.loads(line)
except json.JSONDecodeError:
continue
key = (rec.get("model", ""), rec.get("text_sha256", ""))
self._mem[key] = rec.get("vector", [])
def get(self, *, model: str, text: str) -> list[float] | None:
return self._mem.get((model, _hash(text)))
def put(self, *, model: str, text: str, vector: list[float]) -> None:
h = _hash(text)
self._mem[(model, h)] = vector
self._dirty.append({
"model": model,
"text_sha256": h,
"vector": vector,
"cached_utc": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
})
def flush(self) -> None:
"""Append dirty records to disk atomically.
Loads existing on-disk content, appends new, writes to .tmp, renames.
"""
if not self._dirty:
return
self.path.parent.mkdir(parents=True, exist_ok=True)
existing_lines: list[str] = []
if self.path.is_file():
existing_lines = self.path.read_text(encoding="utf-8").splitlines()
new_lines = existing_lines + [json.dumps(r, ensure_ascii=False) for r in self._dirty]
tmp = self.path.with_suffix(self.path.suffix + ".tmp")
tmp.write_text("\n".join(new_lines) + "\n", encoding="utf-8")
os.replace(tmp, self.path)
self._dirty.clear()
|