diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 50dd86e..0db2c61 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -12,4 +12,9 @@ repos: - id: ruff-format types_or: [ python, pyi ] files: .* + - repo: https://github.com/astral-sh/uv-pre-commit + # uv version. + rev: 0.6.10 + hooks: + - id: uv-lock files: '' \ No newline at end of file diff --git a/baguetter/indices/dense/faiss.py b/baguetter/indices/dense/faiss.py index 0913b39..5efac69 100644 --- a/baguetter/indices/dense/faiss.py +++ b/baguetter/indices/dense/faiss.py @@ -1,5 +1,6 @@ from __future__ import annotations +import os from dataclasses import asdict from typing import TYPE_CHECKING, Any @@ -15,7 +16,6 @@ from collections.abc import Callable from baguetter.types import DTypeLike, Key, ScoringMetric, TextOrVector - from baguetter.utils.file_repository import AbstractFileRepository def _support_nprobe(index: Any) -> bool: @@ -93,7 +93,6 @@ def from_config(cls, config: FaissDenseIndexConfig) -> FaissDenseIndex: def _save( self, path: str, - repository: AbstractFileRepository, ) -> str: """Save the index state and data.""" state = { @@ -106,10 +105,10 @@ def _save( index_file_path, ) = BaseDenseIndex.build_index_file_paths(path or self.name) - with repository.open(state_file_path, "wb") as file: + with open(state_file_path, "wb") as file: np.savez_compressed(file, state=state) - with repository.open(index_file_path, "wb") as file: + with open(index_file_path, "wb") as file: index = faiss.serialize_index(self.faiss_index) np.savez_compressed(file, index=index) return state_file_path @@ -119,7 +118,6 @@ def _load( cls, path: str, *, - repository: AbstractFileRepository, mmap: bool = False, ) -> FaissDenseIndex: """Load the index from saved state. @@ -137,22 +135,22 @@ def _load( """ state_file_path, index_file_path = BaseDenseIndex.build_index_file_paths(path) - if not repository.exists(state_file_path): + if not os.path.exists(state_file_path): msg = f"Index.state {state_file_path} not found in repository." raise FileNotFoundError(msg) - if not repository.exists(index_file_path): + if not os.path.exists(index_file_path): msg = f"Index.index {index_file_path} not found in repository." raise FileNotFoundError(msg) - with repository.open(state_file_path, "rb") as file: + with open(state_file_path, "rb") as file: state = np.load(file, allow_pickle=True)["state"][()] config = state["config"] index = cls.from_config(FaissDenseIndexConfig(**config)) index.key_mapping = state["key_mapping"] mmap_mode = "r" if mmap else None - with repository.open(index_file_path, "rb") as file: + with open(index_file_path, "rb") as file: stored = np.load(file, allow_pickle=True, mmap_mode=mmap_mode) index.faiss_index = faiss.deserialize_index(stored["index"]) diff --git a/baguetter/indices/dense/usearch.py b/baguetter/indices/dense/usearch.py index 9ea6370..5d199b0 100644 --- a/baguetter/indices/dense/usearch.py +++ b/baguetter/indices/dense/usearch.py @@ -1,5 +1,6 @@ from __future__ import annotations +import os import math import tempfile from dataclasses import asdict @@ -17,7 +18,6 @@ from collections.abc import Callable from baguetter.types import DTypeLike, Key, ScoringMetric, TextOrVector - from baguetter.utils.file_repository import AbstractFileRepository def get_normalization_fn(metric: MetricKind, embedding_dim: int) -> Callable[[np.ndarray], np.ndarray]: @@ -141,7 +141,6 @@ def from_config(cls, config: UsearchDenseIndexConfig) -> USearchDenseIndex: def _save( self, path: str, - repository: AbstractFileRepository, ) -> str: """Save the index state and data. @@ -161,12 +160,12 @@ def _save( index_file_path, ) = BaseDenseIndex.build_index_file_paths(path or self.name) - with repository.open(state_file_path, "wb") as file: + with open(state_file_path, "wb") as file: np.savez_compressed(file, state=state) with ( - repository.open(index_file_path, "wb") as file, - tempfile.NamedTemporaryFile() as temp_file, + open(index_file_path, "wb") as file, + tempfile.NamedTemporaryFile(delete=False) as temp_file, ): Index.save(self, temp_file.name) file.write(temp_file.read()) @@ -178,7 +177,6 @@ def _load( cls, path: str, *, - repository: AbstractFileRepository, mmap: bool = False, ) -> USearchDenseIndex: """Load the index from saved state. @@ -200,15 +198,15 @@ def _load( index_file_path, ) = BaseDenseIndex.build_index_file_paths(path) - if not repository.exists(state_file_path): + if not os.path.exists(state_file_path): msg = f"Index.state {state_file_path} not found in repository." raise FileNotFoundError(msg) - if not repository.exists(index_file_path): + if not os.path.exists(index_file_path): msg = f"Index.index {index_file_path} not found in repository." raise FileNotFoundError(msg) - with repository.open(state_file_path, "rb") as file: + with open(state_file_path, "rb") as file: state = np.load(file, allow_pickle=True)["state"][()] config = state["config"] index = cls.from_config(UsearchDenseIndexConfig(**config)) @@ -216,8 +214,8 @@ def _load( index.key_counter = state["id_count"] with ( - repository.open(index_file_path, "rb") as file, - tempfile.NamedTemporaryFile(delete=not mmap) as temp, + open(index_file_path, "rb") as file, + tempfile.NamedTemporaryFile(delete=False) as temp, ): temp.write(file.read()) temp.seek(0) diff --git a/baguetter/indices/sparse/base.py b/baguetter/indices/sparse/base.py index 5665976..e2de3b8 100644 --- a/baguetter/indices/sparse/base.py +++ b/baguetter/indices/sparse/base.py @@ -19,7 +19,6 @@ from collections.abc import Generator from baguetter.types import Key, TextOrTokens - from baguetter.utils.file_repository import AbstractFileRepository class BaseSparseIndex(BaseIndex, abc.ABC): @@ -162,7 +161,6 @@ def vocabulary(self) -> dict[str, int]: def _save( self, path: str, - repository: AbstractFileRepository, ) -> str: """Save the index to the given path. @@ -177,7 +175,7 @@ def _save( "corpus_tokens": self.corpus_tokens, "config": dataclasses.asdict(self.config), } - with repository.open(path, "wb") as f: + with open(path, "wb") as f: np.savez_compressed(f, state=state) return path @@ -185,7 +183,6 @@ def _save( def _load( cls, path: str, - repository: AbstractFileRepository, *, mmap: bool = False, ) -> BaseSparseIndex: @@ -203,12 +200,12 @@ def _load( FileNotFoundError: If the index file is not found. """ - if not repository.exists(path): + if not os.path.exists(path): msg = f"Index {path} not found." raise FileNotFoundError(msg) mmap_mode = "r" if mmap else None - with repository.open(path, "rb") as f: + with open(path, "rb") as f: stored = np.load(f, allow_pickle=True, mmap_mode=mmap_mode) state = stored["state"][()] retriever = cls.from_config(SparseIndexConfig(**state["config"])) diff --git a/baguetter/utils/file_repository.py b/baguetter/utils/file_repository.py index 228ace9..5cedab6 100644 --- a/baguetter/utils/file_repository.py +++ b/baguetter/utils/file_repository.py @@ -1,5 +1,6 @@ from __future__ import annotations +import os from pathlib import Path from fsspec import AbstractFileSystem @@ -103,12 +104,13 @@ def __init__(self, path: str | None = None, **kwargs) -> None: """ super().__init__(**kwargs) - self._base_path = str(Path(path or f"{settings.base_path}/repository").resolve()) - if not self.isdir(self._base_path): - if self.exists(self._base_path): - msg = f"Path '{self._base_path}' exists but is not a directory." + base_path = str(Path(path or f"{settings.base_path}/repository").resolve()) + if not os.path.isdir(base_path): + if os.path.exists(base_path): + msg = f"Path '{base_path}' exists but is not a directory." raise ValueError(msg) - self.mkdir(self._base_path, parents=True, exist_ok=True) + os.makedirs(base_path, exist_ok=True) + self._base_path = base_path def _strip_protocol(self, path: str) -> str: """Strip the protocol from the given path. diff --git a/baguetter/utils/persistable.py b/baguetter/utils/persistable.py index d6baf55..219c829 100644 --- a/baguetter/utils/persistable.py +++ b/baguetter/utils/persistable.py @@ -1,10 +1,10 @@ from __future__ import annotations +import os from abc import ABC, abstractmethod from typing import Any -from baguetter.utils.file_repository import AbstractFileRepository, HuggingFaceFileRepository, LocalFileRepository - +from baguetter.utils.file_repository import HuggingFaceFileRepository class Persistable(ABC): """Abstract base class for objects that can be persisted to and loaded from storage. @@ -18,7 +18,6 @@ class Persistable(ABC): def _load( cls, path: str, - repository: AbstractFileRepository, *, allow_pickle: bool = True, mmap: bool = False, @@ -27,7 +26,6 @@ def _load( Args: path (str): Path of the object to load. - repository (AbstractFileRepository): File repository to load from. allow_pickle (bool, optional): Whether to allow loading pickled objects. Defaults to True. mmap (bool, optional): Whether to memory-map the file. Defaults to False. @@ -37,12 +35,11 @@ def _load( """ @abstractmethod - def _save(self, path: str, repository: AbstractFileRepository) -> str: + def _save(self, path: str) -> str: """Save the object to storage. Args: path (str): Path to save the object to. - repository (AbstractFileRepository): File repository to save to. Returns: str: Path to the saved object. @@ -59,12 +56,8 @@ def save(self, path: str) -> str: str: Path to the saved object. """ - repository = LocalFileRepository() - directory = path.rsplit("/", 1) - if len(directory) > 1: - repository.mkdirs(directory[0], exist_ok=True) - path = self._save(path=path, repository=repository) - return repository.info(path)["name"] + path = self._save(path=path) + return os.path.basename(path) @classmethod def load(cls, path: str, *, mmap: bool = False) -> Any: @@ -78,10 +71,8 @@ def load(cls, path: str, *, mmap: bool = False) -> Any: Any: The loaded object. """ - repository = LocalFileRepository() return cls._load( path=path, - repository=repository, mmap=mmap, ) diff --git a/pyproject.toml b/pyproject.toml index aaf0e7c..7646c92 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,3 +52,8 @@ include = ["baguetter", "baguetter.*"] [tool.setuptools.package-data] baguetter = ["py.typed"] + +[dependency-groups] +dev = [ + "pytest>=8.3.5", +] diff --git a/tests/indices/dense/faiss_test.py b/tests/indices/dense/faiss_test.py index 16540f5..7180bc4 100644 --- a/tests/indices/dense/faiss_test.py +++ b/tests/indices/dense/faiss_test.py @@ -6,7 +6,6 @@ from baguetter.indices.base import SearchResults from baguetter.indices.dense.faiss import FaissDenseIndex -from baguetter.utils.file_repository import LocalFileRepository @pytest.fixture @@ -100,10 +99,9 @@ def test_faiss_save_and_load(sample_data): with tempfile.TemporaryDirectory() as tmpdir: save_path = "faiss_index" - repo = LocalFileRepository(tmpdir) - index._save(repository=repo, path=save_path) + index._save(path=save_path) - loaded_index = FaissDenseIndex._load(save_path, repository=repo) + loaded_index = FaissDenseIndex._load(save_path) assert loaded_index.size == index.size assert loaded_index.config.__dict__ == index.config.__dict__ diff --git a/tests/indices/dense/usearch_test.py b/tests/indices/dense/usearch_test.py index b544e97..b01a510 100644 --- a/tests/indices/dense/usearch_test.py +++ b/tests/indices/dense/usearch_test.py @@ -1,5 +1,6 @@ from __future__ import annotations +import os import tempfile import numpy as np @@ -8,7 +9,6 @@ from baguetter.indices.dense.base import _INDEX_PREFIX from baguetter.indices.dense.config import DenseIndexConfig from baguetter.indices.dense.usearch import USearchDenseIndex -from baguetter.utils.file_repository import LocalFileRepository @pytest.fixture @@ -94,13 +94,12 @@ def test_usearch_save_load(sample_data): with tempfile.TemporaryDirectory() as tmp_dir: save_path = "usearch_new-index" - repository = LocalFileRepository(tmp_dir) # Save the index - index._save(path=save_path, repository=repository) + index._save(path=save_path) # Load the index - loaded_index = USearchDenseIndex._load(path=save_path, repository=repository) + loaded_index = USearchDenseIndex._load(path=save_path) assert loaded_index.config.embedding_dim == 128 assert loaded_index.key_counter == 10 @@ -116,8 +115,8 @@ def test_usearch_save_load(sample_data): ) # Verify file names - assert repository.exists(save_path) - assert repository.exists(f"{_INDEX_PREFIX}{save_path}") + assert os.path.exists(save_path) + assert os.path.exists(f"{_INDEX_PREFIX}{save_path}") def test_usearch_embed_function():