fix: ColBERT linear size for smaller model

AnswerDotAI · Aug 19, 2024 · ecc2fe9 · ecc2fe9
1 parent 7f881ec
commit ecc2fe9
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 3 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -14,7 +14,7 @@ packages = [
 name = "rerankers" 
 
 
-version = "0.5.2"
+version = "0.5.2post1"
 
 description = "A unified API for various document re-ranking models."
 

diff --git a/rerankers/__init__.py b/rerankers/__init__.py
@@ -2,4 +2,4 @@
 from rerankers.documents import Document
 
 __all__ = ["Reranker", "Document"]
-__version__ = "0.5.2"
+__version__ = "0.5.2post1"
diff --git a/rerankers/models/colbert_ranker.py b/rerankers/models/colbert_ranker.py
@@ -88,7 +88,14 @@ class ColBERTModel(BertPreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
         self.bert = BertModel(config)
-        self.linear = nn.Linear(config.hidden_size, 128, bias=False)
+
+        # TODO: Load from artifact.metadata
+        if "small" in config._name_or_path:
+            linear_dim = 96
+        else:
+            linear_dim = 128
+        print("Linear Dim set to: {linear_dim} for downcasting")
+        self.linear = nn.Linear(config.hidden_size, linear_dim, bias=False)
         self.init_weights()
 
     def forward(