Adding back padding to query encoding

lightonai · Aug 19, 2024 · 1123a1c · 1123a1c
1 parent 80b3451
commit 1123a1c
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/giga_cherche/models/colbert.py b/giga_cherche/models/colbert.py
@@ -875,8 +875,8 @@ def tokenize(
         max_length = self.query_length if is_query else self.document_length
         self._first_module().max_seq_length = max_length
 
-        # Handle padding for documents if specified
-        tokenize_args = {"padding": "max_length"} if pad_document else {}
+        # Pad queries (query expansion) and handle padding for documents if specified
+        tokenize_args = {"padding": "max_length"} if pad_document or is_query else {}
 
         # Tokenize the texts
         tokenized_outputs = self._first_module().tokenize(texts, **tokenize_args)