Nearest Neighbors algo implementation with FAISS

faiss_nn.py

import numpy as np
import faiss


class FaissNearestNeighbors:
    def __init__(self, k=5):
        self.index = None
        self.y = None
        self.k = k

    def fit(self, x, y=None):
        self.index = faiss.IndexFlatL2(x.shape[1])
        self.index.add(x.astype(np.float32))
        if y is not None:
            self.y = y

    def predict(self, x, k=None):
        distances, indices = self.index.search(
            x.astype(np.float32),
            k=self.k if k is None else k
        )
        if self.y is not None:
            labels = self.y[indices]
            return (distances, indices, labels)
        return (distances, indices)