train_rem.py

# model input: <query, dist_start, dist_1st, dist_10th, 1st_to_start, 10th_to_start>
# model output: log2(predicted steps)


import matplotlib as mpl
mpl.use('Agg')  # noqa
import matplotlib.pyplot as plt
import numpy as np
import argparse
import os
import pickle
from tqdm import tqdm
import time
import json
from scipy import stats


from benchmark.datasets import DATASETS
from benchmark.algorithms.definitions import get_definitions
from benchmark.plotting.metrics import all_metrics as metrics
from benchmark.plotting.metrics import get_all_recall_values, get_count_at_certain_recall
from benchmark.plotting.utils import (get_plot_label, compute_metrics,
        create_linestyles, create_pointset)
from benchmark.results import (store_results, load_all_results, load_all_results_without_read, 
        get_result_filename, get_unique_algorithms)
from benchmark.dataset_io import knn_result_read
import benchmark.streaming.compute_gt
from benchmark.streaming.load_runbook import load_runbook
from benchmark.utils import read_gt_fromdir

from meta_analysis import read_float_arg_from_filename, read_int_arg_from_filename


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--dataset',
        metavar="DATASET",
        required=True)
    parser.add_argument(
        '--count',
        default=-1,
        type=int)
    parser.add_argument(
        '--definitions',
        metavar='FILE',
        help='load algorithm definitions from FILE',
        default='algos-2021.yaml')
    parser.add_argument(
        '--neurips23track',
        choices=['filter', 'ood', 'sparse', 'streaming', 'none'],
        default='none'
    )
    parser.add_argument(
        '--runbook_path',
        metavar='FILE',
        help='paths to runbooks',
    )
    parser.add_argument(
        '--results_base_path',
        type=str,
        default='results'
    )
    parser.add_argument(
        '--model_base_path',
        type=str,
        required=True
    )
    parser.add_argument(
        '--private-query',
        help='Use the private queries and ground truth',
        action='store_true')
    parser.add_argument(
        '--mode',
        type=str,
        required=True)
    parser.add_argument(
        '--filtered',
        help='Use filtered queries.',
        action='store_true'
    )
    parser.add_argument(
        '--label_file', 
        type=str, 
        default=None,
        help='Path to the label file.'
    )
    parser.add_argument(
        '--filter_label_file', 
        type=str, 
        default=None,
        help='Path to the filter file.'
    )
    parser.add_argument(
        '--num_queries',
        type=int,
        default=0,
        help='Number of queries to run.'
    )
    args = parser.parse_args()

    assert args.mode.startswith("train_rem"), "train_rem.py only supports train_rem mode"

    dataset = DATASETS[args.dataset]()
    dim = dataset.d

    if args.count == -1:
        args.count = dataset.default_count()

    count = int(args.count)

    Q = dataset.get_training_queries().astype(np.float32)
    if args.num_queries > 0:
        Q = Q[:args.num_queries]
    nq = Q.shape[0]
    print(fr"Got {nq} queries")

    max_pts, runbook = load_runbook(args.dataset, dataset.nb, args.runbook_path)
    results = load_all_results(args.dataset, count, neurips23track=args.neurips23track, runbook_path=args.runbook_path, \
                               filtered=args.filtered, label_file=args.label_file, filter_label_file=args.filter_label_file, base_path=args.results_base_path)
    rem_table = {}
    step_latency = 0
    start_time = time.time()

    for i, (fileroot, filename, properties, run) in enumerate(results):
        if not filename.startswith(args.mode):
            continue

        print(f"from {fileroot}/{filename}")
        if fileroot.split("/")[-1].endswith('.txt'):
            label_file, filter_label_file = fileroot.split("/")[-2], fileroot.split("/")[-1]
        else:
            label_file, filter_label_file = "", ""
        gt_dir = benchmark.streaming.compute_gt.gt_dir(dataset, args.runbook_path, label_file, filter_label_file)

        for i in range(0, properties['num_searches']):
            search_step_id = properties['step_' + str(i)]
            step_suffix = str(search_step_id)
            step_latency += properties['latency_step_' + step_suffix] / 10**6
            neighbors = np.array(run['neighbors_step' + step_suffix])
            groundtruths, groundtruth_distances = read_gt_fromdir(gt_dir, step_suffix, count, train=True)

            mean_recall = 0
            for (query_id, (
                neighbors_per_query, groundtruths_per_query, 
            )) in tqdm(enumerate(zip(
                neighbors, groundtruths, 
            )), total=nq, desc="Generating training data"):

                recall_per_query = 0
                for rank, groundtruth in enumerate(groundtruths_per_query):
                    if groundtruth in neighbors_per_query:
                        rank_in_neighbors = np.where(neighbors_per_query == groundtruth)[0][0]
                        recall_per_query += 1
                mean_recall += recall_per_query

            mean_recall = mean_recall / (count * nq)
            ef = read_int_arg_from_filename(filename, "efSearch", 0)
            rem_table[mean_recall] = ef

    train_latency = time.time() - start_time

    # Dump ef_recall_map
    if not os.path.exists(args.model_base_path):
        os.makedirs(args.model_base_path)

    rem_table_path = os.path.join(args.model_base_path, "rem_table.txt")
    with open(rem_table_path, "w") as f:
        # sort rem_table by recall
        rem_table = {k: v for k, v in sorted(rem_table.items(), key=lambda item: item[0])}
        for recall, ef in rem_table.items():
            f.write(f"{recall},{ef}\n")

    latency_path = os.path.join(args.model_base_path, "latency.json")
    with open(latency_path, "w") as f:
        json.dump({
            "run_latency": step_latency,
            "generate_latency": 0,
            "train_latency": train_latency,
            "mse": 0,
        }, f, indent=4)