update offline_ann.py

PaddlePaddle · Apr 30, 2024 · c8218d3 · c8218d3
1 parent 23a9cd7
commit c8218d3
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 4 deletions.
diff --git 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
@@ -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)