fix: add missing module

README.md

@@ -90,7 +90,7 @@ results = model.rank(query, documents, return_documents=True, top_k=3)
 from transformers import AutoModelForSequenceClassification
 
 model = AutoModelForSequenceClassification.from_pretrained(
-    'jinaai/jina-reranker-v2-base-multilingual', num_labels=1, trust_remote_code=True
+    'jinaai/jina-reranker-v2-base-multilingual', trust_remote_code=True
 )
 
 # Example query and documents

modeling_xlm_roberta.py

@@ -902,3 +902,217 @@ class XLMRobertaForSequenceClassification(XLMRobertaPreTrainedModel):
             hidden_states=outputs.hidden_states,
             attentions=outputs.attentions,
         )
+
+
+    @torch.inference_mode()
+    def compute_score(
+        self,
+        sentence_pairs: Union[List[Tuple[str, str]], Tuple[str, str]],
+        batch_size: int = 32,
+        max_length: Optional[int] = None,
+    ) -> List[float]:
+
+        if not hasattr(self, "_tokenizer"):
+            from transformers import AutoTokenizer
+
+            self._tokenizer = AutoTokenizer.from_pretrained(
+                self.name_or_path, trust_remote_code=True
+            )
+
+        assert isinstance(sentence_pairs, list)
+        if isinstance(sentence_pairs[0], str):
+            sentence_pairs = [sentence_pairs]
+
+        all_scores = []
+        for start_index in range(
+            0, len(sentence_pairs), batch_size
+        ):
+            sentences_batch = sentence_pairs[
+                start_index : start_index + batch_size
+            ]
+            inputs = self._tokenizer(
+                sentences_batch,
+                padding=True,
+                truncation=True,
+                return_tensors='pt',
+                max_length=max_length,
+            ).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
+
+    def predict(
+        self,
+        sentence_pairs: Union[List[Tuple[str, str]], Tuple[str, str]],
+        batch_size: int = 32,
+        max_length: Optional[int] = None,
+    ) -> List[float]:
+        # used for beir evaluation
+        return self.compute_score(sentence_pairs, batch_size=batch_size, max_length=max_length)
+
+    def rerank(
+        self,
+        query: str,
+        documents: List[str],
+        batch_size: int = 32,
+        max_length: int = 1024,
+        max_query_length: int = 512,
+        overlap_tokens: int = 80,
+        top_n: Optional[int] = None,
+        **kwargs,
+    ):
+        assert max_length >= max_query_length * 2, (
+            f'max_length ({max_length}) must be greater than or equal to '
+            f'max_query_length ({max_query_length}) * 2'
+        )
+
+        if not hasattr(self, "_tokenizer"):
+            from transformers import AutoTokenizer
+
+            self._tokenizer = AutoTokenizer.from_pretrained(
+                self.name_or_path, trust_remote_code=True
+            )
+
+        # preproc of tokenization
+        sentence_pairs, sentence_pairs_pids = reranker_tokenize_preproc(
+            query,
+            documents,
+            tokenizer=self._tokenizer,
+            max_length=max_length,
+            max_query_length=max_query_length,
+            overlap_tokens=overlap_tokens,
+        )
+
+        tot_scores = []
+        with torch.no_grad():
+            for k in range(0, len(sentence_pairs), batch_size):
+                batch = self._tokenizer.pad(
+                    sentence_pairs[k : k + batch_size],
+                    padding=True,
+                    max_length=max_length,
+                    pad_to_multiple_of=None,
+                    return_tensors="pt",
+                )
+                batch_on_device = {k: v.to(self.device) for k, v in batch.items()}
+                scores = (
+                    self.forward(**batch_on_device, return_dict=True)
+                    .logits.view(
+                        -1,
+                    )
+                    .float()
+                )
+                scores = torch.sigmoid(scores)
+                tot_scores.extend(scores.cpu().numpy().tolist())
+
+        # ranking
+        merge_scores = [0 for _ in range(len(documents))]
+        for pid, score in zip(sentence_pairs_pids, tot_scores):
+            merge_scores[pid] = max(merge_scores[pid], score)
+
+        merge_scores_argsort = np.argsort(merge_scores)[::-1]
+        sorted_documents = []
+        sorted_scores = []
+        for mid in merge_scores_argsort:
+            sorted_scores.append(merge_scores[mid])
+            sorted_documents.append(documents[mid])
+
+        top_n = min(top_n or len(sorted_documents), len(sorted_documents))
+
+        return [
+            {
+                'document': sorted_documents[i],
+                'relevance_score': sorted_scores[i],
+                'index': merge_scores_argsort[i],
+            }
+            for i in range(top_n)
+        ]
+
+
+def reranker_tokenize_preproc(
+    query: str,
+    passages: List[str],
+    tokenizer=None,
+    max_length: int = 1024,
+    max_query_length: int = 512,
+    overlap_tokens: int = 80,
+):
+    from copy import deepcopy
+
+    assert tokenizer is not None, "Please provide a valid tokenizer for tokenization!"
+    sep_id = tokenizer.sep_token_id
+
+    def _merge_inputs(chunk1_raw, chunk2):
+        chunk1 = deepcopy(chunk1_raw)
+        chunk1['input_ids'].append(sep_id)
+        chunk1['input_ids'].extend(chunk2['input_ids'])
+        chunk1['input_ids'].append(sep_id)
+        chunk1['attention_mask'].append(chunk2['attention_mask'][0])
+        chunk1['attention_mask'].extend(chunk2['attention_mask'])
+        chunk1['attention_mask'].append(chunk2['attention_mask'][-1])
+        if 'token_type_ids' in chunk1:
+            token_type_ids = [1 for _ in range(len(chunk2['token_type_ids']) + 2)]
+            chunk1['token_type_ids'].extend(token_type_ids)
+        return chunk1
+
+    # Note: the long query will be truncated to 256 tokens by default
+    query_inputs = tokenizer.encode_plus(
+        query, truncation=True, padding=False, max_length=max_query_length
+    )
+
+    max_passage_inputs_length = max_length - len(query_inputs['input_ids']) - 2
+    # assert (
+    #     max_passage_inputs_length > 100
+    # ), "Your query is too long! Please make sure your query less than 500 tokens!"
+
+    overlap_tokens_implt = min(overlap_tokens, max_passage_inputs_length // 4)
+
+    res_merge_inputs = []
+    res_merge_inputs_pids = []
+    for pid, passage in enumerate(passages):
+        passage_inputs = tokenizer.encode_plus(
+            passage,
+            truncation=False,
+            padding=False,
+            add_special_tokens=False,
+            max_length=0,
+        )
+        passage_inputs_length = len(passage_inputs['input_ids'])
+
+        if passage_inputs_length <= max_passage_inputs_length:
+            qp_merge_inputs = _merge_inputs(query_inputs, passage_inputs)
+            res_merge_inputs.append(qp_merge_inputs)
+            res_merge_inputs_pids.append(pid)
+        else:
+            start_id = 0
+            while start_id < passage_inputs_length:
+                end_id = start_id + max_passage_inputs_length
+                # make sure the length of the last chunk is `max_passage_inputs_length`
+                if end_id >= passage_inputs_length:
+                    sub_passage_inputs = {
+                        k: v[-max_passage_inputs_length:]
+                        for k, v in passage_inputs.items()
+                    }
+                else:
+                    sub_passage_inputs = {
+                        k: v[start_id:end_id] for k, v in passage_inputs.items()
+                    }
+                start_id = (
+                    end_id - overlap_tokens_implt
+                    if end_id < passage_inputs_length
+                    else end_id
+                )
+
+                qp_merge_inputs = _merge_inputs(query_inputs, sub_passage_inputs)
+                res_merge_inputs.append(qp_merge_inputs)
+                res_merge_inputs_pids.append(pid)
+
+    return res_merge_inputs, res_merge_inputs_pids