Spaces:
Paused
Paused
File size: 7,126 Bytes
9792ea7 | 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 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 | # -*- 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
@property
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)
@staticmethod
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,
)
|