finish integration

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,48 @@ 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_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_2",
+    "passage_21",
+  ]
+]
+scores = model.predict(inputs)
 print(scores)
-# [0.5126953125, 0.331298828125, 0.3642578125]
+# [0.4359, 0.3840, 0.3231]
 ```
 
 To reproduce the results with iterative inference, please run:

config.json

@@ -39,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

@@ -16,22 +16,18 @@
 # 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 math
 import torch
 from torch import nn
 from torch.nn import functional as F
-import numpy as np
 from transformers import (
-    AutoTokenizer,
-    is_torch_npu_available,
-    AutoModel,
     PreTrainedModel,
-    PretrainedConfig,
-    AutoConfig,
     BertModel,
     BertConfig,
+    AutoTokenizer,
 )
+import os
 from transformers.modeling_outputs import SequenceClassifierOutput
 from typing import Union, List, Optional
 
@@ -46,12 +42,16 @@ class ListConRankerConfig(BertConfig):
         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
@@ -75,7 +75,7 @@ class ListTransformer(nn.Module):
         super().__init__()
         self.config = config
         self.list_transformer_layer = nn.TransformerEncoderLayer(
-            1792,
+            config.list_con_hidden_size,
             self.config.num_attention_heads,
             batch_first=True,
             activation=F.gelu,
@@ -213,11 +213,10 @@ class ListConRankerModel(PreTrainedModel):
             config.list_transformer_layers,
             config,
         )
-        self.sep_token_id = 102  # [SEP] token ID
 
     def forward(
         self,
-        input_ids: Optional[torch.Tensor] = None,
+        input_ids: torch.Tensor,
         attention_mask: Optional[torch.Tensor] = None,
         token_type_ids: Optional[torch.Tensor] = None,
         position_ids: Optional[torch.Tensor] = None,
@@ -228,36 +227,157 @@ class ListConRankerModel(PreTrainedModel):
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
         **kwargs,
-    ) -> Union[SequenceClassifierOutput, tuple]:
-        # Get device
-        device = input_ids.device if input_ids is not None else inputs_embeds.device
-        self.list_transformer.device = device
-
-        # Forward through base model
+    ) -> Union[tuple[torch.Tensor], SequenceClassifierOutput]:
         if self.training:
-            pass
-        else:
-            ranker_out = self.hf_model(
-                input_ids=input_ids,
-                attention_mask=attention_mask,
-                token_type_ids=token_type_ids,
-                position_ids=position_ids,
-                head_mask=head_mask,
-                inputs_embeds=inputs_embeds,
-                output_attentions=output_attentions,
-                return_dict=True,
-            )
-            last_hidden_state = ranker_out.last_hidden_state
-
-            pair_features = self.average_pooling(last_hidden_state, attention_mask)
-            pair_features = self.linear_in_embedding(pair_features)
+            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,
+        )
 
-            logits, pair_features_after_list_transformer = self.list_transformer(
-                pair_features
+    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()
+
+            # 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:]
+            p_mask = mask[first_sep:]
+            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()
             )
-            logits = self.sigmoid(logits)
 
-            return logits
+            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 = (
@@ -275,15 +395,55 @@ class ListConRankerModel(PreTrainedModel):
         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
+        )
 
-        # Load custom weights
-        linear_path = f"{model_name_or_path}/linear_in_embedding.pt"
-        transformer_path = f"{model_name_or_path}/list_transformer.pt"
+        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:
-            print(f"Warning: Could not load custom weights from {model_name_or_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=True,
+                return_tensors="pt",
+            )
+            logits = self(**inputs)[0]
+            total_logits[batch * batch_size : (batch + 1) * batch_size] = (
+                logits.squeeze(1)
+            )
+        return total_logits

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]",