Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 18 additions & 3 deletions src/python_bindings/bindings.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -594,10 +594,15 @@ class PyTiered_SVSIndex : public PyTieredIndex {
explicit PyTiered_SVSIndex(const SVSParams &svs_params,
const TieredSVSParams &tiered_svs_params, size_t buffer_limit) {

// Create primaryIndexParams and specific params for hnsw tiered index.
// Create primaryIndexParams and specific params for svs tiered index.
VecSimParams primary_index_params = {.algo = VecSimAlgo_SVS,
.algoParams = {.svsParams = svs_params}};

if (primary_index_params.algoParams.svsParams.num_threads == 0) {
primary_index_params.algoParams.svsParams.num_threads =
this->mock_thread_pool.thread_pool_size; // Use the mock thread pool size as default
}

auto tiered_params = this->getTieredIndexParams(buffer_limit);
tiered_params.primaryIndexParams = &primary_index_params;
tiered_params.specificParams.tieredSVSParams = tiered_svs_params;
Expand All @@ -611,6 +616,10 @@ class PyTiered_SVSIndex : public PyTieredIndex {
// Set the created tiered index in the index external context.
this->mock_thread_pool.ctx->index_strong_ref = this->index;
}

size_t SVSLabelCount() {
return this->index->debugInfo().tieredInfo.backendCommonInfo.indexLabelCount;
}
};
#endif

Expand Down Expand Up @@ -672,10 +681,13 @@ PYBIND11_MODULE(VecSim, m) {

py::enum_<VecSimSvsQuantBits>(m, "VecSimSvsQuantBits")
.value("VecSimSvsQuant_NONE", VecSimSvsQuant_NONE)
.value("VecSimSvsQuant_8", VecSimSvsQuant_8)
.value("VecSimSvsQuant_Scalar", VecSimSvsQuant_Scalar)
.value("VecSimSvsQuant_4", VecSimSvsQuant_4)
.value("VecSimSvsQuant_8", VecSimSvsQuant_8)
.value("VecSimSvsQuant_4x4", VecSimSvsQuant_4x4)
.value("VecSimSvsQuant_4x8", VecSimSvsQuant_4x8)
.value("VecSimSvsQuant_4x8_LeanVec", VecSimSvsQuant_4x8_LeanVec)
.value("VecSimSvsQuant_8x8_LeanVec", VecSimSvsQuant_8x8_LeanVec)
.export_values();

py::class_<SVSParams>(m, "SVSParams")
Expand Down Expand Up @@ -705,6 +717,7 @@ PYBIND11_MODULE(VecSim, m) {
py::class_<TieredSVSParams>(m, "TieredSVSParams")
.def(py::init())
.def_readwrite("trainingTriggerThreshold", &TieredSVSParams::trainingTriggerThreshold)
.def_readwrite("updateTriggerThreshold", &TieredSVSParams::updateTriggerThreshold)
.def_readwrite("updateJobWaitTime", &TieredSVSParams::updateJobWaitTime);

py::class_<AlgoParams>(m, "AlgoParams")
Expand Down Expand Up @@ -799,7 +812,9 @@ PYBIND11_MODULE(VecSim, m) {
size_t flat_buffer_size = DEFAULT_BLOCK_SIZE) {
return new PyTiered_SVSIndex(svs_params, tiered_svs_params, flat_buffer_size);
}),
py::arg("svs_params"), py::arg("tiered_svs_params"), py::arg("flat_buffer_size"));
py::arg("svs_params"), py::arg("tiered_svs_params"),
py::arg("flat_buffer_size") = DEFAULT_BLOCK_SIZE)
.def("svs_label_count", &PyTiered_SVSIndex::SVSLabelCount);
#endif

py::class_<PyBatchIterator>(m, "BatchIterator")
Expand Down
22 changes: 22 additions & 0 deletions tests/flow/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,28 @@
import math
from ml_dtypes import bfloat16

# alpha = 0 means the default value, which is 1.2 for L2 and 0.9 for other metrics.
def create_svs_params (dim, num_elements, data_type, metric, quantBits = VecSimSvsQuant_NONE, alpha = 0,
graph_max_degree = 32, construction_window_size = 200, search_window_size = 10,
max_candidate_pool_size = 0, prune_to = 0, epsilon = 0.01, num_threads = 0):
svs_params = SVSParams()

svs_params.dim = dim
svs_params.type = data_type
svs_params.metric = metric
svs_params.quantBits = quantBits
svs_params.alpha = alpha
svs_params.graph_max_degree = graph_max_degree
svs_params.construction_window_size = construction_window_size
svs_params.max_candidate_pool_size = max_candidate_pool_size
svs_params.prune_to = prune_to
svs_params.use_search_history = VecSimOption_AUTO # VecSimOption_AUTO means use the default value
svs_params.search_window_size = search_window_size
svs_params.epsilon = epsilon
svs_params.num_threads = num_threads

return svs_params

def create_hnsw_params(dim, num_elements, metric, data_type, ef_construction=200, m=16, ef_runtime=10, epsilon=0.01,
is_multi=False):
hnsw_params = HNSWParams()
Expand Down
Loading
Loading