From 14c522a22d62c5662a72c15298b1c5e6333d3781 Mon Sep 17 00:00:00 2001
From: ManuelFay <manuel.faysse@illuin.tech>
Date: Mon, 2 Sep 2024 15:48:42 +0200
Subject: [PATCH] add suffix to query processing func

---
 colpali_engine/utils/colpali_processing_utils.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/colpali_engine/utils/colpali_processing_utils.py b/colpali_engine/utils/colpali_processing_utils.py
index 4d2e577f..646fe20d 100644
--- a/colpali_engine/utils/colpali_processing_utils.py
+++ b/colpali_engine/utils/colpali_processing_utils.py
@@ -15,10 +15,14 @@ def process_images(processor, images, max_length: int = 50):
     return batch_doc
 
 
-def process_queries(processor, queries, mock_image, max_length: int = 50):
+def process_queries(processor, queries, mock_image, max_length: int = 50, suffix: str="default_suffix"):
+    if suffix == "default_suffix":
+        suffix = "<pad>" * 10
     texts_query = []
     for query in queries:
-        query = f"Question: {query}<unused0><unused0><unused0><unused0><unused0>"
+        query = f"Question: {query}"
+        # add pad tokens
+        query += suffix
         texts_query.append(query)
 
     batch_query = processor(