File size: 1,494 Bytes
892fa81 | 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 | """Embedding cache — load-once, reuse-many for embedding models.
When a provider needs to generate an embedding, it asks this cache for
the model. The first call loads the model; subsequent calls return the
cached instance. This ensures we never load the same model twice.
"""
from __future__ import annotations
import threading
from typing import Any, Callable, Dict
class EmbeddingCache:
"""Thread-safe cache for embedding models.
Usage:
cache = EmbeddingCache()
model = cache.get_or_load("clip-vit-base", lambda: load_clip_model())
embedding = model.encode(img)
"""
def __init__(self) -> None:
self._cache: Dict[str, Any] = {}
self._lock = threading.RLock()
def get_or_load(self, key: str, loader: Callable[[], Any]) -> Any:
"""Return the cached model, or load it via `loader` and cache it."""
with self._lock:
if key not in self._cache:
self._cache[key] = loader()
return self._cache[key]
def is_loaded(self, key: str) -> bool:
with self._lock:
return key in self._cache
def evict(self, key: str) -> bool:
with self._lock:
return self._cache.pop(key, None) is not None
def clear(self) -> int:
with self._lock:
n = len(self._cache)
self._cache.clear()
return n
def keys(self) -> list[str]:
with self._lock:
return list(self._cache.keys())
|