From c8218d39ae64c9918eba5724f3cd5f7ce97a9132 Mon Sep 17 00:00:00 2001 From: w5688414 Date: Tue, 30 Apr 2024 08:56:27 +0000 Subject: [PATCH] update offline_ann.py --- pipelines/utils/offline_ann.py | 8 ++++++-- pipelines/utils/offline_ann_mm.py | 8 ++++++-- 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/pipelines/utils/offline_ann.py b/pipelines/utils/offline_ann.py index 200a603fe8c5..689d8d562eb2 100644 --- a/pipelines/utils/offline_ann.py +++ b/pipelines/utils/offline_ann.py @@ -140,8 +140,12 @@ def offline_ann(index_name, doc_dir): use_gpu=use_gpu, embed_title=args.embed_title, ) - # Writing docs may take a while. so waitting for 3 seconds for writing docs to be completed. - time.sleep(3) + # Writing docs may take a while. so waitting until writing docs to be completed. + document_count = document_store.get_document_count() + while document_count == 0: + time.sleep(1) + print("Waiting for writing docs to be completed.") + document_count = document_store.get_document_count() # 建立索引库 document_store.update_embeddings(retriever) diff --git a/pipelines/utils/offline_ann_mm.py b/pipelines/utils/offline_ann_mm.py index c72a90a1da92..729f897df0be 100644 --- a/pipelines/utils/offline_ann_mm.py +++ b/pipelines/utils/offline_ann_mm.py @@ -94,8 +94,12 @@ def offline_ann(index_name, doc_dir): query_type="image", document_embedding_models={"text": args.document_embedding_model}, ) - # Writing docs may take a while. so waitting for 3 seconds for writing docs to be completed. - time.sleep(3) + # Writing docs may take a while. so waitting until writing docs to be completed. + document_count = document_store.get_document_count() + while document_count == 0: + time.sleep(1) + print("Waiting for writing docs to be completed.") + document_count = document_store.get_document_count() # 建立索引库 document_store.update_embeddings(retriever_mm)