Integrate with Transformers & SentenceTransformers

README.md

@@ -1,6 +1,9 @@
 ---
 tags:
 - mteb
+- sentence-transformers
+- transformers
+pipeline_tag: text-ranking
 model-index:
 - name: ListConRanker
   results:
@@ -103,9 +106,9 @@ To reduce the discrepancy between training and inference, we propose iterative i
 
 ## How to use
 ```python
-from modules.listconranker import ListConRanker
+from transfoermers import AutoModelForSequenceClassification
 
-reranker = ListConRanker('./', use_fp16=True, list_transformer_layer=2)
+reranker = AutoModelForSequenceClassification('ByteDance/ListConRanker', trust_remote_code=True)
 
 # [query, passages_1, passage_2, ..., passage_n]
 batch = [
@@ -125,15 +128,55 @@ batch = [
 ]
 
 # for conventional inference, please manage the batch size by yourself
-scores = reranker.compute_score(batch)
+scores = reranker.multi_passage(batch)
 print(scores)
 # [[0.5126953125, 0.331298828125, 0.3642578125], [0.63671875, 0.71630859375, 0.42822265625, 0.35302734375]]
 
-# for iterative inferfence, only a batch size of 1 is supported
-# the scores do not indicate similarity but are intended only for ranking
-scores = reranker.iterative_inference(batch[0])
+inputs = tokenizer(
+    [
+        [
+            "query 1",
+            "passage_11",
+        ],
+        [
+            "query 1",
+            "passage_12",
+        ],
+        [
+            "query_2",
+            "passage_21",
+        ],
+    ],
+    return_tensors="pt",
+    padding=True,
+)
+probs, logits = model(**inputs)
+print(probs)
+# tensor([[0.4359], [0.3840]], grad_fn=<ViewBackward0>)
+```
+or using the `sentence_transformers` library:
+```python
+from sentence_transformers import CrossEncoder
+
+model = CrossEncoder('ByteDance/ListConRanker', trust_remote_code=True)
+
+inputs = [
+  [
+    "query 1",
+    "passage_11",
+  ],
+  [
+    "query 1",
+    "passage_12",
+  ],
+  [
+    "query_2",
+    "passage_21",
+  ],
+]
+scores = model.predict(inputs)
 print(scores)
-# [0.5126953125, 0.331298828125, 0.3642578125]
+# [0.43585014, 0.32305932, 0.38395187]
 ```
 
 To reproduce the results with iterative inference, please run:

config.json

@@ -1,7 +1,11 @@
 {
   "architectures": [
-    "BertModel"
+    "ListConRanker"
   ],
+  "auto_map": {
+    "AutoConfig": "listconranker.ListConRankerConfig",
+    "AutoModelForSequenceClassification": "listconranker.ListConRankerModel"
+  },
   "attention_probs_dropout_prob": 0.1,
   "classifier_dropout": null,
   "directionality": "bidi",
@@ -9,6 +13,7 @@
   "hidden_act": "gelu",
   "hidden_dropout_prob": 0.1,
   "hidden_size": 1024,
+  "list_con_hidden_size": 1792,
   "id2label": {
     "0": "LABEL_0"
   },
@@ -34,5 +39,7 @@
   "transformers_version": "4.45.2",
   "type_vocab_size": 2,
   "use_cache": true,
-  "vocab_size": 21128
+  "vocab_size": 21128,
+  "cls_token_id": 101,
+  "sep_token_id": 102
 }

listconranker.py

