adding tqdm to retrieval step

pyg-team · Jan 7, 2025 · c010843 · c010843
1 parent 032c2cd
commit c010843
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/examples/llm/hotpot_qa.py b/examples/llm/hotpot_qa.py
@@ -130,7 +130,11 @@
     https://github.com/pyg-team/pytorch_geometric/pull/9806
     """
     precisions = []
-    for QA_pair in relevant_triples.keys():
+    if args.verbose:
+        loader = relevant_triples.keys()
+    else:
+        loader = tqdm(relevant_triples.keys())
+    for QA_pair in loader:
         golden_triples = relevant_triples[QA_pair]
         q = QA_pair[0]
         retrieved_subgraph = query_loader.query(q)