| from __future__ import annotations
|
|
|
| from collections import OrderedDict
|
| from dataclasses import dataclass
|
| from typing import Any, Callable, Iterable, Hashable
|
|
|
| import torch
|
|
|
|
|
| @dataclass
|
| class _CacheEntry:
|
| value: Any
|
| size_bytes: int
|
|
|
|
|
| class TextEncoderCache:
|
| def __init__(self, max_size_mb: float = 100) -> None:
|
| self.max_size_bytes = int(max_size_mb * 1024 * 1024)
|
| self._entries: "OrderedDict[Hashable, _CacheEntry]" = OrderedDict()
|
| self._size_bytes = 0
|
|
|
| def encode(
|
| self,
|
| encode_fn: Callable[[list[str]], list[Any]],
|
| prompts: Iterable[str] | str,
|
| device: torch.device | str | None = None,
|
| parallel: bool = False,
|
| cache_keys: Iterable[Hashable] | Hashable | None = None,
|
| ) -> list[Any]:
|
| if isinstance(prompts, str):
|
| prompts_list = [prompts]
|
| else:
|
| prompts_list = list(prompts)
|
| if not prompts_list:
|
| return []
|
| if cache_keys is None:
|
| keys_list = prompts_list
|
| else:
|
| if len(prompts_list) == 1 and not isinstance(cache_keys, list):
|
| keys_list = [cache_keys]
|
| else:
|
| keys_list = list(cache_keys)
|
| if len(keys_list) != len(prompts_list):
|
| raise ValueError("cache_keys must match the number of prompts.")
|
|
|
| if not parallel:
|
| results: list[Any] = []
|
| for prompt, cache_key in zip(prompts_list, keys_list):
|
| cached = self._entries.get(cache_key)
|
| if cached is not None:
|
| self._entries.move_to_end(cache_key)
|
| results.append(self._to_device(cached.value, device))
|
| continue
|
| encoded = encode_fn([prompt])
|
| if isinstance(encoded, (list, tuple)):
|
| if not encoded:
|
| raise ValueError("encode_fn returned empty embeddings.")
|
| encoded_item = encoded[0]
|
| else:
|
| encoded_item = encoded
|
| results.append(self._store(cache_key, encoded_item, device))
|
| return results
|
|
|
| results = [None] * len(prompts_list)
|
| missing_prompts: list[str] = []
|
| missing_indices: list[int] = []
|
| missing_keys: list[Hashable] = []
|
|
|
| for idx, (prompt, cache_key) in enumerate(zip(prompts_list, keys_list)):
|
| cached = self._entries.get(cache_key)
|
| if cached is None:
|
| missing_prompts.append(prompt)
|
| missing_indices.append(idx)
|
| missing_keys.append(cache_key)
|
| continue
|
| self._entries.move_to_end(cache_key)
|
| results[idx] = self._to_device(cached.value, device)
|
|
|
| if missing_prompts:
|
| encoded_batch = encode_fn(missing_prompts)
|
| if not isinstance(encoded_batch, list):
|
| encoded_batch = list(encoded_batch)
|
| if len(encoded_batch) != len(missing_prompts):
|
| raise ValueError("encode_fn returned unexpected number of embeddings.")
|
| for cache_key, idx, encoded in zip(missing_keys, missing_indices, encoded_batch):
|
| results[idx] = self._store(cache_key, encoded, device)
|
|
|
| return results
|
|
|
| def _store(self, cache_key: Hashable, encoded: Any, device: torch.device | str | None) -> Any:
|
| cached_value = self._detach_to_cpu(encoded)
|
| size_bytes = self._estimate_size_bytes(cached_value)
|
| if size_bytes <= self.max_size_bytes:
|
| existing = self._entries.pop(cache_key, None)
|
| if existing is not None:
|
| self._size_bytes -= existing.size_bytes
|
| self._entries[cache_key] = _CacheEntry(cached_value, size_bytes)
|
| self._size_bytes += size_bytes
|
| self._purge_if_needed()
|
| else:
|
| if cache_key in self._entries:
|
| self._entries.move_to_end(cache_key)
|
| return self._to_device(encoded, device)
|
|
|
| def _purge_if_needed(self) -> None:
|
| if self._size_bytes <= self.max_size_bytes:
|
| return
|
| while self._entries and self._size_bytes > self.max_size_bytes:
|
| _, entry = self._entries.popitem(last=False)
|
| self._size_bytes -= entry.size_bytes
|
|
|
| def _estimate_size_bytes(self, value: Any) -> int:
|
| if torch.is_tensor(value):
|
| return int(value.numel() * value.element_size())
|
| if isinstance(value, dict):
|
| return sum(self._estimate_size_bytes(v) for v in value.values())
|
| if isinstance(value, (list, tuple)):
|
| return sum(self._estimate_size_bytes(v) for v in value)
|
| return 0
|
|
|
| def _detach_to_cpu(self, value: Any) -> Any:
|
| if torch.is_tensor(value):
|
| if value.device.type == "cpu":
|
| return value.detach()
|
| return value.detach().to("cpu")
|
| if isinstance(value, dict):
|
| return {k: self._detach_to_cpu(v) for k, v in value.items()}
|
| if isinstance(value, tuple):
|
| items = [self._detach_to_cpu(v) for v in value]
|
| if hasattr(value, "_fields"):
|
| return value.__class__(*items)
|
| return tuple(items)
|
| if isinstance(value, list):
|
| return [self._detach_to_cpu(v) for v in value]
|
| return value
|
|
|
| def _to_device(self, value: Any, device: torch.device | str | None) -> Any:
|
| if device is None:
|
| return value
|
| if torch.is_tensor(value):
|
| return value.to(device)
|
| if isinstance(value, dict):
|
| return {k: self._to_device(v, device) for k, v in value.items()}
|
| if isinstance(value, tuple):
|
| items = [self._to_device(v, device) for v in value]
|
| if hasattr(value, "_fields"):
|
| return value.__class__(*items)
|
| return tuple(items)
|
| if isinstance(value, list):
|
| return [self._to_device(v, device) for v in value]
|
| return value
|
|
|
