Update app.py

app.py

@@ -8,43 +8,9 @@ from torch.utils.data import DataLoader
 from tqdm import tqdm
 from transformers import AutoProcessor
 
-from custom_colbert.models.paligemma_colbert_architecture import ColPali
-from custom_colbert.trainer.retrieval_evaluator import CustomEvaluator
-
-
-def process_images(processor, images, max_length: int = 50):
-    texts_doc = ["Describe the image."] * len(images)
-    images = [image.convert("RGB") for image in images]
-
-    batch_doc = processor(
-        text=texts_doc,
-        images=images,
-        return_tensors="pt",
-        padding="longest",
-        max_length=max_length + processor.image_seq_length,
-    )
-    return batch_doc
-
-
-def process_queries(processor, queries, mock_image, max_length: int = 50):
-    texts_query = []
-    for query in queries:
-        query = f"Question: {query}<unused0><unused0><unused0><unused0><unused0>"
-        texts_query.append(query)
-
-    batch_query = processor(
-        images=[mock_image.convert("RGB")] * len(texts_query),
-        # NOTE: the image is not used in batch_query but it is required for calling the processor
-        text=texts_query,
-        return_tensors="pt",
-        padding="longest",
-        max_length=max_length + processor.image_seq_length,
-    )
-    del batch_query["pixel_values"]
-
-    batch_query["input_ids"] = batch_query["input_ids"][..., processor.image_seq_length :]
-    batch_query["attention_mask"] = batch_query["attention_mask"][..., processor.image_seq_length :]
-    return batch_query
+from colpali_engine.models.paligemma_colbert_architecture import ColPali
+from colpali_engine.trainer.retrieval_evaluator import CustomEvaluator
+from colpali_engine.utils.colpali_processing_utils import process_images, process_queries
 
 
 def search(query: str, ds, images):
@@ -71,7 +37,7 @@ def index(file, ds):
     # run inference - docs
     dataloader = DataLoader(
         images,
-        batch_size=8,
+        batch_size=4,
         shuffle=False,
         collate_fn=lambda x: process_images(processor, x),
     )