fix: ColBERT linear size for smaller model

Author: bclavie

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."
 

rerankers/__init__.py

@@ -2,4 +2,4 @@
 from rerankers.documents import Document
 
 __all__ = ["Reranker", "Document"]
-__version__ = "0.5.2"
+__version__ = "0.5.2post1"

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(