Commit 279f2323 authored by 赵威's avatar 赵威

write data

parent e337d89f
import json
import os import os
import random import random
import sys import sys
...@@ -7,7 +8,9 @@ sys.path.append(os.path.realpath(".")) ...@@ -7,7 +8,9 @@ sys.path.append(os.path.realpath("."))
import faiss import faiss
import numpy as np import numpy as np
from bert_serving.client import BertClient from bert_serving.client import BertClient
from utils.cache import redis_client_db
from utils.es import es_scan, get_answer_info_from_es from utils.es import es_scan, get_answer_info_from_es
from utils.files import MODEL_PATH
def cos_sim(vector_a, vector_b): def cos_sim(vector_a, vector_b):
...@@ -52,6 +55,8 @@ if __name__ == "__main__": ...@@ -52,6 +55,8 @@ if __name__ == "__main__":
level_dict[content_level].append(id) level_dict[content_level].append(id)
embedding_dict[id] = bc.encode([content]).tolist()[0] embedding_dict[id] = bc.encode([content]).tolist()[0]
redis_client_db.hmset("answer:level_dict", json.dumps(level_dict))
tmp_tuple = random.choice(list(embedding_dict.items())) tmp_tuple = random.choice(list(embedding_dict.items()))
print(tmp_tuple) print(tmp_tuple)
answer_ids = np.array(list(embedding_dict.keys())).astype("int") answer_ids = np.array(list(embedding_dict.keys())).astype("int")
...@@ -66,9 +71,7 @@ if __name__ == "__main__": ...@@ -66,9 +71,7 @@ if __name__ == "__main__":
print("trained: " + str(index2.is_trained)) print("trained: " + str(index2.is_trained))
print("total index: " + str(index2.ntotal)) print("total index: " + str(index2.ntotal))
base_dir = os.getcwd() index_path = os.path.join(MODEL_PATH, "faiss_answer_similarity.index")
model_dir = os.path.join(base_dir, "_models")
index_path = os.path.join(model_dir, "faiss_answer_similarity.index")
faiss.write_index(index2, index_path) faiss.write_index(index2, index_path)
print(index_path) print(index_path)
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment