Update int8 single hnsw + structural changes (#601)

* Change dataset to 1 mil

* Change dataset to 1 mil batch

* change file

* change dataset url to 1mil all

* change dataset url to 1mil int8

* restructer python scripts in data, add int8 to serializer, add batch-iter-int8 to yml

(cherry picked from commit 5a86c90)

Author: lerman25

.github/workflows/benchmark.yml

@@ -11,7 +11,6 @@ on:
         options:
           - benchmarks-all
           - benchmarks-default
-          - bm-basics-int8-single
           - bm-basics-fp32-single
           - bm-basics-fp32-multi
           - bm-basics-fp64-single
@@ -20,6 +19,7 @@ on:
           - bm-basics-bf16-multi
           - bm-basics-fp16-single
           - bm-basics-fp16-multi
+          - bm-basics-int8-single
           - bm-batch-iter-fp32-single
           - bm-batch-iter-fp32-multi
           - bm-batch-iter-fp64-single
@@ -28,6 +28,7 @@ on:
           - bm-batch-iter-bf16-multi
           - bm-batch-iter-fp16-single
           - bm-batch-iter-fp16-multi
+          - bm-batch-iter-int8-single
           - bm-updated-fp32-single
           - bm-spaces
         description: 'Benchmarks set to run'

tests/benchmark/data/hnsw_indices/hnsw_indices_all.txt

@@ -25,5 +25,5 @@ https://dev.cto.redis.s3.amazonaws.com/VectorSimilarity/dbpedia-cosine-dim768-fp
 https://dev.cto.redis.s3.amazonaws.com/VectorSimilarity/fashion_images_multi_value-cosine-dim512-M64-efc512-fp16.hnsw_v3
 https://dev.cto.redis.s3.amazonaws.com/VectorSimilarity/fashion_images_multi_value-cosine-dim512-fp16-test_vectors.raw
 
-https://dev.cto.redis.s3.amazonaws.com/VectorSimilarity/wipedia-cosine-dim1024-M64-efc512-int8.hnsw_v3 
-https://dev.cto.redis.s3.amazonaws.com/VectorSimilarity/wipedia-cosine-dim1024-int8-test_vectors.raw
+https://dev.cto.redis.s3.amazonaws.com/VectorSimilarity/wipedia_single-cosine-dim1024-M64-efc512-int8.hnsw_v3 
+https://dev.cto.redis.s3.amazonaws.com/VectorSimilarity/wipedia_single-cosine-dim1024-int8-test_vectors.raw

tests/benchmark/data/hnsw_indices/hnsw_indices_basic_int8.txt

@@ -1,2 +1,2 @@
-https://dev.cto.redis.s3.amazonaws.com/VectorSimilarity/wipedia-cosine-dim1024-M64-efc512-int8.hnsw_v3 
-https://dev.cto.redis.s3.amazonaws.com/VectorSimilarity/wipedia-cosine-dim1024-int8-test_vectors.raw
+https://dev.cto.redis.s3.amazonaws.com/VectorSimilarity/wipedia_single-cosine-dim1024-M64-efc512-int8.hnsw_v3 
+https://dev.cto.redis.s3.amazonaws.com/VectorSimilarity/wipedia_single-cosine-dim1024-int8-test_vectors.raw

tests/benchmark/data/scripts/download_dataset.py

@@ -0,0 +1,50 @@
+# This file is a template file for downloading datasets
+# In this version, it downloads the "wipedia_single" dataset used for int8
+# Refrain from pushing changes unless necessary
+from datasets import load_dataset
+import numpy as np
+import os
+import h5py
+from math import ceil
+from tqdm import tqdm
+INT8_KEY = 'emb_int8'
+DATASET = 'wikipedia-1024_eng_v3_single'
+hdf5_output_file_name = "%s.hdf5" %DATASET
+
+lang = "en" #Use the English Wikipedia subset
+
+num_vectors_train = 1_000_000
+num_vectors_test = 10_000
+num_vectors = num_vectors_train + num_vectors_test
+
+dim = 1024
+docs = load_dataset("Cohere/wikipedia-2023-11-embed-multilingual-v3-int8-binary", lang, split="train", streaming=True)
+label_size = 1
+data = np.empty((num_vectors//label_size,label_size, dim), dtype=np.int8)
+
+
+ids = []
+counter = 0
+label_index = 0
+
+with tqdm(total=num_vectors) as progress_bar:
+    for doc in docs:
+        if counter == num_vectors:
+            break
+        ids.append(doc['_id'])
+        emb = doc['emb_int8']
+        data[label_index, counter % label_size] = emb
+        counter += 1
+        if counter % label_size == 0:
+            label_index += 1
+        progress_bar.update(1)
+
+train_data = data[:num_vectors_train // label_size].reshape(-1, dim)
+test_data = data[num_vectors_train // label_size:].reshape(-1, dim)
+
+print(f"Train data shape: {train_data.shape}")
+print(f"Test data shape: {test_data.shape}")
+
+with h5py.File(hdf5_output_file_name, 'w') as hdf5_file:
+    hdf5_file.create_dataset('train', data=train_data)
+    hdf5_file.create_dataset('test', data=test_data)

tests/benchmark/data/serializer.py …sts/benchmark/data/scripts/serializer.pytests/benchmark/data/serializer.py renamed to tests/benchmark/data/scripts/serializer.py tests/benchmark/data/serializer.py renamed to tests/benchmark/data/scripts/serializer.py

@@ -136,13 +136,21 @@
         'metric': VecSimMetric_L2,
         'skipRaw': True,
     },
+    {
+        'filename': 'wikipedia_single-1024_eng_v3',
+        'nickname': 'wipedia_single',
+        'dim': 1024,
+        'type': VecSimType_INT8,
+        'metric': VecSimMetric_Cosine,
+    },
 ]
 
 TYPES_ATTR = {
     VecSimType_FLOAT32: {"size_in_bytes": 4, "vector_type": np.float32},
     VecSimType_FLOAT64: {"size_in_bytes": 8, "vector_type": np.float64},
     VecSimType_BFLOAT16: {"size_in_bytes": 2, "vector_type": bfloat16},
     VecSimType_FLOAT16: {"size_in_bytes": 2, "vector_type": np.float16},
+    VecSimType_INT8: {"size_in_bytes": 1, "vector_type": np.int8}
 }
 
 
@@ -190,7 +198,9 @@ def serialize(files=DEFAULT_FILES):
             elif hnswparams.type == VecSimType_FLOAT16:
                 serialized_file_name = serialized_file_name + '-fp16'
                 serialized_raw_name = serialized_raw_name + '-fp16'
-
+            elif hnswparams.type == VecSimType_INT8:
+                serialized_file_name = serialized_file_name + '-int8'
+                serialized_raw_name = serialized_raw_name + '-int8'
             print('first, exporting test set to binary')
             if not file.get('skipRaw', False):
                 test = f['test']
@@ -201,6 +211,8 @@ def serialize(files=DEFAULT_FILES):
                     test = np.array(test_set, dtype=bfloat16)
                 elif hnswparams.type == VecSimType_FLOAT16:
                     test = test.astype(np.float16)
+                elif hnswparams.type == VecSimType_INT8:
+                    test = test.astype(np.int8)
                 print(f"creating test set of {len(test)} vectors")
                 with open(os.path.join(location, serialized_raw_name + '-test_vectors.raw'), 'wb') as testfile:
                     for vec in test:
@@ -222,6 +234,8 @@ def serialize(files=DEFAULT_FILES):
                 data = np.array(data_set, dtype=bfloat16)
             elif hnswparams.type == VecSimType_FLOAT16:
                 data = data.astype(np.float16)
+            elif hnswparams.type == VecSimType_INT8:
+                data = data.astype(np.int8)
             print(f"creating index with {hnswparams.initialCapacity} vectors")
             for label, cur in enumerate(data):
                 for vec in cur if hnswparams.multi else [cur]:

tests/benchmark/run_files/bm_basics_single_int8.cpp

@@ -9,16 +9,16 @@
 bool BM_VecSimGeneral::is_multi = false;
 
 size_t BM_VecSimGeneral::n_queries = 10000;
-size_t BM_VecSimGeneral::n_vectors = 999424;
+size_t BM_VecSimGeneral::n_vectors = 1000000;
 size_t BM_VecSimGeneral::dim = 1024;
 size_t BM_VecSimGeneral::M = 64;
 size_t BM_VecSimGeneral::EF_C = 512;
 tieredIndexMock BM_VecSimGeneral::mock_thread_pool{};
 
 const char *BM_VecSimGeneral::hnsw_index_file =
-    "tests/benchmark/data/wipedia-cosine-dim1024-M64-efc512-int8.hnsw_v3";
+    "tests/benchmark/data/wipedia_single-cosine-dim1024-M64-efc512-int8.hnsw_v3";
 const char *BM_VecSimGeneral::test_queries_file =
-    "tests/benchmark/data/wipedia-cosine-dim1024-int8-test_vectors.raw";
+    "tests/benchmark/data/wipedia_single-cosine-dim1024-int8-test_vectors.raw";
 
 #define BM_FUNC_NAME(bm_func, algo) bm_func##_##algo##_Single
 #define BM_ADD_LABEL                AddLabel_Single

tests/benchmark/run_files/bm_batch_iterator_single_int8.cpp

@@ -3,17 +3,17 @@
 bool BM_VecSimGeneral::is_multi = false;
 
 size_t BM_VecSimGeneral::n_queries = 10000;
-size_t BM_VecSimGeneral::n_vectors = 999424;
+size_t BM_VecSimGeneral::n_vectors = 1000000;
 size_t BM_VecSimGeneral::dim = 1024;
 size_t BM_VecSimGeneral::M = 64;
 size_t BM_VecSimGeneral::EF_C = 512;
 size_t BM_VecSimGeneral::block_size = 1024;
 tieredIndexMock BM_VecSimGeneral::mock_thread_pool{};
 
 const char *BM_VecSimGeneral::hnsw_index_file =
-    "tests/benchmark/data/wipedia-cosine-dim1024-M64-efc512-int8.hnsw_v3";
+    "tests/benchmark/data/wipedia_single-cosine-dim1024-M64-efc512-int8.hnsw_v3";
 const char *BM_VecSimGeneral::test_queries_file =
-    "tests/benchmark/data/wipedia-cosine-dim1024-int8-test_vectors.raw";
+    "tests/benchmark/data/wipedia_single-cosine-dim1024-int8-test_vectors.raw";
 
 #define BM_FUNC_NAME(bm_func, algo) algo##_##bm_func##_Single