[Medium] KNN (K-Nearest Neighbors) is Dead!

01 KNN is Dead.md

KNN is Dead

All the code snippets for KNN is Dead Medium Post Link

02_fit_hnsw_index.py

import hnswlib
import numpy as np

def fit_hnsw_index(features, ef=100, M=16, save_index_file=False):
    # Convenience function to create HNSW graph
    # features : list of lists containing the embeddings
    # ef, M: parameters to tune the HNSW algorithm
    
    num_elements = len(features)
    labels_index = np.arange(num_elements)
    EMBEDDING_SIZE = len(features[0])
    
    # Declaring index
    # possible space options are l2, cosine or ip
    p = hnswlib.Index(space='l2', dim=EMBEDDING_SIZE)
    
    # Initing index - the maximum number of elements should be known
    p.init_index(max_elements=num_elements, ef_construction=ef, M=M)
    
    # Element insertion
    int_labels = p.add_items(features, labels_index)
    
    # Controlling the recall by setting ef
    # ef should always be > k
    p.set_ef(ef) 
    
    # If you want to save the graph to a file
    if save_index_file:
         p.save_index(save_index_file)
    
    return p

03_querying_neighbors.py

ann_neighbor_indices, ann_distances = p.knn_query(features, k)

04_imports.py

# Imports
# For input data pre-processing
import json
import gzip
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import fasttext.util
fasttext.util.download_model('en', if_exists='ignore') # English pre-trained model
ft = fasttext.load_model('cc.en.300.bin')

# For KNN vs ANN benchmarking
from datetime import datetime
from tqdm import tqdm
from sklearn.neighbors import NearestNeighbors
import hnswlib

05_data.py

# Data: http://deepyeti.ucsd.edu/jianmo/amazon/
data = []
with gzip.open('meta_Cell_Phones_and_Accessories.json.gz') as f:
    for l in f:
        data.append(json.loads(l.strip()))
        
# Pre-Processing: https://colab.research.google.com/drive/1Zv6MARGQcrBbLHyjPVVMZVnRWsRnVMpV#scrollTo=LgWrDtZ94w89
# Convert list into pandas dataframe
df = pd.DataFrame.from_dict(data)
df.fillna('', inplace=True)

# Filter unformatted rows
df = df[~df.title.str.contains('getTime')] 

# Restrict to just 'Cell Phones and Accessories'
df = df[df['main_cat']=='Cell Phones & Accessories']

# Reset index
df.reset_index(inplace=True, drop=True)

# Only keep the title columns
df = df[['title']]

# Check the df
print(df.shape)
df.head()

06_embedding.py

# Title Embedding using FastText Sentence Embedding
df['emb'] = df['title'].apply(ft.get_sentence_vector)

# Extract out the embeddings column as a list of lists for input to our NN algos
X = [item.tolist() for item in df['emb'].values]

07_benchmarking.py

# Number of products for benchmark loop
n_products = [1000, 10000, 100000, len(X)]

# Number of neighbors for benchmark loop
n_neighbors = [10, 100]

# Dictionary to save metric results for each iteration
metrics = {'products':[], 'k':[], 'knn_time':[], 'ann_time':[], 'pct_overlap':[]}

for products in tqdm(n_products):
    # "products" number of products included in the search space
    features = X[:products]
    
    for k in tqdm(n_neighbors):   
        # "K" Nearest Neighbor search
        # KNN 
        knn_start = datetime.now()
        nbrs = NearestNeighbors(n_neighbors=k, metric='euclidean').fit(features)
        knn_distances, knn_neighbor_indices = nbrs.kneighbors(features)
        knn_end = datetime.now()
        metrics['knn_time'].append((knn_end - knn_start).total_seconds())
        
        # HNSW ANN
        ann_start = datetime.now()
        p = fit_hnsw_index(features, ef=k*10)
        ann_neighbor_indices, ann_distances = p.knn_query(features, k)
        ann_end = datetime.now()
        metrics['ann_time'].append((ann_end - ann_start).total_seconds())
        
        # Average Percent Overlap in Nearest Neighbors across all "products"
        pct_overlap_per_product = [len(np.intersect1d(knn_neighbor_indices[i], ann_neighbor_indices[i]))/k for i in range(len(features))]
        metrics['pct_overlap'].append(np.mean(pct_overlap_per_product))
        
        metrics['products'].append(products)
        metrics['k'].append(k)
        
metrics_df = pd.DataFrame(metrics)
metrics_df.to_csv('data/metrics_df.csv', index=False)
metrics_df