@@ -0,0 +1,546 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy of this software
+# and associated documentation files (the "Software"), to deal in the Software without
+# restriction, including without limitation the rights to use, copy, modify, merge, publish,
+# distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all copies or
+# substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR
+# OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
+# ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
+# OTHER DEALINGS IN THE SOFTWARE.
+from __future__ import annotations
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from transformers import (
+    PreTrainedModel,
+    BertModel,
+    BertConfig,
+    AutoTokenizer,
+)
+import os
+from transformers.modeling_outputs import SequenceClassifierOutput
+from typing import Union, List, Optional
+from collections import defaultdict
+import numpy as np
+import math
+
+
+class ListConRankerConfig(BertConfig):
+    """Configuration class for ListConRanker model."""
+
+    model_type = "ListConRanker"
+
+    def __init__(
+        self,
+        list_transformer_layers: int = 2,
+        list_con_hidden_size: int = 1792,
+        num_labels: int = 1,
+        cls_token_id: int = 101,
+        sep_token_id: int = 102,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.list_transformer_layers = list_transformer_layers
+        self.list_con_hidden_size = list_con_hidden_size
+        self.num_labels = num_labels
+        self.cls_token_id = cls_token_id
+        self.sep_token_id = sep_token_id
+
+        self.bert_config = BertConfig(**kwargs)
+        self.bert_config.output_hidden_states = True
+
+
+class QueryEmbedding(nn.Module):
+    def __init__(self, config) -> None:
+        super().__init__()
+        self.query_embedding = nn.Embedding(2, config.list_con_hidden_size)
+        self.layerNorm = nn.LayerNorm(config.list_con_hidden_size)
+
+    def forward(self, x, tags):
+        query_embeddings = self.query_embedding(tags)
+        x += query_embeddings
+        x = self.layerNorm(x)
+        return x
+
+
+class ListTransformer(nn.Module):
+    def __init__(self, num_layer, config) -> None:
+        super().__init__()
+        self.config = config
+        self.list_transformer_layer = nn.TransformerEncoderLayer(
+            config.list_con_hidden_size,
+            self.config.num_attention_heads,
+            batch_first=True,
+            activation=F.gelu,
+            norm_first=False,
+        )
+        self.list_transformer = nn.TransformerEncoder(
+            self.list_transformer_layer, num_layer
+        )
+        self.relu = nn.ReLU()
+        self.query_embedding = QueryEmbedding(config)
+
+        self.linear_score3 = nn.Linear(
+            config.list_con_hidden_size * 2, config.list_con_hidden_size
+        )
+        self.linear_score2 = nn.Linear(
+            config.list_con_hidden_size * 2, config.list_con_hidden_size
+        )
+        self.linear_score1 = nn.Linear(config.list_con_hidden_size * 2, 1)
+
+    def forward(
+        self, pair_features: torch.Tensor, pair_nums: List[int]
+    ) -> torch.Tensor:
+        batch_pair_features = pair_features.split(pair_nums)
+
+        pair_feature_query_passage_concat_list = []
+        for i in range(len(batch_pair_features)):
+            pair_feature_query = (
+                batch_pair_features[i][0].unsqueeze(0).repeat(pair_nums[i] - 1, 1)
+            )
+            pair_feature_passage = batch_pair_features[i][1:]
+            pair_feature_query_passage_concat_list.append(
+                torch.cat([pair_feature_query, pair_feature_passage], dim=1)
+            )
+        pair_feature_query_passage_concat = torch.cat(
+            pair_feature_query_passage_concat_list, dim=0
+        )
+
+        batch_pair_features = nn.utils.rnn.pad_sequence(
+            batch_pair_features, batch_first=True
+        )
+
+        query_embedding_tags = torch.zeros(
+            batch_pair_features.size(0),
+            batch_pair_features.size(1),
+            dtype=torch.long,
+            device=self.device,
+        )
+        query_embedding_tags[:, 0] = 1
+        batch_pair_features = self.query_embedding(
+            batch_pair_features, query_embedding_tags
+        )
+
+        mask = self.generate_attention_mask(pair_nums)
+        query_mask = self.generate_attention_mask_custom(pair_nums)
+        pair_list_features = self.list_transformer(
+            batch_pair_features, src_key_padding_mask=mask, mask=query_mask
+        )
+
+        output_pair_list_features = []
+        output_query_list_features = []
+        pair_features_after_transformer_list = []
+        for idx, pair_num in enumerate(pair_nums):
+            output_pair_list_features.append(pair_list_features[idx, 1:pair_num, :])
+            output_query_list_features.append(pair_list_features[idx, 0, :])
+            pair_features_after_transformer_list.append(
+                pair_list_features[idx, :pair_num, :]
+            )
+
+        pair_features_after_transformer_cat_query_list = []
+        for idx, pair_num in enumerate(pair_nums):
+            query_ft = (
+                output_query_list_features[idx].unsqueeze(0).repeat(pair_num - 1, 1)
+            )
+            pair_features_after_transformer_cat_query = torch.cat(
+                [query_ft, output_pair_list_features[idx]], dim=1
+            )
+            pair_features_after_transformer_cat_query_list.append(
+                pair_features_after_transformer_cat_query
+            )
+        pair_features_after_transformer_cat_query = torch.cat(
+            pair_features_after_transformer_cat_query_list, dim=0
+        )
+
+        pair_feature_query_passage_concat = self.relu(
+            self.linear_score2(pair_feature_query_passage_concat)
+        )
+        pair_features_after_transformer_cat_query = self.relu(
+            self.linear_score3(pair_features_after_transformer_cat_query)
+        )
+        final_ft = torch.cat(
+            [
+                pair_feature_query_passage_concat,
+                pair_features_after_transformer_cat_query,
+            ],
+            dim=1,
+        )
+        logits = self.linear_score1(final_ft).squeeze()
+        return logits, torch.cat(pair_features_after_transformer_list, dim=0)
+
+    def generate_attention_mask(self, pair_num):
+        max_len = max(pair_num)
+        batch_size = len(pair_num)
+        mask = torch.zeros(batch_size, max_len, dtype=torch.bool, device=self.device)
+        for i, length in enumerate(pair_num):
+            mask[i, length:] = True
+        return mask
+
+    def generate_attention_mask_custom(self, pair_num):
+        max_len = max(pair_num)
+        mask = torch.zeros(max_len, max_len, dtype=torch.bool, device=self.device)
+        mask[0, 1:] = True
+        return mask
+
+
+class ListConRankerModel(PreTrainedModel):
+    """
+    ListConRanker model for sequence classification that's compatible with AutoModelForSequenceClassification.
+    """
+
+    config_class = ListConRankerConfig
+    base_model_prefix = "listconranker"
+
+    def __init__(self, config: ListConRankerConfig):
+        super().__init__(config)
+        self.config = config
+        self.num_labels = config.num_labels
+        self.hf_model = BertModel(config.bert_config)
+
+        self.sigmoid = nn.Sigmoid()
+
+        self.linear_in_embedding = nn.Linear(
+            config.hidden_size, config.list_con_hidden_size
+        )
+        self.list_transformer = ListTransformer(
+            config.list_transformer_layers,
+            config,
+        )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        token_type_ids: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs,
+    ) -> Union[tuple[torch.Tensor], SequenceClassifierOutput]:
+        if self.training:
+            raise NotImplementedError("Training not supported; use eval mode.")
+        device = input_ids.device
+        self.list_transformer.device = device
+        # Reorganize by unique queries and their passages
+        (
+            reorganized_input_ids,
+            reorganized_attention_mask,
+            reorganized_token_type_ids,
+            pair_nums,
+            group_indices,
+        ) = self._reorganize_inputs(input_ids, attention_mask, token_type_ids)
+
+        out = self.hf_model(
+            input_ids=reorganized_input_ids,
+            attention_mask=reorganized_attention_mask,
+            token_type_ids=reorganized_token_type_ids,
+            return_dict=True,
+        )
+        feats = out.last_hidden_state
+        pooled = self.average_pooling(feats, reorganized_attention_mask)
+        embedded = self.linear_in_embedding(pooled)
+        logits, _ = self.list_transformer(embedded, pair_nums)
+        probs = self.sigmoid(logits)
+
+        # Restore original order
+        sorted_probs = self._restore_original_order(probs, group_indices)
+        sorted_logits = self._restore_original_order(logits, group_indices)
+        if not return_dict:
+            return (sorted_probs, sorted_logits)
+
+        return SequenceClassifierOutput(
+            loss=None,
+            logits=sorted_logits,
+            hidden_states=out.hidden_states,
+            attentions=out.attentions,
+        )
+
+    def _reorganize_inputs(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        token_type_ids: Optional[torch.Tensor],
+    ) -> tuple[
+        torch.Tensor, torch.Tensor, Optional[torch.Tensor], List[int], List[List[int]]
+    ]:
+        """
+        Group inputs by unique queries: for each query, produce [query] + its passages,
+        then flatten, pad, and return pair sizes and original indices mapping.
+        """
+        batch_size = input_ids.size(0)
+        # Structure: query_key -> {
+        #   'query': (seq, mask, tt),
+        #   'passages': [(seq, mask, tt), ...],
+        #   'indices': [original_index, ...]
+        # }
+        grouped = {}
+
+        for idx in range(batch_size):
+            seq = input_ids[idx]
+            mask = attention_mask[idx]
+            token_type_ids[idx] if token_type_ids is not None else torch.zeros_like(seq)
+
+            sep_idxs = (seq == self.config.sep_token_id).nonzero(as_tuple=True)[0]
+            if sep_idxs.numel() == 0:
+                raise ValueError(f"No SEP in sequence {idx}")
+            first_sep = sep_idxs[0].item()
+            second_sep = sep_idxs[1].item()
+
+            # Extract query and passage
+            q_seq = seq[: first_sep + 1]
+            q_mask = mask[: first_sep + 1]
+            q_tt = torch.zeros_like(q_seq)
+
+            p_seq = seq[first_sep : second_sep + 1]
+            p_mask = mask[first_sep : second_sep + 1]
+            p_seq = p_seq.clone()
+            p_seq[0] = self.config.cls_token_id
+            p_tt = torch.zeros_like(p_seq)
+
+            # Build key excluding CLS/SEP
+            key = tuple(
+                q_seq[
+                    (q_seq != self.config.cls_token_id)
+                    & (q_seq != self.config.sep_token_id)
+                ].tolist()
+            )
+
+            # truncation
+            q_seq = q_seq[: self.config.max_position_embeddings]
+            q_seq[-1] = self.config.sep_token_id
+            p_seq = p_seq[: self.config.max_position_embeddings]
+            p_seq[-1] = self.config.sep_token_id
+            q_mask = q_mask[: self.config.max_position_embeddings]
+            p_mask = p_mask[: self.config.max_position_embeddings]
+            q_tt = q_tt[: self.config.max_position_embeddings]
+            p_tt = p_tt[: self.config.max_position_embeddings]
+
+            if key not in grouped:
+                grouped[key] = {
+                    "query": (q_seq, q_mask, q_tt),
+                    "passages": [],
+                    "indices": [],
+                }
+            grouped[key]["passages"].append((p_seq, p_mask, p_tt))
+            grouped[key]["indices"].append(idx)
+
+        # Flatten according to group insertion order
+        seqs, masks, tts, pair_nums, group_indices = [], [], [], [], []
+        for key, data in grouped.items():
+            q_seq, q_mask, q_tt = data["query"]
+            passages = data["passages"]
+            indices = data["indices"]
+            # record sizes and original positions
+            pair_nums.append(len(passages) + 1)  # +1 for the query
+            group_indices.append(indices)
+
+            # append query then its passages
+            seqs.append(q_seq)
+            masks.append(q_mask)
+            tts.append(q_tt)
+            for p_seq, p_mask, p_tt in passages:
+                seqs.append(p_seq)
+                masks.append(p_mask)
+                tts.append(p_tt)
+
+        # Pad to uniform length
+        max_len = max(s.size(0) for s in seqs)
+        padded_seqs, padded_masks, padded_tts = [], [], []
+        for s, m, t in zip(seqs, masks, tts):
+            ps = torch.zeros(max_len, dtype=s.dtype, device=s.device)
+            pm = torch.zeros(max_len, dtype=m.dtype, device=m.device)
+            pt = torch.zeros(max_len, dtype=t.dtype, device=t.device)
+            ps[: s.size(0)] = s
+            pm[: m.size(0)] = m
+            pt[: t.size(0)] = t
+            padded_seqs.append(ps)
+            padded_masks.append(pm)
+            padded_tts.append(pt)
+
+        rid = torch.stack(padded_seqs)
+        ram = torch.stack(padded_masks)
+        rtt = torch.stack(padded_tts) if token_type_ids is not None else None
+
+        return rid, ram, rtt, pair_nums, group_indices
+
+    def _restore_original_order(
+        self,
+        logits: torch.Tensor,
+        group_indices: List[List[int]],
+    ) -> torch.Tensor:
+        """
+        Map flattened logits back so each original index gets its passage score.
+        """
+        out = torch.zeros(logits.size(0), dtype=logits.dtype, device=logits.device)
+        i = 0
+        for indices in group_indices:
+            for idx in indices:
+                out[idx] = logits[i]
+                i += 1
+        return out.reshape(-1, 1)
+
+    def average_pooling(self, hidden_state, attention_mask):
+        extended_attention_mask = (
+            attention_mask.unsqueeze(-1)
+            .expand(hidden_state.size())
+            .to(dtype=hidden_state.dtype)
+        )
+        masked_hidden_state = hidden_state * extended_attention_mask
+        sum_embeddings = torch.sum(masked_hidden_state, dim=1)
+        sum_mask = extended_attention_mask.sum(dim=1)
+        return sum_embeddings / sum_mask
+
+    @classmethod
+    def from_pretrained(
+        cls, model_name_or_path, config: Optional[ListConRankerConfig] = None, **kwargs
+    ):
+        model = super().from_pretrained(model_name_or_path, config=config, **kwargs)
+        model.hf_model = BertModel.from_pretrained(
+            model_name_or_path, config=model.config.bert_config, **kwargs
+        )
+
+        linear_path = os.path.join(model_name_or_path, "linear_in_embedding.pt")
+        transformer_path = os.path.join(model_name_or_path, "list_transformer.pt")
+
+        try:
+            model.linear_in_embedding.load_state_dict(torch.load(linear_path))
+            model.list_transformer.load_state_dict(torch.load(transformer_path))
+        except FileNotFoundError as e:
+            raise e
+
+        return model
+
+    def multi_passage(
+        self,
+        sentences: List[List[str]],
+        batch_size: int = 32,
+        tokenizer: AutoTokenizer = AutoTokenizer.from_pretrained(
+            "ByteDance/ListConRanker"
+        ),
+    ):
+        """
+        Process multiple passages for each query.
+        :param sentences: List of lists, where each inner list contains sentences for a query.
+        :return: Tensor of logits for each passage.
+        """
+        pairs = []
+        for batch in sentences:
+            if len(batch) < 2:
+                raise ValueError("Each query must have at least one passage.")
+            query = batch[0]
+            passages = batch[1:]
+            for passage in passages:
+                pairs.append((query, passage))
+
+        total_batches = (len(pairs) + batch_size - 1) // batch_size
+        total_logits = torch.zeros(len(pairs), dtype=torch.float, device=self.device)
+        for batch in range(total_batches):
+            batch_pairs = pairs[batch * batch_size : (batch + 1) * batch_size]
+            inputs = tokenizer(
+                batch_pairs,
+                padding=True,
+                truncation=False,
+                return_tensors="pt",
+            )
+
+            for k, v in inputs.items():
+                inputs[k] = v.to(self.device)
+
+            logits = self(**inputs)[0]
+            total_logits[batch * batch_size : (batch + 1) * batch_size] = (
+                logits.squeeze(1)
+            )
+        return total_logits
+
+    def multi_passage_in_iterative_inference(
+        self,
+        sentences: List[str],
+        stop_num: int = 20,
+        decrement_rate: float = 0.2,
+        min_filter_num: int = 10,
+        tokenizer: AutoTokenizer = AutoTokenizer.from_pretrained(
+            "ByteDance/ListConRanker"
+        ),
+    ):
+        """
+        Process multiple passages for one query in iterative inference.
+        :param sentences: List contains sentences for a query.
+        :return: Tensor of logits for each passage.
+        """
+        if stop_num < 1:
+            raise ValueError("stop_num must be greater than 0")
+        if decrement_rate <= 0 or decrement_rate >= 1:
+            raise ValueError("decrement_rate must be in (0, 1)")
+        if min_filter_num < 1:
+            raise ValueError("min_filter_num must be greater than 0")
+
+        query = sentences[0]
+        passage = sentences[1:]
+
+        filter_times = 0
+        passage2score = defaultdict(list)
+        while len(passage) > stop_num:
+            batch = [[query] + passage]
+            pred_scores = self.multi_passage(
+                batch, batch_size=len(batch[0]) - 1, tokenizer=tokenizer
+            ).tolist()
+            pred_scores_argsort = np.argsort(
+                pred_scores
+            ).tolist()  # Sort in increasing order
+
+            passage_len = len(passage)
+            to_filter_num = math.ceil(passage_len * decrement_rate)
+            if to_filter_num < min_filter_num:
+                to_filter_num = min_filter_num
+
+            have_filter_num = 0
+            while have_filter_num < to_filter_num:
+                idx = pred_scores_argsort[have_filter_num]
+                passage2score[passage[idx]].append(pred_scores[idx] + filter_times)
+                have_filter_num += 1
+            while (
+                pred_scores[pred_scores_argsort[have_filter_num - 1]]
+                == pred_scores[pred_scores_argsort[have_filter_num]]
+            ):
+                idx = pred_scores_argsort[have_filter_num]
+                passage2score[passage[idx]].append(pred_scores[idx] + filter_times)
+                have_filter_num += 1
+            next_passage = []
+            next_passage_idx = have_filter_num
+            while next_passage_idx < len(passage):
+                idx = pred_scores_argsort[next_passage_idx]
+                next_passage.append(passage[idx])
+                next_passage_idx += 1
+            passage = next_passage
+            filter_times += 1
+
+        batch = [[query] + passage]
+        pred_scores = self.multi_passage(
+            batch, batch_size=len(batch[0]) - 1, tokenizer=tokenizer
+        ).tolist()
+
+        cnt = 0
+        while cnt < len(passage):
+            passage2score[passage[cnt]].append(pred_scores[cnt] + filter_times)
+            cnt += 1
+
+        passage = sentences[1:]
+        final_score = []
+        for i in range(len(passage)):
+            p = passage[i]
+            final_score.append(passage2score[p][0])
+        return final_score

tokenizer_config.json

@@ -47,7 +47,7 @@
   "do_lower_case": true,
   "mask_token": "[MASK]",
   "max_length": 512,
-  "model_max_length": 1000000000000000019884624838656,
+  "model_max_length": 512,
   "never_split": null,
   "pad_to_multiple_of": null,
   "pad_token": "[PAD]",