Spaces:
Sleeping
Sleeping
Download scripts/lib/embedding_cache.py from defenestrate/asa: direct link, hf CLI and curl.
- Browser
- Download file 2.52 kB
-
https://huggingface.co/spaces/defenestrate/asa/resolve/main/scripts/lib/embedding_cache.py
- Command line
-
hf download hf://spaces/defenestrate/asa/scripts/lib/embedding_cache.py
-
curl -L -o embedding_cache.py https://huggingface.co/spaces/defenestrate/asa/resolve/main/scripts/lib/embedding_cache.py
2.52 kB
| """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() | |