Open wangbq18 opened 3 years ago
def search_annoy(self, owner_name, question, num=5): ''' Author: xiaoyichao param {type} Description: 使用Annoy 召回 ''' sentences = read_vec2bin.read_bert_sents(owner_name=owner_name) annoy_index_path = os.path.join( dir_name, '../es/search_model/%s_annoy.index' % owner_name) **encodearrary = self.sentenceBERT.get_bert([question])** tc_index = AnnoyIndex(f=512, metric='angular') tc_index.load(annoy_index_path) items = tc_index.get_nns_by_vector( encodearrary[0], num, include_distances=True) sim_questions = [sentences[num_annoy] for num_annoy in items[0]] # sims = items[1] # index_nums = items[0] return sim_questions def search_faiss(self, owner_name, question, num=5): ''' Author: xiaoyichao param {type} Description: 使用Faiss 召回 ''' sentences = read_vec2bin.read_bert_sents(owner_name=owner_name) faiss_index_path = os.path.join( dir_name, '../es/search_model/%s_faiss.index' % owner_name) index = faiss.read_index(faiss_index_path) **question_vec = np.array(bc.encode([question])).astype('float32')** index.nprobe = 1 sims, index_nums = index.search(question_vec, num) sim_questions = [sentences[num_faiss] for num_faiss in index_nums[0]] # index_nums = index_nums[0].tolist() # sims = sims[0].tolist() return sim_questions
好的,我周末看一下。感谢!
def search_annoy(self, owner_name, question, num=5): ''' Author: xiaoyichao param {type} Description: 使用Annoy 召回 ''' sentences = read_vec2bin.read_bert_sents(owner_name=owner_name) annoy_index_path = os.path.join( dir_name, '../es/search_model/%s_annoy.index' % owner_name) **encodearrary = self.sentenceBERT.get_bert([question])** tc_index = AnnoyIndex(f=512, metric='angular') tc_index.load(annoy_index_path) items = tc_index.get_nns_by_vector( encodearrary[0], num, include_distances=True) sim_questions = [sentences[num_annoy] for num_annoy in items[0]] # sims = items[1] # index_nums = items[0] return sim_questions def search_faiss(self, owner_name, question, num=5): ''' Author: xiaoyichao param {type} Description: 使用Faiss 召回 ''' sentences = read_vec2bin.read_bert_sents(owner_name=owner_name) faiss_index_path = os.path.join( dir_name, '../es/search_model/%s_faiss.index' % owner_name) index = faiss.read_index(faiss_index_path) **question_vec = np.array(bc.encode([question])).astype('float32')** index.nprobe = 1 sims, index_nums = index.search(question_vec, num) sim_questions = [sentences[num_faiss] for num_faiss in index_nums[0]] # index_nums = index_nums[0].tolist() # sims = sims[0].tolist() return sim_questions
已经修复,谢谢!