generate_training_data_multi.py

# model input: <K, 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 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


all_results = []


def parse_count_list(s):
    """Parse a string of comma-separated integers into a list of integers"""
    return [int(x.strip()) for x in s.split(',')]


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--dataset',
        metavar="DATASET",
        required=True)
    parser.add_argument(
        '--counts',
        nargs='*',
        required=True,
        help='counts to use')
    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(
        '--private-query',
        help='Use the private queries and ground truth',
        action='store_true')
    parser.add_argument(
        '--training_data_path',
        type=str,
        required=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'), "generate_training_data.py only supports train mode"

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

    if not args.counts:
        args.counts = [dataset.default_count()]

    counts = [int(c) for c in args.counts]
    max_count = max(counts)
    print(f"Processing counts: {counts}, max count: {max_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)
    
    # Calculate total rows needed for combined dataset
    total_rows = 0
    for count in counts:
        all_results.append(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))
        for i, (fileroot, filename, properties, run) in enumerate(all_results[-1]):
            if not filename.startswith(args.mode):
                continue
                
            for j in range(0, properties['num_searches']):
                step_suffix = str(properties['step_' + str(j)])
                total_cmps = np.array(run['total_cmps_step' + step_suffix])[:nq]
                
                if args.mode.startswith('train_opt') or args.mode.startswith('train_darth'):
                    for total_cmps_per_query in total_cmps:
                        total_rows += total_cmps_per_query - 1
                else:
                    total_rows += nq
            # We only need to check the first result file to get the structure
            break

    print(f"Total rows across all counts: {total_rows}")

    # Initialize combined data arrays
    if args.mode.startswith('train_opt'):
        cols = 16  # Added K value to the input
    elif args.mode.startswith('train_darth'):
        cols = 13  # Added K value to the input
    else:
        cols = dim + 7  # Added K value to the input (was dim + 6)

    combined_input_data = np.zeros((total_rows, cols), dtype=np.float32)
    combined_output_data = np.zeros((total_rows, 1), dtype=np.float32)
    gt_cmps = np.full((nq, max_count), np.inf, dtype=np.float32)  # Initialize with inf
    curr_row = 0

    # Calculate total run latency and generate latency across all counts
    total_run_latency = 0
    overall_start_time = time.time()  # Start time for generate latency calculation

    # Process each count value and accumulate data
    for count in counts:
        print(f"Processing count: {count}")
        all_results.append(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))
        for i, (fileroot, filename, properties, run) in enumerate(all_results[-1]):
            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)

            nums_active_nodes = [0]
            for step, entry in enumerate(runbook):
                if entry['operation'] == 'search':
                    nums_active_nodes.append(nums_active_nodes[-1])
                elif entry['operation'] == 'insert':
                    nums_active_nodes.append(nums_active_nodes[-1] + (entry['end'] - entry['start']))
                elif entry['operation'] == 'delete':
                    nums_active_nodes.append(nums_active_nodes[-1] - (entry['end'] - entry['start']))
                else:
                    raise Exception(f'Undefined runbook operation {entry["operation"]}')

            for i in range(0, properties['num_searches']):
                start_time = time.time()

                search_step_id = properties['step_' + str(i)]
                step_suffix = str(search_step_id)
                N = nums_active_nodes[search_step_id]
                step_latency = properties['latency_step_' + step_suffix] / 10**6
                print(f"step {search_step_id}, N = {N}, latency = {step_latency} s")
                
                # Accumulate run latency
                total_run_latency += step_latency

                neighbors = np.array(run['neighbors_step' +  step_suffix])[:nq]
                total_cmps = np.array(run['total_cmps_step' + step_suffix])[:nq]
                total_latency = np.array(run['total_latency_step' + step_suffix])[:nq]
                cmps = np.array(run['cmps_step' + step_suffix])[:nq]
                lats = np.array(run['lats_step' + step_suffix])[:nq]

                dists_start = np.array(run['dists_start_step' + step_suffix])[:nq]
                dists_1st = np.array(run['dists_1st_step' + step_suffix])[:nq]

                if args.mode.startswith("train_opt"): # train_opt
                    dist_1st_hops = np.array(run['dist_1st_hops_step' + step_suffix])[:nq]
                    dist_1st_cmps = np.array(run['dist_1st_cmps_step' + step_suffix])[:nq]
                elif args.mode.startswith("train_darth"): # train_darth, train_darth_opt
                    dists_kth = np.array(run['dists_kth_step' + step_suffix])[:nq]
                else: # train
                    dists_10th = np.array(run['dists_10th_step' + step_suffix])[:nq]

                if args.mode.startswith("train_opt"): # train_opt
                    dists_visited = np.array(run['dists_visited_step' + step_suffix])[:nq]
                    cmps_visited = np.array(run['cmps_visited_step' + step_suffix])[:nq]
                    hops_visited = np.array(run['hops_visited_step' + step_suffix])[:nq]
                elif args.mode.startswith("train_darth"): # train_darth, train_darth_opt
                    cmps_visited = np.array(run['cmps_visited_step' + step_suffix])[:nq]
                    hops_visited = np.array(run['hops_visited_step' + step_suffix])[:nq]
                    inserts_visited = np.array(run['inserts_visited_step' + step_suffix])[:nq]

                if args.mode.startswith('train_opt'):
                    traversal_window_stats = np.array(run['traversal_window_stats_step' + step_suffix])[:nq]

                if args.mode.startswith('train_darth'):
                    result_set_stats = np.array(run['result_set_stats_step' + step_suffix])[:nq]

                groundtruths, groundtruth_distances = read_gt_fromdir(gt_dir, step_suffix, count, train=True)

                mean_recall = 0
                
                # allocate input_data and output_data for this count/step
                if args.mode.startswith('train_opt') or args.mode.startswith('train_darth'):
                    rows = 0
                    for total_cmps_per_query in total_cmps:
                        rows += total_cmps_per_query - 1
                else:
                    rows = nq

                input_data = np.zeros((rows, cols), dtype=np.float32)
                output_data = np.zeros((rows, 1), dtype=np.float32)
                curr_row_in_this_batch = 0

                for (query_id, (
                    neighbors_per_query, groundtruths_per_query, groundtruth_distances_per_query, 
                    cmps_per_query, lats_per_query, 
                    total_cmps_per_query, total_latency_per_query,
                    dist_start
                )) in tqdm(enumerate(zip(
                    neighbors, groundtruths, groundtruth_distances,
                    cmps, lats, 
                    total_cmps, total_latency,
                    dists_start
                )), total=nq, desc="Generating training data"):

                    recall_per_query = 0
                    true_cmps_per_query = []
                    true_cmps_per_query_sorted = []
                    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
                            true_cmps_per_query.append(cmps_per_query[rank_in_neighbors])
                            true_cmps_per_query_sorted.append(cmps_per_query[rank_in_neighbors])
                        else:
                            true_cmps_per_query.append(total_cmps_per_query)
                    mean_recall += recall_per_query
                    true_cmps_per_query_sorted.sort()

                    assert(len(true_cmps_per_query) == count)
                    np.copyto(gt_cmps[query_id, :count], np.array(true_cmps_per_query))

                    if args.mode.startswith('train_opt'):
                        assert(len(hops_visited[query_id]) == total_cmps_per_query - 1)

                        # model input: <K, curr_hops, curr_cmps, curr_dist, dist_1st_hops, dist_1st_cmps, dist_1st, dist_start> + traversal_window_stats
                        # model output: probability that Top1 has been collected
                        input_data_rows = np.concatenate((
                            np.ones((total_cmps_per_query - 1, 1)) * query_id,
                            np.ones((total_cmps_per_query - 1, 1)) * count,  # K value
                            np.array(hops_visited[query_id]).reshape(-1, 1),
                            np.array(cmps_visited[query_id]).reshape(-1, 1),
                            np.array(dists_visited[query_id]).reshape(-1, 1),
                            np.array(dist_1st_hops[query_id]).reshape(-1, 1),
                            np.array(dist_1st_cmps[query_id]).reshape(-1, 1),
                            np.array(dists_1st[query_id]).reshape(-1, 1),
                            np.ones((total_cmps_per_query - 1, 1)) * dist_start,
                            np.array(traversal_window_stats[query_id]).reshape(-1, 7),
                        ), axis=1)

                        output_rows = np.array(cmps_visited[query_id]) >= max(true_cmps_per_query)
                        output_rows = output_rows.reshape(-1, 1)

                        # copy to input_data and output_data
                        np.copyto(input_data[curr_row_in_this_batch : curr_row_in_this_batch + total_cmps_per_query - 1], input_data_rows)
                        np.copyto(output_data[curr_row_in_this_batch : curr_row_in_this_batch + total_cmps_per_query - 1], output_rows)

                        curr_row_in_this_batch += total_cmps_per_query - 1

                    elif args.mode.startswith('train_darth'):
                        assert(len(hops_visited[query_id]) == total_cmps_per_query - 1)

                        # model input: <K, curr_hops, curr_cmps, curr_inserts, dist_start, dist_1st, dist_kth> + result_set_stats
                        # model output: curr recall
                        input_data_rows = np.concatenate((
                            np.ones((total_cmps_per_query - 1, 1)) * query_id,
                            np.ones((total_cmps_per_query - 1, 1)) * count,  # K value
                            np.array(hops_visited[query_id]).reshape(-1, 1),
                            np.array(cmps_visited[query_id]).reshape(-1, 1),
                            np.array(inserts_visited[query_id]).reshape(-1, 1),
                            np.ones((total_cmps_per_query - 1, 1)) * dist_start,
                            np.array(dists_1st[query_id]).reshape(-1, 1),
                            np.array(dists_kth[query_id]).reshape(-1, 1),
                            np.array(result_set_stats[query_id]).reshape(-1, 5),
                        ), axis=1)

                        recalls = np.searchsorted(np.array(true_cmps_per_query_sorted), np.array(cmps_visited[query_id]), side='right')
                        output_rows = recalls / count
                        output_rows = output_rows.reshape(-1, 1)

                        # copy to input_data and output_data
                        np.copyto(input_data[curr_row_in_this_batch : curr_row_in_this_batch + total_cmps_per_query - 1], input_data_rows)
                        np.copyto(output_data[curr_row_in_this_batch : curr_row_in_this_batch + total_cmps_per_query - 1], output_rows)

                        curr_row_in_this_batch += total_cmps_per_query - 1

                    else:
                        # model input: K + query + <dist_start, dist_1st, dist_10th, 1st_to_start, 10th_to_start>
                        # model output: log2(predicted steps)
                        input_data_row = [query_id, count] + Q[query_id].tolist() + [  # Added K value at the beginning
                            dist_start, dists_1st[query_id], dists_10th[query_id],
                            dists_1st[query_id] / dist_start, dists_10th[query_id] / dist_start,
                        ]

                        if recall_per_query > 0:
                            output_data_row = [np.log2(true_cmps_per_query_sorted[-1])]
                        else:
                            output_data_row = [np.log2(total_cmps_per_query)]

                        # copy to input_data and output_data
                        np.copyto(input_data[curr_row_in_this_batch], np.array(input_data_row))
                        np.copyto(output_data[curr_row_in_this_batch], np.array(output_data_row))

                        curr_row_in_this_batch += 1

                mean_recall = mean_recall / (count * nq)
                print(f"recall: {mean_recall:.4f}")

                # Copy data to combined arrays
                np.copyto(combined_input_data[curr_row : curr_row + rows], input_data[:rows])
                np.copyto(combined_output_data[curr_row : curr_row + rows], output_data[:rows])
                curr_row += rows

                print(f"input shape: {input_data.shape}")
                print(f"output shape: {output_data.shape}")

                generate_latency = time.time() - start_time
                print(f"run_latency: {step_latency} s, generate_latency: {generate_latency} s, total: {step_latency + generate_latency} s")

            print("")

    # Save combined data
    training_data_path = args.training_data_path
    if not os.path.exists(training_data_path):
        os.makedirs(training_data_path)

    combined_input_data_path = os.path.join(training_data_path, f"input_data_step{step_suffix}.npy")
    combined_output_data_path = os.path.join(training_data_path, f"output_data_step{step_suffix}.npy")
    gt_cmps_path = os.path.join(training_data_path, f"gt_cmps_step{step_suffix}.npy")
    combined_latency_path = os.path.join(training_data_path, f"latency_step{step_suffix}.json")

    print(f"Combined input shape: {combined_input_data.shape}")
    print(f"Combined output shape: {combined_output_data.shape}")
    print(f"gt_cmps shape: {gt_cmps.shape}")

    print(f"Combined input examples: {combined_input_data[:10]}")
    print(f"Combined output examples: {combined_output_data[:10]}")
    print(f"gt_cmps examples: {gt_cmps[:10]}")

    np.save(combined_input_data_path, combined_input_data)
    np.save(combined_output_data_path, combined_output_data)
    np.save(gt_cmps_path, gt_cmps)
    print(f"Saved combined training data to {combined_input_data_path} and {combined_output_data_path}")
    print(f"Saved gt_cmps to {gt_cmps_path}")

    # Close all open HDF5 files to prevent warnings at program exit
    for results in all_results:
        if hasattr(results, '__iter__'):
            # For generators, we need to iterate through them to close files
            try:
                for _, _, _, f in results:
                    if hasattr(f, 'close'):
                        f.close()
            except:
                # Generator might be exhausted, that's okay
                pass

    total_generate_latency = time.time() - overall_start_time
    with open(combined_latency_path, "w") as f:
        json.dump({
            "run_latency": total_run_latency,
            "generate_latency": total_generate_latency,
        }, f, indent=4)