From 3c59ed779e3b30ab4877ae94242c9df076df681a Mon Sep 17 00:00:00 2001
From: Gustaf Rydholm <gustaf.rydholm@gmail.com>
Date: Sat, 13 Apr 2024 13:54:58 +0200
Subject: Fix tqdm logging

---
 rag/upload.py | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

(limited to 'rag/upload.py')

diff --git a/rag/upload.py b/rag/upload.py
index 3c5a100..8567142 100644
--- a/rag/upload.py
+++ b/rag/upload.py
@@ -7,9 +7,6 @@ from tqdm import tqdm
 
 from rag.retriever.retriever import Retriever
 
-log.remove()
-log.add(lambda msg: tqdm.write(msg, end=""), colorize=True)
-
 
 @click.command()
 @click.option(
@@ -22,10 +19,12 @@ def main(directory: str):
     log.info(f"Uploading pfs found in directory {directory}...")
     retriever = Retriever()
     pdfs = Path(directory).glob("**/*.pdf")
-    for path in tqdm(pdfs):
+    for path in tqdm(list(pdfs)):
         retriever.add_pdf(path=path)
 
 
 if __name__ == "__main__":
+    log.remove()
+    log.add(lambda msg: tqdm.write(msg, end=""), colorize=True)
     load_dotenv()
     main()
-- 
cgit v1.2.3-70-g09d2