redisbench-admin 0.11.63__py3-none-any.whl → 0.11.65__py3-none-any.whl

redisbench_admin/run/ann/pkg/ann_benchmarks/results.py

@@ -0,0 +1,71 @@
+from __future__ import absolute_import
+
+import h5py
+import json
+import os
+import re
+import traceback
+
+
+def get_result_filename(dataset=None, count=None, definition=None,
+                        query_arguments=None, batch_mode=False, id=0):
+    d = ['results']
+    if dataset:
+        d.append(dataset)
+    if count:
+        d.append(str(count))
+    if definition:
+        d.append(definition.algorithm + ('-batch' if batch_mode else ''))
+        data = definition.arguments + query_arguments
+        for i in range(len(data)):
+            if isinstance(data[i], dict):
+                data[i] = {k:data[i][k] for k in data[i] if data[i][k] is not None and k != 'auth'}
+        data.append('client')
+        data.append(id)
+        d.append(re.sub(r'\W+', '_', json.dumps(data, sort_keys=True)).strip('_') + ".hdf5")
+    return os.path.join(*d)
+
+
+def store_results(dataset, count, definition, query_arguments, attrs, results,
+                  batch, id=0):
+    fn = get_result_filename(
+        dataset, count, definition, query_arguments, batch, id)
+    head, tail = os.path.split(fn)
+    if not os.path.isdir(head):
+        os.makedirs(head)
+    f = h5py.File(fn, 'w')
+    for k, v in attrs.items():
+        f.attrs[k] = v
+    times = f.create_dataset('times', (len(results),), 'f')
+    neighbors = f.create_dataset('neighbors', (len(results), count), 'i')
+    distances = f.create_dataset('distances', (len(results), count), 'f')
+    for i, (time, ds) in enumerate(results):
+        times[i] = time
+        neighbors[i] = [n for n, d in ds] + [-1] * (count - len(ds))
+        distances[i] = [d for n, d in ds] + [float('inf')] * (count - len(ds))
+    f.close()
+
+
+def load_all_results(dataset=None, count=None, batch_mode=False):
+    for root, _, files in os.walk(get_result_filename(dataset, count)):
+        for fn in files:
+            if os.path.splitext(fn)[-1] != '.hdf5':
+                continue
+            try:
+                f = h5py.File(os.path.join(root, fn), 'r+')
+                properties = dict(f.attrs)
+                if batch_mode != properties['batch_mode']:
+                    continue
+                yield properties, f
+                f.close()
+            except:
+                print('Was unable to read', fn)
+                traceback.print_exc()
+
+
+def get_unique_algorithms():
+    algorithms = set()
+    for batch_mode in [False, True]:
+        for properties, _ in load_all_results(batch_mode=batch_mode):
+            algorithms.add(properties['algo'])
+    return algorithms

redisbench_admin/run/ann/pkg/ann_benchmarks/runner.py

