write diary data

doc_similarity/diary_similarity.py

@@ -20,16 +20,15 @@ if __name__ == "__main__":
     embedding_dict = {}
     for item in get_diary_info_from_es(["id", "answer", "content_level"]):
         count += 1
-        if count < 1000:
-            id = int(item["_id"])
-            print(count, id)
-            content = item["_source"]["answer"]
-            content_level = str(item["_source"]["content_level"])
-            level_dict[content_level].append(id)
-            try:
-                embedding_dict[id] = bc.encode([content]).tolist()[0]
-            except Exception as e:
-                pass
+        id = int(item["_id"])
+        print(count, id)
+        content = item["_source"]["answer"]
+        content_level = str(item["_source"]["content_level"])
+        level_dict[content_level].append(id)
+        try:
+            embedding_dict[id] = bc.encode([content]).tolist()[0]
+        except Exception as e:
+            pass
 
     # redis_client_db.hmset("diary:level_dict", json.dumps(level_dict))
 
@@ -47,9 +46,9 @@ if __name__ == "__main__":
     print("trained: " + str(index2.is_trained))
     print("total index: " + str(index2.ntotal))
 
-    # index_path = os.path.join(MODEL_PATH, "faiss_diary_similarity.index")
-    # faiss.write_index(index2, index_path)
-    # print(index_path)
+    index_path = os.path.join(MODEL_PATH, "faiss_diary_similarity.index")
+    faiss.write_index(index2, index_path)
+    print(index_path)
 
     id = tmp_tuple[0]
     emb = np.array([embedding_dict[id]]).astype("float32")