mirror of
https://github.com/AstrBotDevs/AstrBot
synced 2026-07-15 17:30:13 +08:00
79 lines
2.6 KiB
Python
79 lines
2.6 KiB
Python
from typing import Any, overload
|
|
|
|
import numpy as np
|
|
|
|
class Index:
|
|
d: int
|
|
ntotal: int
|
|
code_size: int
|
|
nprobe: int
|
|
|
|
def add(self, x: np.ndarray) -> None: ...
|
|
def add_with_ids(self, x: np.ndarray, ids: np.ndarray) -> None: ...
|
|
def search(
|
|
self,
|
|
x: np.ndarray,
|
|
k: int,
|
|
*,
|
|
params: Any = ...,
|
|
D: np.ndarray | None = ...,
|
|
I: np.ndarray | None = ...,
|
|
) -> tuple[np.ndarray, np.ndarray]: ...
|
|
def remove_ids(self, x: np.ndarray) -> int: ...
|
|
@overload
|
|
def reconstruct(self, key: int) -> np.ndarray: ...
|
|
@overload
|
|
def reconstruct(self, key: int, x: np.ndarray) -> None: ...
|
|
def reconstruct(
|
|
self, key: int, x: np.ndarray | None = ...
|
|
) -> np.ndarray | None: ...
|
|
@overload
|
|
def reconstruct_n(self, n0: int, ni: int) -> np.ndarray: ...
|
|
@overload
|
|
def reconstruct_n(self, n0: int, ni: int, x: np.ndarray) -> None: ...
|
|
def reconstruct_n(
|
|
self, n0: int = ..., ni: int = ..., x: np.ndarray | None = ...
|
|
) -> np.ndarray | None: ...
|
|
def range_search(
|
|
self, x: np.ndarray, thresh: float, *, params: Any = ...
|
|
) -> tuple[np.ndarray, np.ndarray, np.ndarray]: ...
|
|
def add_sa_codes(self, codes: np.ndarray, ids: np.ndarray | None = ...) -> None: ...
|
|
def sa_encode(self, x: np.ndarray) -> np.ndarray: ...
|
|
def sa_decode(self, codes: np.ndarray) -> np.ndarray: ...
|
|
|
|
class IndexFlatL2(Index):
|
|
def __init__(self, d: int) -> None: ...
|
|
|
|
class IndexIDMap(Index):
|
|
index: Index
|
|
|
|
def __init__(self, index: Index) -> None: ...
|
|
|
|
def read_index(path: str) -> Index: ...
|
|
def write_index(index: Index, path: str | None = ...) -> None: ...
|
|
def normalize_L2(x: np.ndarray) -> None: ...
|
|
|
|
# Additional concrete-ish classes exposed by some faiss builds (SWIG helpers
|
|
# expose `downcast_*` helpers to convert generic objects to these concrete
|
|
# types). We keep these minimal — only the names are important for typing.
|
|
class IndexBinary(Index):
|
|
def __init__(self, d: int) -> None: ...
|
|
|
|
class InvertedLists:
|
|
def __len__(self) -> int: ...
|
|
|
|
class AdditiveQuantizer: ...
|
|
class Quantizer: ...
|
|
class VectorTransform: ...
|
|
|
|
# SWIG-provided downcast helpers (present in some faiss Python builds).
|
|
def downcast_IndexBinary(obj: Any) -> IndexBinary: ...
|
|
def downcast_InvertedLists(obj: Any) -> InvertedLists: ...
|
|
def downcast_AdditiveQuantizer(obj: Any) -> AdditiveQuantizer: ...
|
|
def downcast_Quantizer(obj: Any) -> Quantizer: ...
|
|
def downcast_VectorTransform(obj: Any) -> VectorTransform: ...
|
|
def downcast_index(obj: Any) -> Index: ...
|
|
|
|
# version exposed by runtime
|
|
__version__: str
|