| 18 | |
| 19 | |
| 20 | class FaissSearch(): |
| 21 | def __init__(self, data_vector=None, sport_mode=True): |
| 22 | self.data_vector = data_vector |
| 23 | self.d = self.data_vector.shape[1] |
| 24 | if self.data_vector is None: |
| 25 | self.data_vector = self._load_dataset() |
| 26 | self._init( sport_mode ) |
| 27 | |
| 28 | def _load_dataset(self): |
| 29 | pass |
| 30 | |
| 31 | def _normalize(self, vector): |
| 32 | normalize_L2( vector ) |
| 33 | return vector |
| 34 | |
| 35 | def _init(self, sport_mode, nlist=cfg.faiss.n_clusters, nprobe=cfg.faiss.nprobe): |
| 36 | train_vector = self.data_vector.astype('float32') |
| 37 | train_vector = self._normalize( train_vector ) |
| 38 | if sport_mode: |
| 39 | # nlist = 3 # 聚类中心的个数 |
| 40 | quantizer = faiss.IndexFlatIP( self.d ) # the other index，需要以其他index作为基础 |
| 41 | self.index = faiss.IndexIVFFlat(quantizer, self.d, nlist, faiss.METRIC_INNER_PRODUCT) |
| 42 | # by default it performs inner-product search |
| 43 | assert not self.index.is_trained |
| 44 | self.index.train( train_vector ) |
| 45 | assert self.index.is_trained |
| 46 | self.index.nprobe = nprobe # default nprobe is 1, try a few more |
| 47 | self.index.add( train_vector ) |
| 48 | else: |
| 49 | self.index = faiss.IndexFlatIP( self.d ) |
| 50 | self.index.train( train_vector ) |
| 51 | assert self.index.is_trained |
| 52 | self.index.add( train_vector ) |
| 53 | |
| 54 | def predict(self, vector, topn=3): |
| 55 | vector = self._normalize( vector.astype('float32') ) |
| 56 | t1=time.time() |
| 57 | D, I = self.index.search(vector, topn) |
| 58 | t2 = time.time() |
| 59 | # print('faiss kmeans result times {}'.format(t2-t1)) |
| 60 | self.logging.info( "[ESEngine] faiss kmeans result times:{}".format( t2-t1 ) ) |
| 61 | #print(I[:topn]) # 表示最相近的前3个的index |
| 62 | #print(D[:topn]) # 表示最相近的前3个的相似度的值 |
| 63 | # return I, D |
| 64 | return [ dict(zip(i, d)) for i,d in zip(I, D) ] |
| 65 | |
| 66 | |
| 67 | |