<jupyter_start><jupyter_text>Annoy is an in-memory cpu-only approximate nearest neighbor index which I'm using in this demo. We will be running some performance experiments by grid searching over two Annoy parameters, `n_trees` and `search_k` assuming the dataset size is constant and assuming a fixed latency bugdet (off the top of my head 1-2ms should be reasonable).The parameters are explained in more detail in https://github.com/spotify/annoytradeoffs.This notebook assumes you have a prepared encoder and a directory of wav files.<jupyter_code>%load_ext autoreload
%autoreload 2

!pip3 install -q pandas matplotlib
import warnings; warnings.filterwarnings("ignore"); import os; os.environ['KALDI_ROOT'] = '/tmp' # ignore warnings
from pathlib import Path
import time

import matplotlib.pyplot as plt
import pandas as pd
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader, ConcatDataset
from tqdm import tqdm

import boiler.index
import boiler.dataset
class args:
    wav_dir = Path('/home/proger/coub-crawler/monthlyLog/wav')
    encoder = Path('/home/proger/boiler/exp/p_t64_b512/vqvae_223/BagTopVQVAE/encoder.pt')
    index_dir = Path('/home/proger/boiler/annoy')
    device = 'cuda'
coubs = sorted(args.wav_dir.glob('*.wav'))
coub_dataset = ConcatDataset([boiler.dataset.WavFile(wav) for wav in coubs])
coub_loader = DataLoader(coub_dataset, batch_size=64, shuffle=False, num_workers=4)
model = torch.jit.load(str(args.encoder))
embeddings = []
for batch in tqdm(coub_loader):
    embeddings.append(model(batch.to(args.device)).detach().cpu())
embeddings = torch.cat(embeddings)<jupyter_output>100%|██████████| 146/146 [00:05<00:00, 26.06it/s]<jupyter_text>Tuning guidelines:> Basically it's recommended to set `n_trees` as large as possible given the amount of memory you can afford,> and it's recommended to set `search_k` as large as possible given the time constraints you have for the queries.`search_k` defaults to `n*n_trees`.<jupyter_code>experiments = []

n = 10
query = embeddings[1]  # maybe try more queries?
baseline_result = (1 - F.cosine_similarity(query.unsqueeze(0), embeddings)).argsort(descending=False)[:n].tolist()

for n_trees in tqdm([5, 10, 20, 50, 100, 200, 300, 500, 1000]):
    index, path = boiler.index.make_index(embeddings, args.index_dir, n_trees)
    size = path.stat().st_size
    
    for search_k in [n*n_trees//2, n*n_trees, 2*n*n_trees, 3*n*n_trees, 10*n*n_trees]:
        def run():
            a = time.perf_counter()
            result = index.get_nns_by_vector(query, n=n, search_k=search_k, include_distances=False)
            b = time.perf_counter()
            recall = len(set(baseline_result).intersection(set(result)))/len(baseline_result)
            return b-a, recall
        
        stats = torch.tensor([run() for _ in range(1000)])
        latency, recall = stats.mean(dim=0).tolist()
        latency_std, recall_std = stats.std(dim=0).tolist()
        
        experiments.append(dict(n_trees=n_trees,
                                 search_k=search_k,
                                 size=size,
                                 latency=latency, latency_std=latency_std,
                                 recall=recall, recall_std=recall_std))
        
    del index
pd.set_option('display.max_rows', 500)
df = pd.DataFrame(experiments)
df<jupyter_output><empty_output><jupyter_text>I will be dropping all runs that have even a slightest drop in recall:<jupyter_code>df = df[~(df.recall < 1)]<jupyter_output><empty_output><jupyter_text>Annoy mmaps a single index only once into the memory and it can be shared between multiple workers. We only need to make sure a single copy fits into our RAM budget assuming the dataset size is constant.<jupyter_code>df.plot.scatter(x='n_trees', y='size')
plt.axhline(0.5*1e9, color='red', label='ram cap (0.5 gigabytes)')<jupyter_output><empty_output><jupyter_text>Let's look at the latency.<jupyter_code>df.latency.plot()
plt.axhline(0.002, color='red', label='hard budget cap (2ms)')
plt.axhline(0.001, color='orange', label='budget cap (1ms)')
plt.legend()

df[(df.latency<0.002)]<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11536608863198458, 'fraction_numerical': 0.02384393063583815, 'mean_word_length': 4.16542288557214, 'pattern_counts': {'":': 0, '<': 23, '<?xml version=': 0, '>': 22, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '552608', 'license': 'MIT', 'n_tokens_mistral': 1472, 'n_tokens_neox': 1361, 'n_words': 358, 'path': 'benchmark-annoy.ipynb', 'repo_name': 'proger/boiler', 'size': '4152'}