write data

doc_similarity/diary_similarity.py

@@ -10,7 +10,7 @@ import numpy as np
 from bert_serving.client import BertClient
 from utils.cache import redis_client_db
 from utils.es import get_diary_info_from_es
-from utils.files import MODEL_PATH
+from utils.files import MODEL_PATH, DATA_PATH
 
 
 def save_result():
@@ -20,14 +20,13 @@ def save_result():
     faiss_index = faiss.read_index(index_path)
     print(faiss_index)
 
-    # level_dict = {"6": set([]), "5": set([]), "4": set([]), "3.5": set([]), "3": set([])}
+    with open(os.path.join(DATA_PATH, "diary_similarity.log"), "w") as f:
         count = 0
         for item in get_diary_info_from_es(["id", "answer", "content_level"]):
             count += 1
             id = int(item["_id"])
             content = item["_source"]["answer"]
             content_level = str(item["_source"]["content_level"])
-        # level_dict[content_level].add(id)
             try:
                 emb = np.array([bc.encode([content]).tolist()[0]]).astype("float32")
                 D, I = faiss_index.search(emb, 10)
@@ -38,7 +37,8 @@ def save_result():
                     tmp_id = ids[index]
                     if i <= 1.0 and tmp_id != id:
                         res.append(tmp_id)
-            print(count, id, content_level, res)
+                if res:
+                    f.write("{}:{}:{}".format(content_level, id, ",".join(res)))
             except Exception as e:
                 print(e)