diff --git a/byaldi/colpali.py b/byaldi/colpali.py index cc11dcb..edb6f9f 100644 --- a/byaldi/colpali.py +++ b/byaldi/colpali.py @@ -10,6 +10,7 @@ from colpali_engine.models import ColPali, ColPaliProcessor, ColQwen2, ColQwen2Processor from pdf2image import convert_from_path from PIL import Image +from tqdm.auto import tqdm from byaldi.objects import Result @@ -357,8 +358,7 @@ def index( raise ValueError( f"Number of metadata entries ({len(metadata)}) does not match number of documents ({len(items)})" ) - for i, item in enumerate(items): - print(f"Indexing file: {item}") + for i, item in enumerate(tqdm(items, desc="Indexing files")): doc_id = doc_ids[i] if doc_ids else self.highest_doc_id + 1 doc_metadata = metadata[doc_id] if metadata else None self.add_to_index( @@ -741,4 +741,4 @@ def encode_query(self, query: Union[str, List[str]]) -> torch.Tensor: return embeddings.cpu() def get_doc_ids_to_file_names(self): - return self.doc_ids_to_file_names + return self.doc_ids_to_file_names \ No newline at end of file