Release/v1.3.0 alpha.1 #6
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,136 @@ | ||
| 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) |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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 | ||
|
Comment on lines
+33
to
+43
There was a problem hiding this comment. IVF FLAT / IVF GAS 에서의 nlist 값이 다른데 의도하신걸까요? |
||
| nprobe: 6 | ||
| train_centroids: true | ||
| centroids_path: centroids/clip-vit-b-32/centroids.npy | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Food가 빠져있는 것 같습니다?!
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
이제 추가했습니다!
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
166e5a6