diff --git a/README.md b/README.md index 419095aa4..31193ff87 100644 --- a/README.md +++ b/README.md @@ -63,8 +63,9 @@ enVector supports two types of benchmark cases: ├── README.md ├── scripts │ ├── get_kmeans_centroids.py # create kmeans centroids +│ ├── prepare_dataset.py # download and prepare ground truth neighbors for GAS dataset +│ ├── prepare_random_dataset.py # download and prepare ground truth neighbors for random dataset │ ├── requirements.txt # python requirements -│ ├── prepare_dataset.py # download and prepare ground truth neighbors for dataset │ └── run_benchmark.sh # benchmark script └── vectordb_bench/config-files # benchmark config file └── envector_{benchmark_case}_config.yml @@ -116,7 +117,7 @@ Run the following commands to run enVector with VectorDBBench's built-in benchma ./scripts/run_benchmark.sh --index-type IVF_FLAT --config-file envector_{benchmark_case}_config.yml # IVF-FLAT ``` -For more details, please refer to `envector_{benchmark_case}_config.yml` in scripts directory for benchmarks with enVector, or you can use the following command: +For more details, please refer to `envector_{benchmark_case}_config.yml` in `vectordb_bench/config-files` directory for benchmarks with enVector, or you can use the following command: ```bash python -m vectordb_bench.cli.vectordbbench envectorflat \ @@ -147,12 +148,11 @@ Prepare the following artifacts for the ANN benchmark with `scripts/prepare_data - prepare ground-truth neighbors - download centroids for the GAS index for corresponding to the embedding model -For the ANN benchmark, we provide two datasets via HuggingFace: -- `PUBMED768D400K`: [cryptolab-playground/pubmed-arxiv-abstract-embedding-gemma-300m](https://huggingface.co/datasets/cryptolab-playground/pubmed-arxiv-abstract-embedding-gemma-300m) -- `BLOOMBERG768D368K`: [cryptolab-playground/Bloomberg-Financial-News-embedding-gemma-300m](https://huggingface.co/datasets/cryptolab-playground/Bloomberg-Financial-News-embedding-gemma-300m) -- `PRODUCTS512D400K` -- `FASHION512D200K` -- `FOOD512D75K` +For the ANN benchmark, we provide four datasets via HuggingFace: +- `pubmed768d400k`: [cryptolab-playground/pubmed-arxiv-abstract-embedding-gemma-300m](https://huggingface.co/datasets/cryptolab-playground/pubmed-arxiv-abstract-embedding-gemma-300m) +- `bloomberg768d368k`: [cryptolab-playground/Bloomberg-Financial-News-embedding-gemma-300m](https://huggingface.co/datasets/cryptolab-playground/Bloomberg-Financial-News-embedding-gemma-300m) +- `products512d400k`: [cryptolab-playground/amazon-products-clip-vit-b-32](https://huggingface.co/datasets/cryptolab-playground/amazon-products-clip-vit-b-32) +- `food512d101k`: [cryptolab-playground/food101-clip-vit-b-32](https://huggingface.co/datasets/cryptolab-playground/food101-clip-vit-b-32) Also, we provide centroids for the corresponding embedding model used in the ANN benchmark: - GAS Centroids: [cryptolab-playground/gas-centroids](https://huggingface.co/datasets/cryptolab-playground/gas-centroids) @@ -165,8 +165,7 @@ pip install -r ./scripts/requirements.txt # Prepare GAS dataset python ./scripts/prepare_dataset.py \ - -d cryptolab-playground/pubmed-arxiv-abstract-embedding-gemma-300m \ - -e embeddinggemma-300m + -d pubmed768d400k ``` Then, you can find the generated files as follows: @@ -196,12 +195,12 @@ Run the provided shell scripts (`./scripts/run_benchmark.sh`) as the following: For more details, please refer to `run_benchmark.sh` or `envector_{benchmark_case}_config.yml` in scripts directory for benchmarks with enVector with ANN (GAS), or you can use the following command: ```bash -python -m vectordb_bench.cli.vectordbbench envectorivfflat \ +python -m vectordb_bench.cli.vectordbbench envectorivfgas \ --config-file envector_pubmed_config.yml # or -python -m vectordb_bench.cli.vectordbbench envectorivfflat \ +python -m vectordb_bench.cli.vectordbbench envectorivfgas \ --uri "localhost:50050" \ --eval-mode mm \ ... \ @@ -211,9 +210,6 @@ python -m vectordb_bench.cli.vectordbbench envectorivfflat \ --nprobe 6 ``` -Note that, **`NUM_PER_BATCH` should be set to the database size** when using IVF-based ANN index for enVector currently. -We will support adjustable `NUM_PER_BATCH` for ANN soon. - ## 🎯 Advanced Usage ### Prepare Other Datasets diff --git a/scripts/get_kmeans_centroids.py b/scripts/get_kmeans_centroids.py new file mode 100644 index 000000000..22829b45f --- /dev/null +++ b/scripts/get_kmeans_centroids.py @@ -0,0 +1,95 @@ +""" +Get KMeans centroids for a given dataset. +""" + +import argparse +import os + +import faiss +import numpy as np +import pandas as pd +from numpy.linalg import norm + + +def get_args(): + parser = argparse.ArgumentParser(description="KMeans Centroid Calculation") + parser.add_argument( + "--nlist", + type=int, + default=256, + help="Number of clusters for KMeans", + ) + parser.add_argument( + "--file-path", + type=str, + default="/tmp/vectordb_bench/dataset/openai/openai_medium_500k", + help="Path to the dataset directory", + ) + parser.add_argument( + "--out-path", + type=str, + default="/tmp/vectordb_bench/centroids/kmeans-centroids/openai_medium_500k", + help="Path to the output directory", + ) + parser.add_argument( + "--dim", + type=int, + default=512, + help="Dimension of the embeddings.", + ) + return parser.parse_args() + + +def load_dataset(file_path): + print("Loading dataset from:", file_path) + + # load parquet files + train_vectors = pd.read_parquet(f"{file_path}/train.parquet") + + # sort by id + train_vectors.sort_values(by="id", inplace=True) + train_ids = train_vectors["id"].to_numpy(dtype=np.int64) + train_vectors = np.vstack(train_vectors["emb"].values) + train_vectors /= norm(train_vectors, axis=1, keepdims=True) + print(f"train_vectors shape: {train_vectors.shape}") + + return train_vectors.astype(np.float32) + + +def main(): + args = get_args() + + nlist = args.nlist + seed = 42 + + # prepare dataset + train_vectors = load_dataset(args.file_path) + dim = train_vectors.shape[1] + assert dim == args.dim, f"Expected dimension {args.dim}, but got {dim}" + print("✅ Load dataset complete.") + + # kmeans using faiss + kmeans = faiss.Kmeans(dim, nlist, niter=25, seed=seed, verbose=True, gpu=True) + kmeans.train(train_vectors) + print("✅ KMeans training complete.") + + # allocate + _, labels = kmeans.index.search(train_vectors, 1) + labels = labels.flatten() + centroids = kmeans.centroids + print(f"Labels shape: {labels.shape}") + print(f"Centroids shape: {centroids.shape}") + + # normalize + centroids /= norm(centroids, axis=1, keepdims=True) + print(f"Norm: {norm(centroids, axis=1)}") + + # save centroids + os.makedirs(args.out_path, exist_ok=True) + file_name = os.path.join(args.out_path, f"centroids_{nlist}.npy") + np.save(file_name, centroids) + print(f"✅ Centroids saved to {file_name}") + + +if __name__ == "__main__": + main() diff --git a/scripts/prepare_dataset.py b/scripts/prepare_dataset.py index df0e6d798..8657a4d4e 100644 --- a/scripts/prepare_dataset.py +++ b/scripts/prepare_dataset.py @@ -13,6 +13,23 @@ import wget from datasets import load_dataset +SUPPORTED_CASES = { + "pubmed768d400k": { + "dataset_name": "cryptolab-playground/pubmed-arxiv-abstract-embedding-gemma-300m", + "embedding_model": "embeddinggemma-300m", + }, + "bloomberg768d368k": { + "dataset_name": "cryptolab-playground/Bloomberg-Financial-News-embedding-gemma-300m", + "embedding_model": "embeddinggemma-300m", + }, + "products512d400k": { + "dataset_name": "cryptolab-playground/amazon-products-clip-vit-b-32", + "embedding_model": "clip-vit-b-32", + }, + "food512d101k": {"dataset_name": "cryptolab-playground/food101-clip-vit-b-32", "embedding_model": "clip-vit-b-32"}, +} +SUPPORTED_EMBEDDING_MODELS = ["embeddinggemma-300m", "clip-vit-b-32"] + def get_args(): parser = argparse.ArgumentParser(description="Prepare dataset and ground truth neighbors for benchmarking.") @@ -20,25 +37,15 @@ def get_args(): "-d", "--dataset-name", type=str, - default="cryptolab-playground/pubmed-arxiv-abstract-embedding-gemma-300m", + default="pubmed768d400k", help="Huggingface dataset name to download.", - choices=[ - "cryptolab-playground/pubmed-arxiv-abstract-embedding-gemma-300m", - "cryptolab-playground/Bloomberg-Financial-News-embedding-gemma-300m", - ], + choices=list(SUPPORTED_CASES.keys()), ) parser.add_argument( "--dataset-dir", type=str, - default=os.path.join(os.environ.get("DATASET_LOCAL_DIR", "/tmp/vectordb_bench/dataset"), "pubmed768d400k"), - help="Dataset directory to save the dataset and neighbors. Default: 'pubmed768d400k' in DATASET_LOCAL_DIR.", - ) - parser.add_argument( - "-e", - "--embedding-model", - type=str, - default="embeddinggemma-300m", - help="Embedding model name to download centroids for.", + default=None, + help="Dataset directory to save the dataset and neighbors. Default: in DATASET_LOCAL_DIR.", ) parser.add_argument( "--centroids-dir", @@ -52,7 +59,7 @@ def get_args(): def download_dataset(dataset_name: str, output_dir: str = "./dataset/pubmed768d400k") -> None: """Download dataset from Huggingface and save as Parquet files.""" # load dataset - ds = load_dataset(dataset_name) + ds = load_dataset(SUPPORTED_CASES[dataset_name]["dataset_name"]) train = ds["train"].to_pandas() test = ds["test"].to_pandas() @@ -62,6 +69,7 @@ def download_dataset(dataset_name: str, output_dir: str = "./dataset/pubmed768d4 test_table = pa.Table.from_pandas(test) pq.write_table(test_table, f"{output_dir}/test.parquet") + print(f"Saved train and test parquet data to {output_dir}.") def prepare_neighbors( @@ -89,12 +97,13 @@ def prepare_neighbors( table = pa.Table.from_pandas(df) pq.write_table(table, f"{data_dir}/neighbors.parquet") + print(f"Saved neighbors data to {data_dir}.") def download_centroids(embedding_model: str, dataset_dir: str) -> None: """Download pre-computed centroids and for IVF_GAS index.""" - if embedding_model != "embeddinggemma-300m": + if embedding_model not in SUPPORTED_EMBEDDING_MODELS: raise ValueError(f"Centroids for {embedding_model} currently not available.") # BASE URL: https://huggingface.co/datasets/cryptolab-playground/gas-centroids @@ -103,13 +112,20 @@ def download_centroids(embedding_model: str, dataset_dir: str) -> None: # download os.makedirs(os.path.join(dataset_dir, embedding_model), exist_ok=True) wget.download(f"{dataset_link}/centroids.npy", out=os.path.join(dataset_dir, embedding_model, "centroids.npy")) - print(f"\nDownloaded centroids to {os.path.join(dataset_dir, embedding_model)}") + print(f"\nSaved centroids data to {os.path.join(dataset_dir, embedding_model)}") if __name__ == "__main__": args = get_args() + + base_dataset_dir = ( + os.environ.get("DATASET_LOCAL_DIR", "/tmp/vectordb_bench/dataset") + if args.dataset_dir is None + else args.dataset_dir + ) + args.dataset_dir = os.path.join(base_dataset_dir, args.dataset_name) os.makedirs(args.dataset_dir, exist_ok=True) download_dataset(args.dataset_name, args.dataset_dir) prepare_neighbors(args.dataset_dir) - download_centroids(args.embedding_model, args.centroids_dir) + download_centroids(SUPPORTED_CASES[args.dataset_name]["embedding_model"], args.centroids_dir) diff --git a/scripts/prepare_random_dataset.py b/scripts/prepare_random_dataset.py new file mode 100644 index 000000000..30d8ac6a6 --- /dev/null +++ b/scripts/prepare_random_dataset.py @@ -0,0 +1,140 @@ +""" +Prepare random dataset and ground truth neighbors for test purposes. +""" + +import argparse +import os + +import faiss +import numpy as np +import pandas as pd +import pyarrow as pa +import pyarrow.parquet as pq + + +def get_args(): + parser = argparse.ArgumentParser(description="Prepare random dataset for benchmarking.") + parser.add_argument( + "--dataset-dir", + type=str, + default=os.path.join(os.environ.get("DATASET_LOCAL_DIR", "/tmp/vectordb_bench/dataset"), "random512d1m"), + help="Directory to save the random vectors.", + ) + parser.add_argument( + "--dataset-size", + type=int, + default=1_000_000, + help="Number of dataset embeddings to use.", + ) + parser.add_argument( + "--query-size", + type=int, + default=1_000, + help="Number of query embeddings to use. 1,000 is recommended in VectorDBBench.", + ) + parser.add_argument( + "--dim", + type=int, + default=512, + help="Dimension of the embeddings.", + ) + + return parser.parse_args() + + +def get_random_data(num_data, dim, seed): + rng = np.random.default_rng(seed) + + data = rng.uniform(low=-1.0, high=1.0, size=(num_data, dim)) + + # L2 normalize + norm = np.linalg.norm(data, axis=1, keepdims=True) + norm = np.maximum(norm, 1e-10) + data /= norm + + print(data.shape) + return data.astype(np.float32) + + +def npy_to_parquet( + vector: np.ndarray, + dataset_dir: str = "./dataset/random512d1m", + mode: str = "train", +) -> None: + """Convert downloaded .npy embeddings to Parquet format.""" + print("Preparing embeddings from numpy array...") + os.makedirs(dataset_dir, exist_ok=True) + + ids = np.arange(len(vector)) + id_array = pa.array(ids, type=pa.int64()) + + list_arrays = [vector[i].tolist() for i in range(len(vector))] + vector_array = pa.array(list_arrays, type=pa.list_(pa.float64())) + + assert len(id_array) == len(vector_array) + + table = pa.Table.from_arrays([id_array, vector_array], names=["id", "emb"]) + + out_path = os.path.join(dataset_dir, f"{mode}.parquet") + print(f"Saving parquet to {out_path}") + pq.write_table(table, out_path) + + +def prepare_neighbors( + data_dir: str = "./dataset/random512d1m", +) -> None: + """Prepare ground truth neighbors using brute-force flat search and save as Parquet.""" + # load dataset + train = pd.read_parquet(f"{data_dir}/train.parquet") + test = pd.read_parquet(f"{data_dir}/test.parquet") + + train = np.stack(train["emb"].to_list()).astype("float32") + test = np.stack(test["emb"].to_list()).astype("float32") + dim = train.shape[1] + + # flat search + index = faiss.IndexFlatIP(dim) + index.add(train) + + k = len(test) + distances, indices = index.search(test, k) + print(f"Distances: {distances.shape}, Indices: {indices.shape}") + + assert all(indices[:, 0] == np.arange(len(test))) ### first N vectors + + # save flat search result as neighbors + df = pd.DataFrame({"id": np.arange(len(indices)), "neighbors_id": indices.tolist()}) + + table = pa.Table.from_pandas(df) + pq.write_table(table, f"{data_dir}/neighbors.parquet") + print(f"Saving parquet to {data_dir}/neighbors.parquet") + + +if __name__ == "__main__": + args = get_args() + + # generate random data and save as .npy + vectors = get_random_data( + num_data=args.dataset_size, + dim=args.dim, + seed=42, + ) + + # prepare train parquet file from numpy arrays + npy_to_parquet( + vector=vectors, + dataset_dir=args.dataset_dir, + mode="train", + ) + + # prepare test set + test_vectors = vectors[: args.query_size] ### first N vectors + + npy_to_parquet( + vector=test_vectors, + dataset_dir=args.dataset_dir, + mode="test", + ) + + # prepare neighbors + prepare_neighbors(data_dir=args.dataset_dir) diff --git a/vectordb_bench/backend/clients/api.py b/vectordb_bench/backend/clients/api.py index a2232df17..f385ce610 100644 --- a/vectordb_bench/backend/clients/api.py +++ b/vectordb_bench/backend/clients/api.py @@ -1,13 +1,13 @@ from abc import ABC, abstractmethod from contextlib import contextmanager -from enum import Enum +from enum import StrEnum from pydantic import BaseModel, SecretStr, validator from vectordb_bench.backend.filter import Filter, FilterOp -class MetricType(str, Enum): +class MetricType(StrEnum): L2 = "L2" COSINE = "COSINE" IP = "IP" @@ -16,7 +16,7 @@ class MetricType(str, Enum): JACCARD = "JACCARD" -class IndexType(str, Enum): +class IndexType(StrEnum): HNSW = "HNSW" HNSW_SQ = "HNSW_SQ" HNSW_BQ = "HNSW_BQ" @@ -46,7 +46,7 @@ class IndexType(str, Enum): IVFGAS = "IVF_VCT" # enVector custom index type -class SQType(str, Enum): +class SQType(StrEnum): SQ6 = "SQ6" SQ8 = "SQ8" BF16 = "BF16" diff --git a/vectordb_bench/backend/clients/elastic_cloud/config.py b/vectordb_bench/backend/clients/elastic_cloud/config.py index 4d9ec32d4..2587a6532 100644 --- a/vectordb_bench/backend/clients/elastic_cloud/config.py +++ b/vectordb_bench/backend/clients/elastic_cloud/config.py @@ -1,4 +1,4 @@ -from enum import Enum +from enum import StrEnum from pydantic import BaseModel, SecretStr @@ -16,7 +16,7 @@ def to_dict(self) -> dict: } -class ESElementType(str, Enum): +class ESElementType(StrEnum): float = "float" # 4 byte byte = "byte" # 1 byte, -128 to 127 diff --git a/vectordb_bench/config-files/envector_food_config.yml b/vectordb_bench/config-files/envector_food_config.yml new file mode 100644 index 000000000..ad2853b10 --- /dev/null +++ b/vectordb_bench/config-files/envector_food_config.yml @@ -0,0 +1,46 @@ +# Custom Case +_base_dataset: &base_dataset + case_type: PerformanceCustomDataset + custom_case_name: FOOD512D101K + custom_case_description: FOOD512D101K benchmark (512D, 101K vectors) + custom_dataset_name: FOOD512D101K + custom_dataset_dir: "" + custom_dataset_size: 101000 + custom_dataset_dim: 512 + custom_dataset_file_count: 1 + custom_dataset_use_shuffled: false + custom_dataset_with_gt: true + k: 10 + +# envector server settings +_base_envector: &base_envector + uri: localhost:50050 + eval_mode: mm + drop_old: true + load: true + +# FLAT +envectorflat: + <<: [*base_dataset, *base_envector] + index_name: food101_flat + db_label: FOOD512D101K-FLAT + +# IVF-FLAT with trained k-means centroids +envectorivfflat: + <<: [*base_dataset, *base_envector] + index_name: food101_ivfflat + db_label: FOOD512D101K-IVFFLAT + nlist: 128 + nprobe: 6 + train_centroids: true + centroids_path: food/centroids/centroids_128.npy + +# GAS: enVector-customized ANN +envectorivfgas: + <<: [*base_dataset, *base_envector] + index_name: food101_ivfgas + db_label: FOOD512D101K-IVFGAS + nlist: 1024 + nprobe: 6 + train_centroids: true + centroids_path: centroids/clip-vit-b-32/centroids.npy \ No newline at end of file diff --git a/vectordb_bench/config-files/envector_products_config.yml b/vectordb_bench/config-files/envector_products_config.yml index 550404208..b4977fc8e 100644 --- a/vectordb_bench/config-files/envector_products_config.yml +++ b/vectordb_bench/config-files/envector_products_config.yml @@ -40,7 +40,7 @@ envectorivfgas: <<: [*base_dataset, *base_envector] index_name: products_ivfgas db_label: PRODUCTS512D400K-IVFGAS - nlist: 32768 - nprobe: 6 + nlist: 1024 + nprobe: 16 train_centroids: true centroids_path: centroids/clip-vit-b-32/centroids.npy \ No newline at end of file