@@ -0,0 +1,333 @@
+import argparse
+import json
+import logging
+import os
+import threading
+import time
+import traceback
+import inspect
+import h5py
+
+import colors
+import docker
+import numpy
+import psutil
+
+from ann_benchmarks.algorithms.definitions import (Definition,
+                                                   instantiate_algorithm)
+from ann_benchmarks.datasets import get_dataset, DATASETS
+from ann_benchmarks.distance import metrics, dataset_transform
+from ann_benchmarks.results import get_result_filename, store_results
+
+
+def run_individual_query(algo, X_train, X_test, distance, count, run_count,
+                         batch):
+    prepared_queries = \
+        (batch and hasattr(algo, "prepare_batch_query")) or \
+        ((not batch) and hasattr(algo, "prepare_query"))
+
+    best_search_time = float('inf')
+    for i in range(run_count):
+        print('Run %d/%d...' % (i + 1, run_count))
+        # a bit dumb but can't be a scalar since of Python's scoping rules
+        n_items_processed = [0]
+
+        def single_query(v):
+            if prepared_queries:
+                algo.prepare_query(v, count)
+                start = time.time()
+                algo.run_prepared_query()
+                total = (time.time() - start)
+                candidates = algo.get_prepared_query_results()
+            else:
+                start = time.time()
+                candidates = algo.query(v, count)
+                total = (time.time() - start)
+            candidates = [(int(idx), float(metrics[distance]['distance'](v, X_train[idx])))  # noqa
+                          for idx in candidates]
+            n_items_processed[0] += 1
+            if n_items_processed[0] % 1000 == 0:
+                print('Processed %d/%d queries...' % (n_items_processed[0], len(X_test)))
+            if len(candidates) > count:
+                print('warning: algorithm %s returned %d results, but count'
+                      ' is only %d)' % (algo, len(candidates), count))
+            return (total, candidates)
+
+        def batch_query(X):
+            if prepared_queries:
+                algo.prepare_batch_query(X, count)
+                start = time.time()
+                algo.run_batch_query()
+                total = (time.time() - start)
+            else:
+                start = time.time()
+                algo.batch_query(X, count)
+                total = (time.time() - start)
+            results = algo.get_batch_results()
+            candidates = [[(int(idx), float(metrics[distance]['distance'](v, X_train[idx])))  # noqa
+                           for idx in single_results]
+                          for v, single_results in zip(X, results)]
+            return [(total / float(len(X)), v) for v in candidates]
+
+        if batch:
+            results = batch_query(X_test)
+        else:
+            results = [single_query(x) for x in X_test]
+
+        total_time = sum(time for time, _ in results)
+        total_candidates = sum(len(candidates) for _, candidates in results)
+        search_time = total_time / len(X_test)
+        avg_candidates = total_candidates / len(X_test)
+        best_search_time = min(best_search_time, search_time)
+        print("qps:", len(X_test)/total_time)
+
+    verbose = hasattr(algo, "query_verbose")
+    attrs = {
+        "batch_mode": batch,
+        "best_search_time": best_search_time,
+        "candidates": avg_candidates,
+        "expect_extra": verbose,
+        "name": str(algo),
+        "run_count": run_count,
+        "distance": distance,
+        "count": int(count)
+    }
+    additional = algo.get_additional()
+    for k in additional:
+        attrs[k] = additional[k]
+    return (attrs, results)
+
+
+def run(definition, dataset, count, run_count, batch, build_only, test_only, num_clients, id):
+    algo = instantiate_algorithm(definition)
+    assert not definition.query_argument_groups \
+           or hasattr(algo, "set_query_arguments"), """\
+error: query argument groups have been specified for %s.%s(%s), but the \
+algorithm instantiated from it does not implement the set_query_arguments \
+function""" % (definition.module, definition.constructor, definition.arguments)
+
+    D, dimension = get_dataset(dataset)
+    X_train, X_test = dataset_transform(D)
+    distance = D.attrs['distance']
+    print('got a train set of size (%d * %d)' % (X_train.shape[0], dimension))
+
+    hybrid_buckets = None
+    if 'bucket_names' in D.attrs:
+        hybrid_buckets = {}
+        bucket_names = D.attrs['bucket_names']
+        for bucket_name in bucket_names:
+            bucket_dict = {}
+            bucket_dict['ids'] = numpy.array(D[f'{bucket_name}_ids'])
+            bucket_dict['text'] = D[bucket_name]['text'][()]
+            bucket_dict['number'] = D[bucket_name]['number'][()]
+            hybrid_buckets[bucket_name] = bucket_dict
+
+    try:
+        prepared_queries = False
+        if hasattr(algo, "supports_prepared_queries"):
+            prepared_queries = algo.supports_prepared_queries()
+
+        if not test_only:
+            per_client = len(X_train) // num_clients
+            offset = per_client * (id - 1)
+            fit_kwargs = {}
+            if "offset" and "limit" in inspect.getfullargspec(algo.fit)[0]:
+                fit_kwargs['offset']=offset
+                if num_clients != id:
+                    fit_kwargs['limit']=offset + per_client
+            if hybrid_buckets:
+                fit_kwargs['hybrid_buckets']=hybrid_buckets
+
+            t0 = time.time()
+            memory_usage_before = algo.get_memory_usage()
+            algo.fit(X_train, **fit_kwargs)
+            build_time = time.time() - t0
+            index_size = algo.get_memory_usage() - memory_usage_before
+            print('Built index in', build_time)
+            print('Index size: ', index_size)
+
+        query_argument_groups = definition.query_argument_groups
+        # Make sure that algorithms with no query argument groups still get run
+        # once by providing them with a single, empty, harmless group
+        if not query_argument_groups:
+            query_argument_groups = [[]]
+
+        if not build_only:
+            print('got %d queries' % len(X_test))
+            per_client = len(X_test) // num_clients
+            offset = per_client * (id - 1)
+            if (num_clients != id):
+                X_test = X_test[offset : offset + per_client]
+            else:
+                X_test = X_test[offset:]
+            print('running %d out of them' % len(X_test))
+
+            for pos, query_arguments in enumerate(query_argument_groups, 1):
+                print("Running query argument group %d of %d..." %
+                      (pos, len(query_argument_groups)))
+                if query_arguments:
+                    algo.set_query_arguments(*query_arguments)
+                if hybrid_buckets:
+                    text = hybrid_buckets[D.attrs['selected_bucket']]['text'].decode()
+                    print("setting hybrid text query", text)
+                    algo.set_hybrid_query(text)
+                descriptor, results = run_individual_query(
+                    algo, X_train, X_test, distance, count, run_count, batch)
+                if test_only:
+                    try:
+                        fn = get_result_filename(dataset, count)
+                        fn = os.path.join(fn, definition.algorithm, 'build_stats')
+                        f = h5py.File(fn, 'r')
+                        descriptor["build_time"] = f.attrs["build_time"]
+                        descriptor["index_size"] = f.attrs["index_size"]
+                        f.close()
+                    except:
+                        descriptor["build_time"] = 0
+                        descriptor["index_size"] = 0
+                else:
+                    descriptor["build_time"] = build_time
+                    descriptor["index_size"] = index_size
+                descriptor["algo"] = definition.algorithm
+                descriptor["dataset"] = dataset
+                store_results(dataset, count, definition, query_arguments,
+                              descriptor, results, batch, id)
+    finally:
+        algo.done()
+
+
+def run_from_cmdline():
+    parser = argparse.ArgumentParser('''
+
+            NOTICE: You probably want to run.py rather than this script.
+
+''')
+    parser.add_argument(
+        '--dataset',
+        choices=DATASETS.keys(),
+        help=f'Dataset to benchmark on.',
+        required=True)
+    parser.add_argument(
+        '--algorithm',
+        help='Name of algorithm for saving the results.',
+        required=True)
+    parser.add_argument(
+        '--module',
+        help='Python module containing algorithm. E.g. "ann_benchmarks.algorithms.annoy"',
+        required=True)
+    parser.add_argument(
+        '--constructor',
+        help='Constructer to load from modulel. E.g. "Annoy"',
+        required=True)
+    parser.add_argument(
+        '--count',
+        help='K: Number of nearest neighbours for the algorithm to return.',
+        required=True,
+        type=int)
+    parser.add_argument(
+        '--runs',
+        help='Number of times to run the algorihm. Will use the fastest run-time over the bunch.',
+        required=True,
+        type=int)
+    parser.add_argument(
+        '--batch',
+        help='If flag included, algorithms will be run in batch mode, rather than "individual query" mode.',
+        action='store_true')
+    parser.add_argument(
+        '--build-only',
+        action='store_true',
+        help='building index only, not testing with queries')
+    parser.add_argument(
+        '--test-only',
+        action='store_true',
+        help='querying index only, not building it (should be built first)')
+    parser.add_argument(
+        'build',
+        help='JSON of arguments to pass to the constructor. E.g. ["angular", 100]'
+        )
+    parser.add_argument(
+        'queries',
+        help='JSON of arguments to pass to the queries. E.g. [100]',
+        nargs='*',
+        default=[])
+    args = parser.parse_args()
+    algo_args = json.loads(args.build)
+    print(algo_args)
+    query_args = [json.loads(q) for q in args.queries]
+
+    definition = Definition(
+        algorithm=args.algorithm,
+        docker_tag=None,  # not needed
+        module=args.module,
+        constructor=args.constructor,
+        arguments=algo_args,
+        query_argument_groups=query_args,
+        disabled=False
+    )
+    run(definition, args.dataset, args.count, args.runs, args.batch, args.build_only, args.test_only, 1, 1)
+
+
+def run_docker(definition, dataset, count, runs, timeout, batch, cpu_limit,
+               mem_limit=None):
+    cmd = ['--dataset', dataset,
+           '--algorithm', definition.algorithm,
+           '--module', definition.module,
+           '--constructor', definition.constructor,
+           '--runs', str(runs),
+           '--count', str(count)]
+    if batch:
+        cmd += ['--batch']
+    cmd.append(json.dumps(definition.arguments))
+    cmd += [json.dumps(qag) for qag in definition.query_argument_groups]
+
+    client = docker.from_env()
+    if mem_limit is None:
+        mem_limit = psutil.virtual_memory().available
+
+    container = client.containers.run(
+        definition.docker_tag,
+        cmd,
+        volumes={
+            os.path.abspath('ann_benchmarks'):
+                {'bind': '/home/app/ann_benchmarks', 'mode': 'ro'},
+            os.path.abspath('data'):
+                {'bind': '/home/app/data', 'mode': 'ro'},
+            os.path.abspath('results'):
+                {'bind': '/home/app/results', 'mode': 'rw'},
+        },
+        cpuset_cpus=cpu_limit,
+        mem_limit=mem_limit,
+        detach=True)
+    logger = logging.getLogger(f"annb.{container.short_id}")
+
+    logger.info('Created container %s: CPU limit %s, mem limit %s, timeout %d, command %s' % \
+                (container.short_id, cpu_limit, mem_limit, timeout, cmd))
+
+    def stream_logs():
+        for line in container.logs(stream=True):
+            logger.info(colors.color(line.decode().rstrip(), fg='blue'))
+
+    t = threading.Thread(target=stream_logs, daemon=True)
+    t.start()
+
+    try:
+        return_value = container.wait(timeout=timeout)
+        _handle_container_return_value(return_value, container, logger)
+    except:
+        logger.error('Container.wait for container %s failed with exception' % container.short_id)
+        traceback.print_exc()
+    finally:
+        container.remove(force=True)
+
+def _handle_container_return_value(return_value, container, logger):
+    base_msg = 'Child process for container %s' % (container.short_id)
+    if type(return_value) is dict: # The return value from container.wait changes from int to dict in docker 3.0.0
+        error_msg = return_value['Error']
+        exit_code = return_value['StatusCode']
+        msg = base_msg + 'returned exit code %d with message %s' %(exit_code, error_msg)
+    else:
+        exit_code = return_value
+        msg = base_msg + 'returned exit code %d' % (exit_code)
+
+    if exit_code not in [0, None]:
+        logger.error(colors.color(container.logs().decode(), fg='red'))
+        logger.error(msg)

