calculate_ann_metrics.py

import asyncio
import heapq
from typing import List, Tuple

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
import torch
from tqdm import tqdm

import opensearch_indexer
from corpus import WikipediaCorpus
from corpus_embed import EMBEDDING_MODEL_NAME
from embeddings import EmbeddingModel, batch_embeddings_loader


def create_opensearch_knn_query_clause(knn_field: str, vector: List[any], k: int):
    return {
        "knn": {
            knn_field: {
                "vector": vector,
                "k": k
            }
        }
    }


def calculate_exact_nn(corpus_embeddings_files_pattern: str, query_embeddings: torch.Tensor, comparison_batch_size: int,
                       max_elements: int) -> List[List[Tuple[int, float]]]:
    """batch-wise calculation of exact NN for the given query embeddings and corpus embeddings loaded from the file system"""

    def cos_sim(a: torch.Tensor, b: torch.Tensor):
        # assumes tensors are normalized already
        return torch.mm(a, b.transpose(0, 1))

    global_top_similar = [[] for _ in range(len(query_embeddings))]

    progress_bar = tqdm(unit=' docs compared')
    for batch_start_idx, batch_tensor in batch_embeddings_loader(corpus_embeddings_files_pattern, comparison_batch_size, max_elements):
        similarities = cos_sim(query_embeddings, batch_tensor)

        # get the top 10 similarities and adjust for batch_start_idx
        top_k_values, top_k_indices = torch.topk(similarities, k=min(len(similarities[0]), 10), dim=1, largest=True, sorted=False)
        top_k_values = top_k_values.cpu().tolist()
        top_k_indices = top_k_indices.cpu().tolist()

        for query_idx in range(len(query_embeddings)):
            for idx, score in zip(top_k_indices[query_idx], top_k_values[query_idx]):
                global_idx = batch_start_idx + idx
                if len(global_top_similar[query_idx]) < 10:
                    heapq.heappush(global_top_similar[query_idx], (score, global_idx))
                else:
                    heapq.heappushpop(global_top_similar[query_idx], (score, global_idx))

        progress_bar.update(len(batch_tensor))

    return [sorted([(idx, score) for score, idx in query_result], reverse=True, key=lambda x: x[1])
            for query_result in global_top_similar]


async def search_approximate_nn(query_embeddings: torch.Tensor, opensearch_host: str, opensearch_port: int, index_name: str,
                                knn_field: str, k: int, size: int=100) -> List[Tuple[int, List[float], List[str], List[str]]]:
    os = opensearch_indexer.OpenSearchIndexer(opensearch_host, opensearch_port, index_name)
    queries = [{
        'query': create_opensearch_knn_query_clause(knn_field, embedding.tolist(), k),
        'size': size
    } for embedding in query_embeddings]
    results = await os.search_all_queries(queries)
    await os.close()
    return results


def calculate_recall(exact_ids: List[str], approximate_ids: List[any], k: int):
    exact_ids = exact_ids[:k]
    approximate_ids = approximate_ids[:k]
    return float(len([id for id in approximate_ids if id in exact_ids])) / len(exact_ids)

def load_queries():
    return pd.read_csv(QUERIES_FILE)["question"].tolist()


def retrieve_ann_results(corpus_embeddings_file_pattern: str, query_embeddings: torch.Tensor, exact_distance_batch_size: int,
                         max_items: int, opensearch_host: str, opensearch_port: int, opensearch_index: str,
                         k_values: list[int], num_runs: int):
    combined_results = []
    exact_nn = calculate_exact_nn(corpus_embeddings_file_pattern, query_embeddings, exact_distance_batch_size , max_items)
    for run in range(num_runs):
        for k in k_values:
            print(f"Run {run}, index={opensearch_index}, max_items={max_items}, k={k}")
            run_results = asyncio.run(
                search_approximate_nn(query_embeddings, opensearch_host, opensearch_port, opensearch_index, "text_knn_lucene", k=k))
            # drop/ignore run 0 as a warm-up run
            if run > 0:
                for idx, (approx_result, exact_result) in enumerate(zip(run_results, exact_nn)):
                    exact_ids, exact_scores = map(list, zip(*exact_result))
                    combined_results.append({
                        'run': run,
                        'index': opensearch_index,
                        'max_items': max_items,
                        'k': k,
                        'query_idx': idx,
                        'time_taken': approx_result[0],
                        'approx_ids': approx_result[1],
                        'approx_scores': approx_result[2],
                        'exact_ids': [str(id) for id in exact_ids],
                        'exact_scores': exact_scores
                    })
    return pd.DataFrame(combined_results)

