Commit 911a2f60 authored by 赵威's avatar 赵威

add limit

parent 723d710f
...@@ -42,14 +42,15 @@ if __name__ == "__main__": ...@@ -42,14 +42,15 @@ if __name__ == "__main__":
count = 0 count = 0
embedding_dict = {} embedding_dict = {}
for item in get_answer_info_from_es(["id", "answer", "content_level"]): for item in get_answer_info_from_es(["id", "answer", "content_level"]):
count += 1 if count < 1000:
print(count) count += 1
id = int(item["_id"]) id = int(item["_id"])
content = item["_source"]["answer"] print(count, id)
content_level = str(item["_source"]["content_level"]) content = item["_source"]["answer"]
# print(id, content_level, content) content_level = str(item["_source"]["content_level"])
level_dict[content_level].append(id) # print(id, content_level, content)
embedding_dict[id] = bc.encode([content]) level_dict[content_level].append(id)
embedding_dict[id] = bc.encode([content])
answer_ids = np.array(list(embedding_dict.keys())).astype("int") answer_ids = np.array(list(embedding_dict.keys())).astype("int")
answer_embeddings = np.array(list(embedding_dict.values())).astype("float32") answer_embeddings = np.array(list(embedding_dict.values())).astype("float32")
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment