feat: add state dict loading function for UniCoilEncoder model

Author: RaykKretzschmar

tests/test_models/test_coil.py

@@ -109,6 +109,17 @@ def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor | None =
         return out
 
 
+def _load_unicoil_state_dict(model: UniCoilEncoder, hf_model: str) -> None:
+    state_dict_path = hf_hub_download(hf_model, filename="pytorch_model.bin")
+    state_dict = torch.load(state_dict_path, map_location="cpu")
+    state_dict = {
+        key.removeprefix("coil_encoder."): value
+        for key, value in state_dict.items()
+        if key.startswith(("coil_encoder.bert.", "coil_encoder.tok_proj."))
+    }
+    model.load_state_dict(state_dict, strict=False)
+
+
 @pytest.mark.model
 @pytest.mark.parametrize(
     "hf_model",
@@ -175,6 +186,7 @@ def test_same_as_unicoil(hf_model: str):
     ]
 
     orig_model = UniCoilEncoder.from_pretrained(hf_model)
+    _load_unicoil_state_dict(orig_model, hf_model)
     orig_tokenizer = AutoTokenizer.from_pretrained(hf_model)
     orig_query_encoded = orig_tokenizer(
         query, padding=True, truncation=True, max_length=512, return_tensors="pt", return_token_type_ids=False