def plot_recall_vs_time_taken(df):
    plt.figure(figsize=(12, 6))
    recall_positions = [1, 3, 5, 10]
    df_long = pd.melt(df, id_vars=['k', 'time_taken'], value_vars=[f'recall@{i}' for i in recall_positions], var_name='recall', value_name='value')
    sns.lineplot(data=df_long, x='value', y='time_taken', hue='recall', marker='o')
    sns.scatterplot(data=df_long, x='value', y='time_taken', hue='recall', style='recall', legend=False, s=100)

    for i in range(len(df)):
        x = np.mean([k_df[f"recall@{x}"][i] for x in recall_positions])
        y = k_df['time_taken'][i]
        plt.text(x, y, f"k={k_df['k'][i]}", fontsize=9, ha='right', va='bottom')

    plt.xlabel('Recall')
    plt.ylabel('Time (ms)')
    plt.title('Recall vs. Query Time for Values of k')

    plt.grid(visible=True)

    plt.savefig('recall_vs_time_plot_seaborn.png')

    plt.show()


# the files containing the numpy arrays that represent the embeddings
EMBEDDINGS_FILE_PATTERN = 'embeddings/wikipedia_embeddings_{i}.npy'

# the file containing the queries
QUERIES_FILE = 'queries.csv'

# for exact NN we always compare the query embeddings of all questions against this BATCH_SIZE of corpus items
EXACT_NN_BATCH_SIZE = 512

# OpenSearch coordinates for ANN retrieval
OPENSEARCH_HOST = "localhost"
OPENSEARCH_PORT = 9200

# first N items / embeddings from the corpus to consider for exact retrieval
CORPUS_MAX_ITEMS = WikipediaCorpus().num_entries()
# CORPUS_MAX_ITEMS = WikipediaCorpus().num_entries() / 10
# CORPUS_MAX_ITEMS = WikipediaCorpus().num_entries() / 100

# the OpenSearch index for approximate retrieval, needs to match the corpus
OPENSEARCH_INDEX = "wikipedia_en_100pct_2shards"
# OPENSEARCH_INDEX = "wikipedia_en_10pct"
# OPENSEARCH_INDEX = "wikipedia_en_1pct"

# the k values to iterate
K_VALUES = [10, 20, 40, 80, 160, 240, 320, 480, 640]

# the runs to perform, we take the mean of multiple runs to limit the influence of outliers
RUNS = 5

OUTPUT_FILE = f"{OPENSEARCH_INDEX}_recall.csv"


if __name__ == '__main__':
    queries = load_queries()
    embedding_model = EmbeddingModel(EMBEDDING_MODEL_NAME)
    query_embeddings = embedding_model.tokenize_and_embed(queries)
    all_ann_results = retrieve_ann_results(corpus_embeddings_file_pattern=EMBEDDINGS_FILE_PATTERN,
                                           query_embeddings=query_embeddings,
                                           exact_distance_batch_size=EXACT_NN_BATCH_SIZE,
                                           max_items=CORPUS_MAX_ITEMS,
                                           opensearch_host=OPENSEARCH_HOST,
                                           opensearch_port=OPENSEARCH_PORT,
                                           opensearch_index=OPENSEARCH_INDEX,
                                           k_values=K_VALUES,
                                           num_runs=RUNS)
    # add recall@N columns
    for pos in range(1, 11):
        all_ann_results[f"recall@{pos}"] = all_ann_results.apply(lambda row: calculate_recall(row['exact_ids'], row['approx_ids'], k=pos), axis=1)

    aggregations = {col: 'mean'
                    for col in all_ann_results.columns
                    if col == 'time_taken' or col.startswith('recall@')}

    k_query_df = all_ann_results.groupby(['index', 'k', 'query_idx'], as_index=False).agg(aggregations)
    k_df = k_query_df.groupby(['index', 'k'], as_index=False).agg(aggregations)
    plot_recall_vs_time_taken(k_df)

    k_df.to_csv(OUTPUT_FILE, index=False)