redisbench_admin/run/ann/pkg/create_dataset.py

@@ -0,0 +1,12 @@
+import argparse
+from ann_benchmarks.datasets import DATASETS, get_dataset_fn
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        '--dataset',
+        choices=DATASETS.keys(),
+        required=True)
+    args = parser.parse_args()
+    fn = get_dataset_fn(args.dataset)
+    DATASETS[args.dataset](fn)

redisbench_admin/run/ann/pkg/create_hybrid_dataset.py

@@ -0,0 +1,147 @@
+from operator import ne
+import click
+from ann_benchmarks.datasets import get_dataset, DATASETS
+from ann_benchmarks.algorithms.bruteforce import BruteForceBLAS
+import struct
+import numpy as np
+import click
+import h5py
+from joblib import Parallel, delayed
+import multiprocessing
+import scipy.spatial
+
+def calc_i(i, x, bf, test, neighbors, distances, count, orig_ids):
+    if i % 1000 == 0:
+        print('%d/%d...' % (i, len(test)))
+    res = list(bf.query_with_distances(x, count))
+    res.sort(key=lambda t: t[-1])
+    neighbors[i] = [orig_ids[j] for j, _ in res]
+    distances[i] = [d for _, d in res]
+
+def create_buckets(train):
+    bucket_0_5 = []
+    bucket_1 = []
+    bucket_2 = []
+    bucket_5 = []
+    bucket_10 = []
+    bucket_20 = []
+    bucket_50 = []
+    other_bucket = []
+    buckets = [bucket_0_5, bucket_1, bucket_2, bucket_5, bucket_10, bucket_20, bucket_50, other_bucket]
+    bucket_names=['0.5', '1', '2', '5', '10', '20', '50', 'other']
+    for i in range(train.shape[0]):
+        if i % 200 == 19:      # 0.5%
+            bucket_0_5.append(i)
+        elif i % 100 == 17:    # 1%
+            bucket_1.append(i)
+        elif i % 50 == 9:     # 2%
+            bucket_2.append(i)
+        elif i % 20 == 7:     # 5%
+            bucket_5.append(i)
+        elif i % 10 == 3:     # 10%
+            bucket_10.append(i)
+        elif i % 2 == 0:      # 50%
+            bucket_50.append(i)
+        elif i % 5 <= 1:      # 20%
+            bucket_20.append(i)
+        else:
+            other_bucket.append(i)
+    print(len(bucket_0_5), len(bucket_1), len(bucket_2), len(bucket_5), len(bucket_10), len(bucket_20), len(bucket_50), len(other_bucket))
+    numeric_values = {}
+    text_values = {}
+    for i, bucket_name in enumerate(bucket_names):
+        numeric_values[bucket_name] = i
+        text_values[bucket_name] = f'text_{i}'
+    print(numeric_values)
+    print(text_values)
+    return buckets, bucket_names, numeric_values, text_values
+
+@click.command()
+@click.option('--data_set', type=click.Choice(DATASETS.keys(), case_sensitive=False), default='glove-100-angular')
+@click.option('--percentile', type=click.Choice(['0.5', '1', '2', '5', '10', '20', '50'], case_sensitive=False), default=None)
+def create_ds(data_set, percentile):
+    ds, dimension= get_dataset(data_set)
+    train = ds['train']
+    test = ds['test']
+    distance = ds.attrs['distance']
+    count=len(ds['neighbors'][0])
+    print(count)
+    print(train.shape)
+    buckets, bucket_names, numeric_values, text_values = create_buckets(train)
+
+    if percentile is not None:
+        i = ['0.5', '1', '2', '5', '10', '20', '50'].index(percentile)
+        bucket = buckets[i]
+        fn=f'{data_set}-hybrid-{bucket_names[i]}.hdf5'
+        with h5py.File(fn, 'w') as f:
+            f.attrs['type'] = 'dense'
+            f.attrs['distance'] = ds.attrs['distance']
+            f.attrs['dimension'] = len(test[0])
+            f.attrs['point_type'] = 'float'
+            f.attrs['bucket_names'] = bucket_names
+            f.attrs['selected_bucket'] = bucket_names[i]
+            for bucket_name in bucket_names:
+                grp = f.create_group(bucket_name)
+                grp["text"] = text_values[bucket_name]
+                grp["number"] = numeric_values[bucket_name]
+
+            f.create_dataset('train', train.shape, dtype=train.dtype)[:] = train
+            f.create_dataset('test', test.shape, dtype=test.dtype)[:] = test
+            # Write the id buckets so on ingestion we will know what data to assign for each id.
+
+            for j, id_bucket in enumerate(buckets):
+                np_bucket = np.array(id_bucket, dtype=np.int32)
+                f.create_dataset(f'{bucket_names[j]}_ids', np_bucket.shape, dtype=np_bucket.dtype)[:] = np_bucket
+
+            neighbors = f.create_dataset(f'neighbors', (len(test), count), dtype='i')
+            distances = f.create_dataset(f'distances', (len(test), count), dtype='f')
+
+            # Generate ground truth only for the relevan bucket.
+            train_bucket = np.array(bucket, dtype = np.int32)
+            train_set = np.empty((len(bucket), train.shape[1]), dtype=np.float32)
+            for id in range(len(bucket)):
+                train_set[id] = train[bucket[id]]
+            bf = BruteForceBLAS(distance, precision=train.dtype)
+            bf.fit(train_set)
+            Parallel(n_jobs=multiprocessing.cpu_count(),  require='sharedmem')(delayed(calc_i)(i, x, bf, test, neighbors, distances, count, train_bucket) for i, x in enumerate(test))
+
+    else:
+        for i, bucket in enumerate(buckets):
+            fn=f'{data_set}-hybrid-{bucket_names[i]}.hdf5'
+            with h5py.File(fn, 'w') as f:
+                f.attrs['type'] = 'dense'
+                f.attrs['distance'] = ds.attrs['distance']
+                f.attrs['dimension'] = len(test[0])
+                f.attrs['point_type'] = 'float'
+                f.attrs['bucket_names'] = bucket_names
+                f.attrs['selected_bucket'] = bucket_names[i]
+                for bucket_name in bucket_names:
+                    grp = f.create_group(bucket_name)
+                    grp["text"] = text_values[bucket_name]
+                    grp["number"] = numeric_values[bucket_name]
+
+                f.create_dataset('train', train.shape, dtype=train.dtype)[:] = train
+                f.create_dataset('test', test.shape, dtype=test.dtype)[:] = test
+                # Write the id buckets so on ingestion we will know what data to assign for each id.
+                for j, id_bucket in enumerate(buckets):
+                    np_bucket = np.array(id_bucket, dtype=np.int32)
+                    f.create_dataset(f'{bucket_names[j]}_ids', np_bucket.shape, dtype=np_bucket.dtype)[:] = np_bucket
+
+                neighbors = f.create_dataset(f'neighbors', (len(test), count), dtype='i')
+                distances = f.create_dataset(f'distances', (len(test), count), dtype='f')
+
+                # Generate ground truth only for the relevan bucket.
+                train_bucket = np.array(bucket, dtype = np.int32)
+                train_set = np.empty((len(bucket), train.shape[1]), dtype=np.float32)
+                for id in range(len(bucket)):
+                    train_set[id] = train[bucket[id]]
+                print(train_set.shape)
+                bf = BruteForceBLAS(distance, precision=train.dtype)
+                bf.fit(train_set)
+                Parallel(n_jobs=multiprocessing.cpu_count(),  require='sharedmem')(delayed(calc_i)(i, x, bf, test, neighbors, distances, count, train_bucket) for i, x in enumerate(test))
+                print(neighbors[1])
+                print(distances[1])
+
+
+if __name__ == "__main__":
+    create_ds()

