fix dimenstions again

config.json

@@ -12,8 +12,8 @@
   "gradient_checkpointing": false,
   "hidden_act": "gelu",
   "hidden_dropout_prob": 0.1,
-  "hidden_size": 1792,
-  "base_hidden_size": 1024,
+  "hidden_size": 1024,
+  "list_con_hidden_size": 1792,
   "id2label": {
     "0": "LABEL_0"
   },

listconranker.py

@@ -1,20 +1,20 @@
 # 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 
+# 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 
+# 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 
+# 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.
 
 import math
@@ -23,47 +23,45 @@ 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, 
+    AutoTokenizer,
+    is_torch_npu_available,
+    AutoModel,
+    PreTrainedModel,
     PretrainedConfig,
     AutoConfig,
     BertModel,
-    BertConfig
+    BertConfig,
 )
 from transformers.modeling_outputs import SequenceClassifierOutput
 from typing import Union, List, Optional
 
 
-class ListConRankerConfig(PretrainedConfig):
+class ListConRankerConfig(BertConfig):
     """Configuration class for ListConRanker model."""
-    
+
     model_type = "ListConRanker"
-    
+
     def __init__(
         self,
         list_transformer_layers: int = 2,
-        hidden_size: int = 1792,
-        base_hidden_size: int = 1024,
+        list_con_hidden_size: int = 1792,
         num_labels: int = 1,
-        **kwargs
+        **kwargs,
     ):
         super().__init__(**kwargs)
         self.list_transformer_layers = list_transformer_layers
-        self.hidden_size = hidden_size
-        self.base_hidden_size = base_hidden_size
+        self.list_con_hidden_size = list_con_hidden_size
         self.num_labels = num_labels
 
         self.bert_config = BertConfig(**kwargs)
-        self.bert_config.hidden_size = self.base_hidden_size
         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.hidden_size)
-        self.layerNorm = nn.LayerNorm(config.hidden_size)
+        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)
@@ -71,40 +69,70 @@ class QueryEmbedding(nn.Module):
         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(1792, 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.list_transformer_layer = nn.TransformerEncoderLayer(
+            1792,
+            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.hidden_size * 2, config.hidden_size)
-        self.linear_score2 = nn.Linear(config.hidden_size * 2, config.hidden_size)
-        self.linear_score1 = nn.Linear(config.hidden_size * 2, 1)
+        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 = pair_features.size(0)
-        pair_nums = [x + 1 for x in pair_nums]
+    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_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)
+            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)
+        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 = 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)
+        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)
+        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 = []
@@ -112,20 +140,39 @@ class ListTransformer(nn.Module):
         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_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()
+            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):
@@ -147,6 +194,7 @@ class ListConRankerModel(PreTrainedModel):
     """
     ListConRanker model for sequence classification that's compatible with AutoModelForSequenceClassification.
     """
+
     config_class = ListConRankerConfig
     base_model_prefix = "listconranker"
 
@@ -155,14 +203,17 @@ class ListConRankerModel(PreTrainedModel):
         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.base_hidden_size, config.hidden_size)
+        self.linear_in_embedding = nn.Linear(
+            config.hidden_size, config.list_con_hidden_size
+        )
         self.list_transformer = ListTransformer(
-            config.list_transformer_layers, 
-            config, 
+            config.list_transformer_layers,
+            config,
         )
+        self.sep_token_id = 102  # [SEP] token ID
 
     def forward(
         self,
@@ -176,55 +227,63 @@ class ListConRankerModel(PreTrainedModel):
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
-        **kwargs
-    ) -> Union[SequenceClassifierOutput, tuple]:      
+        **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
         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)
+                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)
 
-            logits, pair_features_after_list_transformer = self.list_transformer(pair_features)
+            logits, pair_features_after_list_transformer = self.list_transformer(
+                pair_features
+            )
             logits = self.sigmoid(logits)
 
             return logits
 
     def average_pooling(self, hidden_state, attention_mask):
-        extended_attention_mask = attention_mask.unsqueeze(-1).expand(hidden_state.size()).to(dtype=hidden_state.dtype)
+        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)
-        
+    def from_pretrained(
+        cls, model_name_or_path, config: Optional[ListConRankerConfig] = None, **kwargs
+    ):
+        model = super().from_pretrained(model_name_or_path, config=config, **kwargs)
+
         # Load custom weights
         linear_path = f"{model_name_or_path}/linear_in_embedding.pt"
         transformer_path = f"{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}")
-        
+
         return model