对于请求反馈使用线程来提速
""" Function: get similarity query Author: dengyx DateTime: 20201019 """ import jieba import time import tqdm import threading import queue import numpy as np from gensim.models import KeyedVectors import logging logging.basicConfig(format='%(asctime)s : %(levelname)s : %(message)s', level=logging.INFO) from utils.soaring_vector.soaring_vector.soaring_vector import SoaringVectorClient, IndexType, FieldType, IndexData, SearchQuery, Vector, FloatVector client = SoaringVectorClient("172.16.24.150", 8098, 1000) print("health : ", client.health()) index_name = "seo-query-v10dim" if client.exist(index_name): api_index = client.get(index_name) print(index_name + " is exist") else: schema = {'query': FieldType.STRING_TYPE, 'id': FieldType.STRING_TYPE} api_index = client.create(index_name, "query search match", IndexType.FloatFlatIP, 'query', schema, 10, thread=12) client.set_alias(index_name, "seo-phrase-match") print(api_index.info) class QuerySimilarity(object): def __init__(self,): # self.query_path = r'data/test.txt' self.query_path = r'data/seo_search_word_copy.txt' self.w2c_path = r'resources_10dim/word2vec.model' self.query_features = r'resources/features.pkl' self.tables = r'resources/hashtables.pkl' self.table_num = 3 self.Hashcode_fun = 6 self.query2id = {} self.thread_num = 8 print('加载词向量...') t1 = time.time() self.model = KeyedVectors.load(self.w2c_path, mmap='r') t2 = time.time() print('词向量加载时间：{:.2f}s'.format(t2-t1)) with open(self.query_path, 'r', encoding='utf8') as fr: self.content = fr.readlines() for each in self.content: item = each.strip().split(' ') query_id = item[0] query = item[-1] self.query2id[query] = query_id def cosine_sim(self, x, y): num = x.dot(y.T) denom = np.linalg.norm(x) * np.linalg.norm(y) return num / denom def feature_extract(self, query): """ word -> feature :param query: :return: """ vec = [] tokens = jieba.lcut(query) for word in tokens: if word in self.model: vec.append(self.model[word]) else: vec.append([0]*10) # print('{} {} {} not in word2vec'.format(query, tokens, word)) vec = np.array(vec) mean_vec = np.mean(vec, axis=0) if len(mean_vec) != 10: print('向量纬度不是100') return mean_vec def upload_data(self): """ clean segment stopwords :return: """ self.counter = 0 # self.query2id = {} data_map_buffer = dict() for each in self.content: item = each.strip().split(' ') query_id = item[0] query = item[-1] # self.query2id[query] = query_id current_feature = self.feature_extract(query) vector = self.l2_norm(current_feature).tolist() data = {'query': query, 'id': query_id} data_map_buffer[query] = IndexData(data, vector) if len(data_map_buffer) > 1000: api_index.put(data_map_buffer) self.counter += len(data_map_buffer) data_map_buffer = dict() logging.info('put ' + str(self.counter)) if len(data_map_buffer) > 0: api_index.put(data_map_buffer) self.counter += len(data_map_buffer) logging.info('put ' + str(self.counter)) data_map_buffer = dict() print('数据上传完成') def l2_norm(self, m): dist = np.sqrt((m ** 2).sum(-1))[..., np.newaxis] m /= dist return m def download(self): with open(self.query_path, 'r', encoding='utf8') as fr: content = fr.readlines() new_content = [] for each in tqdm.tqdm(content): each_item = each.strip().split(' ') phrase = each_item[-1] api_vector = dict(api_index.get(phrase).data.vector.vector).get(phrase).floatVector.values query = SearchQuery(vector=Vector(floatVector=FloatVector(values=api_vector))) res = api_index.search(query, 0, 40) line = '' for ret in res.result: items = sorted(ret.item, key=lambda v: v.score, reverse=True) for item in items[1:31]: line += self.query2id[item.key] + '，' to_save = each.strip() + ' ' + line[:-1] + ' ' new_content.append(to_save) save_path = r'data/query_top30_20201021.txt' with open(save_path, 'w', encoding='utf8') as fw: fw.writelines(new_content) print('数据保存成功:{}'.format(save_path)) def run(self, q, fw): while True: if q.empty(): return else: sample = q.get() each_item = sample.strip().split(' ') phrase = each_item[-1] api_vector = dict(api_index.get(phrase).data.vector.vector).get(phrase).floatVector.values query = SearchQuery(vector=Vector(floatVector=FloatVector(values=api_vector))) res = api_index.search(query, 0, 40) line = '' # result = [] for ret in res.result: items = sorted(ret.item, key=lambda v: v.score, reverse=True) for item in items[1:31]: line += self.query2id[item.key] + '，' # result.append(item.key) to_save = sample.strip() + ' ' + line[:-1] + ' ' # print(result) # print(to_save) print(each_item[0]) fw.write(to_save) def main(self, data_path): q = queue.Queue() save_path = r'data/query_top30_20201022.txt' fw = open(save_path, 'a', encoding='utf8') # split_num = 250000 # with open(self.query_path, 'r', encoding='utf8') as fr: # content = fr.readlines() # for i in range(0, len(content), split_num): # split_data = content[i:i+split_num] # with open('data/split_data/group_{}.txt'.format(i), 'w', encoding='utf8') as fw: # fw.writelines(split_data) with open(data_path, 'r', encoding='utf8') as fr: content = fr.readlines() for d in tqdm.tqdm(content): q.put(d) print('数据放入队列完毕') t1 = time.time() threads = [] print('数据预测中...') for i in range(self.thread_num): t = threading.Thread(target=self.run, args=(q, fw)) threads.append(t) for i in range(self.thread_num): threads[i].start() for i in range(self.thread_num): threads[i].join() t2 = time.time() print('处理速度：{:.4f}sample/s'.format(len(content)/(t2-t1))) print('数据写入完毕') # Press the green button in the gutter to run the script. if __name__ == '__main__': data_path = r'data/seo_search_word_copy.txt' qs = QuerySimilarity() qs.main(data_path) # qs.upload_data()