write data to diary

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, DATA_PATH
+from utils.files import DATA_PATH, MODEL_PATH
 
 
 def save_result():
@@ -18,31 +18,30 @@ def save_result():
 
     index_path = os.path.join(MODEL_PATH, "faiss_diary_similarity.index")
     faiss_index = faiss.read_index(index_path)
-    print(faiss_index)
-
-    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"])
-            try:
-                emb = np.array([bc.encode([content]).tolist()[0]]).astype("float32")
-                D, I = faiss_index.search(emb, 10)
-                distances = D.tolist()[0]
-                ids = I.tolist()[0]
-                res = []
-                for (index, i) in enumerate(distances):
-                    tmp_id = ids[index]
-                    if i <= 1.0 and tmp_id != id:
-                        res.append(str(tmp_id))
-                if res:
-                    data = "{}:{}:{}".format(content_level, str(id), ",".join(res))
-                    print(data)
-                    f.write(data)
-            except Exception as e:
-                print(e)
+
+    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"])
+        try:
+            emb = np.array([bc.encode([content]).tolist()[0]]).astype("float32")
+            D, I = faiss_index.search(emb, 10)
+            distances = D.tolist()[0]
+            ids = I.tolist()[0]
+            res = []
+            for (index, i) in enumerate(distances):
+                tmp_id = ids[index]
+                if i <= 1.0 and tmp_id != id:
+                    res.append(str(tmp_id))
+            if res:
+                data = "{}:{}:{}".format(content_level, str(id), ",".join(res))
+                print(data)
+        except Exception as e:
+            print(e)
+    print("done")
+
 
 if __name__ == "__main__":
     # bc = BertClient("172.16.44.82", check_length=False)