Setting the prefixes from stanford-nlp

lightonai · Sep 12, 2024 · 54788cf · 54788cf
1 parent a9ff1c7
commit 54788cf
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/pylate/models/colbert.py b/pylate/models/colbert.py
@@ -261,6 +261,9 @@ def __init__(
                         use_auth_token,
                     )
                 )
+                # Setting the prefixes from stanford-nlp models
+                self.query_prefix = "[unused0]"
+                self.document_prefix = "[unused1]"
                 logger.warning("Loaded the ColBERT model from Stanford NLP.")
             else:
                 # Add a linear projection layer to the model in order to project the embeddings to the desired size