summaryrefslogtreecommitdiff
path: root/rag/upload.py
diff options
context:
space:
mode:
Diffstat (limited to 'rag/upload.py')
-rw-r--r--rag/upload.py7
1 files changed, 3 insertions, 4 deletions
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()