implement compute_score api

modeling_bert.py

@@ -421,7 +421,7 @@ class JinaBertSelfAttention(nn.Module):
         # seem a bit unusual, but is taken from the original Transformer paper.
         attention_probs = self.dropout(attention_probs)
 
-        # Add the alibi matrix to the attention_scores after the call to softmax 
+        # Add the alibi matrix to the attention_scores after the call to softmax
         attention_scores += bias
 
         # Mask heads if we want to
@@ -435,7 +435,7 @@ class JinaBertSelfAttention(nn.Module):
         context_layer = context_layer.view(new_context_layer_shape)
 
         outputs = (
-            (context_layer, attention_probs if output_attention_probs else attention_scores) 
+            (context_layer, attention_probs if output_attention_probs else attention_scores)
             if output_attentions else (context_layer,)
         )
 
@@ -2072,6 +2072,57 @@ class JinaBertForSequenceClassification(JinaBertPreTrainedModel):
             attentions=outputs.attentions,
         )
 
+    @torch.inference_mode()
+    def compute_score(
+        self,
+        sentence_pairs: Union[List[Tuple[str, str]], Tuple[str, str]],
+        batch_size: int = 32,
+        device: Optional[torch.device] = None,
+        **tokenizer_kwargs,
+    ):
+        assert isinstance(sentence_pairs, list)
+        if isinstance(sentence_pairs[0], str):
+            sentence_pairs = [sentence_pairs]
+
+        if not hasattr(self, 'tokenizer'):
+            from transformers import AutoTokenizer
+
+            self.tokenizer = AutoTokenizer.from_pretrained(self._name_or_path)
+
+        is_training = self.training
+        self.eval()
+
+        if device is not None:
+            self.to(device)
+
+        all_scores = []
+        for start_index in range(
+            0, len(sentence_pairs), batch_size
+        ):
+            sentences_batch = sentence_pairs[
+                start_index : start_index + (batch_size or self._eval_batch_size)
+            ]
+            inputs = self._tokenizer(
+                sentences_batch,
+                padding=True,
+                truncation=True,
+                return_tensors='pt',
+                **tokenizer_kwargs,
+            ).to(self.device)
+
+            scores = (
+                self.forward(**inputs, return_dict=True)
+                .logits.view(
+                    -1,
+                )
+                .float()
+            )
+            all_scores.extend(scores.cpu().numpy().tolist())
+
+        if len(all_scores) == 1:
+            return all_scores[0]
+        return all_scores
+
 
 @add_start_docstrings(
     """