redisbench_admin/run/ann/pkg/create_text_to_image_ds.py

@@ -0,0 +1,117 @@
+from ann_benchmarks.algorithms.bruteforce import BruteForceBLAS
+import struct
+import numpy as np
+import click
+import h5py
+from joblib import Parallel, delayed
+import multiprocessing
+
+def read_fbin(filename, start_idx=0, chunk_size=None):
+    """ Read *.fbin file that contains float32 vectors
+    Args:
+        :param filename (str): path to *.fbin file
+        :param start_idx (int): start reading vectors from this index
+        :param chunk_size (int): number of vectors to read. 
+                                 If None, read all vectors
+    Returns:
+        Array of float32 vectors (numpy.ndarray)
+    """
+    with open(filename, "rb") as f:
+        nvecs, dim = np.fromfile(f, count=2, dtype=np.int32)
+        nvecs = (nvecs - start_idx) if chunk_size is None else chunk_size
+        arr = np.fromfile(f, count=nvecs * dim, dtype=np.float32, 
+                          offset=start_idx * 4 * dim)
+    return arr.reshape(nvecs, dim)
+ 
+ 
+def read_ibin(filename, start_idx=0, chunk_size=None):
+    """ Read *.ibin file that contains int32 vectors
+    Args:
+        :param filename (str): path to *.ibin file
+        :param start_idx (int): start reading vectors from this index
+        :param chunk_size (int): number of vectors to read.
+                                 If None, read all vectors
+    Returns:
+        Array of int32 vectors (numpy.ndarray)
+    """
+    with open(filename, "rb") as f:
+        nvecs, dim = np.fromfile(f, count=2, dtype=np.int32)
+        nvecs = (nvecs - start_idx) if chunk_size is None else chunk_size
+        arr = np.fromfile(f, count=nvecs * dim, dtype=np.int32, 
+                          offset=start_idx * 4 * dim)
+    return arr.reshape(nvecs, dim)
+ 
+ 
+def write_fbin(filename, vecs):
+    """ Write an array of float32 vectors to *.fbin file
+    Args:s
+        :param filename (str): path to *.fbin file
+        :param vecs (numpy.ndarray): array of float32 vectors to write
+    """
+    assert len(vecs.shape) == 2, "Input array must have 2 dimensions"
+    with open(filename, "wb") as f:
+        nvecs, dim = vecs.shape
+        f.write(struct.pack('<i', nvecs))
+        f.write(struct.pack('<i', dim))
+        vecs.astype('float32').flatten().tofile(f)
+ 
+        
+def write_ibin(filename, vecs):
+    """ Write an array of int32 vectors to *.ibin file
+    Args:
+        :param filename (str): path to *.ibin file
+        :param vecs (numpy.ndarray): array of int32 vectors to write
+    """
+    assert len(vecs.shape) == 2, "Input array must have 2 dimensions"
+    with open(filename, "wb") as f:
+        nvecs, dim = vecs.shape
+        f.write(struct.pack('<i', nvecs))
+        f.write(struct.pack('<i', dim))
+        vecs.astype('int32').flatten().tofile(f)
+
+def calc_i(i, x, bf, test, neighbors, distances, count):
+    if i % 1000 == 0:
+        print('%d/%d...' % (i, len(test)))
+    res = list(bf.query_with_distances(x, count))
+    res.sort(key=lambda t: t[-1])
+    neighbors[i] = [j for j, _ in res]
+    distances[i] = [d for _, d in res]
+
+
+def calc(bf, test, neighbors, distances, count):
+    Parallel(n_jobs=multiprocessing.cpu_count(),  require='sharedmem')(delayed(calc_i)(i, x, bf, test, neighbors, distances, count) for i, x in enumerate(test))
+
+
+def write_output(train, test, fn, distance, point_type='float', count=100):
+    n = 0
+    f = h5py.File(fn, 'w')
+    f.attrs['type'] = 'dense'
+    f.attrs['distance'] = distance
+    f.attrs['dimension'] = len(train[0])
+    f.attrs['point_type'] = point_type
+    print('train size: %9d * %4d' % train.shape)
+    print('test size:  %9d * %4d' % test.shape)
+    f.create_dataset('train', (len(train), len(
+        train[0])), dtype=train.dtype)[:] = train
+    f.create_dataset('test', (len(test), len(
+        test[0])), dtype=test.dtype)[:] = test
+    neighbors = f.create_dataset('neighbors', (len(test), count), dtype='i')
+    distances = f.create_dataset('distances', (len(test), count), dtype='f')
+    bf = BruteForceBLAS(distance, precision=train.dtype)
+
+    bf.fit(train)
+    calc(bf, test, neighbors, distances, count)
+    f.close()
+
+@click.command()
+@click.option('--size', default=10, help='Number of vectors in milions.')
+@click.option('--distance', default='angular', help='distance metric.')
+@click.option('--test_set', required=True, type=str)
+@click.option('--train_set', required=True, type=str)
+def create_ds(size, distance, test_set, train_set):
+    test_set = read_fbin(test_set)
+    train_set= read_fbin(train_set, chunk_size=size*1000000)
+    write_output(train=train_set, test=test_set, fn=f'Text-to-Image-{size}M.hdf5', distance=distance, point_type='float', count=100)
+
+if __name__ == "__main__":
+    create_ds()