save index

doc_similarity/answer_similarity.py

 import os
-import sys
 import random
+import sys
 
 sys.path.append(os.path.realpath("."))
+
+import faiss
 import numpy as np
 from bert_serving.client import BertClient
 from utils.es import es_scan, get_answer_info_from_es
-import faiss
 
 
 def cos_sim(vector_a, vector_b):
@@ -43,13 +44,11 @@ if __name__ == "__main__":
     count = 0
     embedding_dict = {}
     for item in get_answer_info_from_es(["id", "answer", "content_level"]):
-        if count < 1000:
         count += 1
         id = int(item["_id"])
         print(count, id)
         content = item["_source"]["answer"]
         content_level = str(item["_source"]["content_level"])
-            # print(id, content_level, content)
         level_dict[content_level].append(id)
         embedding_dict[id] = bc.encode([content]).tolist()[0]
 
@@ -67,6 +66,12 @@ if __name__ == "__main__":
     print("trained: " + str(index2.is_trained))
     print("total index: " + str(index2.ntotal))
 
+    base_dir = os.getcwd()
+    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)
+    print(index_path)
+
     id = tmp_tuple[0]
     emb = np.array([embedding_dict[id]]).astype("float32")
     print(emb)
@@ -78,4 +83,3 @@ if __name__ == "__main__":
         if i <= 0.1:
             res.append(ids[index])
     print(res, "\n")
-    print(D)