add index

doc_similarity/answer_similarity.py

@@ -53,7 +53,8 @@ if __name__ == "__main__":
             level_dict[content_level].append(id)
             embedding_dict[id] = bc.encode([content]).tolist()[0]
 
-    print(random.choice(list(embedding_dict.items())))
+    tmp_tuple = random.choice(list(embedding_dict.items()))
+    print(tmp_tuple)
     answer_ids = np.array(list(embedding_dict.keys())).astype("int")
     answer_embeddings = np.array(list(embedding_dict.values())).astype("float32")
     print(answer_embeddings.shape)
@@ -66,7 +67,6 @@ if __name__ == "__main__":
     print("trained: " + str(index2.is_trained))
     print("total index: " + str(index2.ntotal))
 
-    for i in [1020398]:
-        D, I = index2.search(np.array(answer_embeddings[i]).astype("float32"))
-        res = I.tolist()
-        print(res, "\n")
+    D, I = index2.search(np.array(answer_embeddings[tmp_tuple[0]]).astype("float32"))
+    res = I.tolist()
+    print(res, "\n")