Spaces:
Paused
Paused
| # -*- coding: utf-8 -*- | |
| """A file embedding cache implementation for storing and retrieving | |
| embeddings in binary files.""" | |
| import hashlib | |
| import json | |
| import os | |
| from typing import Any, List | |
| import numpy as np | |
| from ._cache_base import EmbeddingCacheBase | |
| from .._logging import logger | |
| from ..types import ( | |
| Embedding, | |
| JSONSerializableObject, | |
| ) | |
| class FileEmbeddingCache(EmbeddingCacheBase): | |
| """The embedding cache class that stores each embeddings vector in | |
| binary files.""" | |
| def __init__( | |
| self, | |
| cache_dir: str = "./.cache/embeddings", | |
| max_file_number: int | None = None, | |
| max_cache_size: int | None = None, | |
| ) -> None: | |
| """Initialize the file embedding cache class. | |
| Args: | |
| cache_dir (`str`, defaults to `"./.cache/embeddings"`): | |
| The directory to store the embedding files. | |
| max_file_number (`int | None`, defaults to `None`): | |
| The maximum number of files to keep in the cache directory. If | |
| exceeded, the oldest files will be removed. | |
| max_cache_size (`int | None`, defaults to `None`): | |
| The maximum size of the cache directory in MB. If exceeded, | |
| the oldest files will be removed until the size is within the | |
| limit. | |
| """ | |
| self._cache_dir = os.path.abspath(cache_dir) | |
| self.max_file_number = max_file_number | |
| self.max_cache_size = max_cache_size | |
| def cache_dir(self) -> str: | |
| """The cache directory where the embedding files are stored.""" | |
| if not os.path.exists(self._cache_dir): | |
| os.makedirs(self._cache_dir, exist_ok=True) | |
| return self._cache_dir | |
| async def store( | |
| self, | |
| embeddings: List[Embedding], | |
| identifier: JSONSerializableObject, | |
| overwrite: bool = False, | |
| **kwargs: Any, | |
| ) -> None: | |
| """Store the embeddings with the given identifier. | |
| Args: | |
| embeddings (`List[Embedding]`): | |
| The embeddings to store. | |
| identifier (`JSONSerializableObject`): | |
| The identifier to distinguish the embeddings, which will be | |
| used to generate a hashable filename, so it should be | |
| JSON serializable (e.g. a string, number, list, dict). | |
| overwrite (`bool`, defaults to `False`): | |
| Whether to overwrite existing embeddings with the same | |
| identifier. If `True`, existing embeddings will be replaced. | |
| """ | |
| filename = self._get_filename(identifier) | |
| path_file = os.path.join(self.cache_dir, filename) | |
| if os.path.exists(path_file): | |
| if not os.path.isfile(path_file): | |
| raise RuntimeError( | |
| f"Path {path_file} exists but is not a file.", | |
| ) | |
| if overwrite: | |
| np.save(path_file, embeddings) | |
| await self._maintain_cache_dir() | |
| else: | |
| np.save(path_file, embeddings) | |
| await self._maintain_cache_dir() | |
| async def retrieve( | |
| self, | |
| identifier: JSONSerializableObject, | |
| ) -> List[Embedding] | None: | |
| """Retrieve the embeddings with the given identifier. If not found, | |
| return `None`. | |
| Args: | |
| identifier (`JSONSerializableObject`): | |
| The identifier to retrieve the embeddings, which will be | |
| used to generate a hashable filename, so it should be | |
| JSON serializable (e.g. a string, number, list, dict). | |
| """ | |
| filename = self._get_filename(identifier) | |
| path_file = os.path.join(self.cache_dir, filename) | |
| if os.path.exists(path_file): | |
| return np.load(os.path.join(self.cache_dir, filename)).tolist() | |
| return None | |
| async def remove(self, identifier: JSONSerializableObject) -> None: | |
| """Remove the embeddings with the given identifier. | |
| Args: | |
| identifier (`JSONSerializableObject`): | |
| The identifiers to remove the embeddings, which will be | |
| used to generate a hashable filename, so it should be | |
| JSON serializable (e.g. a string, number, list, dict). | |
| """ | |
| filename = self._get_filename(identifier) | |
| path_file = os.path.join(self.cache_dir, filename) | |
| if os.path.exists(path_file): | |
| os.remove(path_file) | |
| else: | |
| raise FileNotFoundError(f"File {path_file} does not exist.") | |
| async def clear(self) -> None: | |
| """Clear the cache directory by removing all files.""" | |
| for filename in os.listdir(self.cache_dir): | |
| if filename.endswith(".npy"): | |
| os.remove(os.path.join(self.cache_dir, filename)) | |
| def _get_cache_size(self) -> float: | |
| """Get the current size of the cache directory in MB.""" | |
| total_size = 0 | |
| for filename in os.listdir(self.cache_dir): | |
| if filename.endswith(".npy"): | |
| path_file = os.path.join(self.cache_dir, filename) | |
| if os.path.isfile(path_file): | |
| total_size += os.path.getsize(path_file) | |
| return total_size / (1024.0 * 1024.0) | |
| def _get_filename(identifier: JSONSerializableObject) -> str: | |
| """Generate a filename based on the identifier.""" | |
| json_str = json.dumps(identifier, ensure_ascii=False) | |
| return hashlib.sha256(json_str.encode("utf-8")).hexdigest() + ".npy" | |
| async def _maintain_cache_dir(self) -> None: | |
| """Maintain the cache directory by removing old files if the number of | |
| files exceeds the maximum limit or if the cache size exceeds the | |
| maximum size.""" | |
| files = [ | |
| (_.name, _.stat().st_mtime) | |
| for _ in os.scandir(self.cache_dir) | |
| if _.is_file() and _.name.endswith(".npy") | |
| ] | |
| files.sort(key=lambda x: x[1]) | |
| if self.max_file_number and len(files) > self.max_file_number: | |
| for file_name, _ in files[: 0 - self.max_file_number]: | |
| os.remove(os.path.join(self.cache_dir, file_name)) | |
| logger.info( | |
| "Remove cached embedding file %s for limited number " | |
| "of files (%d).", | |
| file_name, | |
| self.max_file_number, | |
| ) | |
| files = files[0 - self.max_file_number :] | |
| if ( | |
| self.max_cache_size is not None | |
| and self._get_cache_size() > self.max_cache_size | |
| ): | |
| removed_files = [] | |
| for filename, _ in files: | |
| os.remove(os.path.join(self.cache_dir, filename)) | |
| removed_files.append(filename) | |
| if self._get_cache_size() <= self.max_cache_size: | |
| break | |
| if removed_files: | |
| logger.info( | |
| "Remove %d cached embedding file(s) for limited " | |
| "cache size (%d MB).", | |
| len(removed_files), | |
| self.max_cache_size, | |
| ) | |