fix shapes

listconranker.py

@@ -39,12 +39,11 @@ from typing import Union, List, Optional
 class ListConRankerConfig(PretrainedConfig):
     """Configuration class for ListConRanker model."""
     
-    model_type = "listconranker"
+    model_type = "ListConRanker"
     
     def __init__(
         self,
         list_transformer_layers: int = 2,
-        num_attention_heads: int = 8,
         hidden_size: int = 1792,
         base_hidden_size: int = 1024,
         num_labels: int = 1,
@@ -52,12 +51,12 @@ class ListConRankerConfig(PretrainedConfig):
     ):
         super().__init__(**kwargs)
         self.list_transformer_layers = list_transformer_layers
-        self.num_attention_heads = num_attention_heads
         self.hidden_size = hidden_size
         self.base_hidden_size = base_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):
@@ -85,7 +84,8 @@ class ListTransformer(nn.Module):
         self.linear_score2 = nn.Linear(config.hidden_size * 2, config.hidden_size)
         self.linear_score1 = nn.Linear(config.hidden_size * 2, 1)
 
-    def forward(self, pair_features, pair_nums):
+    def forward(self, pair_features: torch.Tensor):
+        pair_nums = pair_features.size(0)
         pair_nums = [x + 1 for x in pair_nums]
         batch_pair_features = pair_features.split(pair_nums)
 
@@ -154,7 +154,7 @@ class ListConRankerModel(PreTrainedModel):
         super().__init__(config)
         self.config = config
         self.num_labels = config.num_labels
-        self.hf_model = BertModel(config)
+        self.hf_model = BertModel(config.bert_config)
         
         self.sigmoid = nn.Sigmoid()
 
@@ -176,17 +176,8 @@ class ListConRankerModel(PreTrainedModel):
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
-        pair_num: Optional[torch.Tensor] = None,
         **kwargs
     ) -> Union[SequenceClassifierOutput, tuple]:      
-        # Handle pair_num parameter
-        if pair_num is not None:
-            pair_nums = pair_num.tolist()
-        else:
-            # Default behavior if pair_num is not provided
-            batch_size = input_ids.size(0) if input_ids is not None else inputs_embeds.size(0)
-            pair_nums = [1] * batch_size
-        
         # Get device
         device = input_ids.device if input_ids is not None else inputs_embeds.device
         self.list_transformer.device = device
@@ -195,20 +186,7 @@ class ListConRankerModel(PreTrainedModel):
         if self.training:
             pass
         else:
-            split_batch = 400
-            if sum(pair_nums) > split_batch:
-                last_hidden_state_list = []
-                input_ids_list = input_ids.split(split_batch)
-                attention_mask_list = attention_mask.split(split_batch)
-                for i in range(len(input_ids_list)):
-                    last_hidden_state = self.hf_model(
-                        input_ids=input_ids_list[i],
-                          attention_mask=attention_mask_list[i],
-                            return_dict=True).hidden_states[-1]
-                    last_hidden_state_list.append(last_hidden_state)
-                last_hidden_state = torch.cat(last_hidden_state_list, dim=0)
-            else:
-                ranker_out = self.hf_model(
+            ranker_out = self.hf_model(
                     input_ids=input_ids,
                     attention_mask=attention_mask,
                     token_type_ids=token_type_ids,
@@ -217,12 +195,12 @@ class ListConRankerModel(PreTrainedModel):
                     inputs_embeds=inputs_embeds,
                     output_attentions=output_attentions,
                     return_dict=True)
-                last_hidden_state = ranker_out.last_hidden_state
+            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, pair_nums)
+            logits, pair_features_after_list_transformer = self.list_transformer(pair_features)
             logits = self.sigmoid(logits)
 
             return logits
@@ -249,4 +227,4 @@ class ListConRankerModel(PreTrainedModel):
         except FileNotFoundError:
             print(f"Warning: Could not load custom weights from {model_name_or_path}")
         
-        return model
+        return model