add limit

doc_similarity/answer_similarity.py

@@ -42,14 +42,15 @@ if __name__ == "__main__":
     count = 0
     embedding_dict = {}
     for item in get_answer_info_from_es(["id", "answer", "content_level"]):
-        count += 1
-        print(count)
-        id = int(item["_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])
+        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])
 
     answer_ids = np.array(list(embedding_dict.keys())).astype("int")
     answer_embeddings = np.array(list(embedding_dict.values())).astype("float32")