Spaces:
Running
Running
| import faiss | |
| import numpy as np | |
| class FaissNeighbors: | |
| def __init__(self): | |
| self.index = None | |
| self.y = None | |
| def fit(self, X, y): | |
| self.index = faiss.IndexFlatL2(X.shape[1]) | |
| self.index.add(X.astype(np.float32)) | |
| self.y = y | |
| def get_distances_and_indices(self, X, top_K=1000): | |
| distances, indices = self.index.search(X.astype(np.float32), k=top_K) | |
| return np.copy(distances), np.copy(indices), np.copy(self.y[indices]) | |
| def get_nearest_labels(self, X, top_K=1000): | |
| distances, indices = self.index.search(X.astype(np.float32), k=top_K) | |
| return np.copy(self.y[indices]) | |
| class FaissCosineNeighbors: | |
| def __init__(self): | |
| self.cindex = None | |
| self.y = None | |
| def fit(self, X, y): | |
| self.cindex = faiss.index_factory(X.shape[1], "Flat", faiss.METRIC_INNER_PRODUCT) | |
| X = np.copy(X) | |
| X = X.astype(np.float32) | |
| faiss.normalize_L2(X) | |
| self.cindex.add(X) | |
| self.y = y | |
| def get_distances_and_indices(self, Q, topK): | |
| Q = np.copy(Q) | |
| faiss.normalize_L2(Q) | |
| distances, indices = self.cindex.search(Q.astype(np.float32), k=topK) | |
| return np.copy(distances), np.copy(indices), np.copy(self.y[indices]) | |
| def get_nearest_labels(self, Q, topK=1000): | |
| Q = np.copy(Q) | |
| faiss.normalize_L2(Q) | |
| distances, indices = self.cindex.search(Q.astype(np.float32), k=topK) | |
| return np.copy(self.y[indices]) |