Upload model

modeling.py

@@ -155,12 +155,12 @@ class EntityFusionLayer(nn.Module):
 
 
 class KPRMixin:
-    def _forward(self, **inputs: Tensor | dict[str, Tensor]) -> tuple[Tensor] | tuple[Tensor, Tensor] | ModelOutput:
+    def _forward(self, **inputs: dict[str, Tensor]) -> tuple[Tensor] | tuple[Tensor, Tensor] | ModelOutput:
         return_dict = inputs.pop("return_dict", True)
 
         if self.training:
-            query_embeddings = self.encode(inputs["queries"])
-            passage_embeddings = self.encode(inputs["passages"])
+            query_embeddings = self.encode(**inputs["queries"])
+            passage_embeddings = self.encode(**inputs["passages"])
 
             query_embeddings = self._dist_gather_tensor(query_embeddings)
             passage_embeddings = self._dist_gather_tensor(passage_embeddings)
@@ -179,13 +179,13 @@ class KPRMixin:
                 return (loss, scores)
 
         else:
-            sentence_embeddings = self.encode(inputs).unsqueeze(1)
+            sentence_embeddings = self.encode(**inputs).unsqueeze(1)
             if return_dict:
                 return ModelOutput(sentence_embeddings=sentence_embeddings)
             else:
                 return (sentence_embeddings,)
 
-    def encode(self, inputs: dict[str, Tensor]) -> Tensor:
+    def encode(self, **inputs: dict[str, Tensor]) -> Tensor:
         entity_ids = inputs.pop("entity_ids", None)
         entity_position_ids = inputs.pop("entity_position_ids", None)
         entity_embeds = inputs.pop("entity_embeds", None)