Making configurable

configuration_hf_alibaba_nlp_gte.py

@@ -116,6 +116,8 @@ class GteConfig(PretrainedConfig):
         use_memory_efficient_attention=False,
         logn_attention_scale=False,
         logn_attention_clip1=False,
+        add_pooling_layer=True,
+        num_labels=0,
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -142,4 +144,7 @@ class GteConfig(PretrainedConfig):
         self.unpad_inputs = unpad_inputs
         self.use_memory_efficient_attention = use_memory_efficient_attention
         self.logn_attention_scale = logn_attention_scale
-        self.logn_attention_clip1 = logn_attention_clip1
+        self.logn_attention_clip1 = logn_attention_clip1
+
+        self.add_pooling_layer = add_pooling_layer
+        self.num_labels = num_labels

modeling_hf_alibaba_nlp_gte.py

@@ -970,8 +970,9 @@ class GteForSequenceClassification(GtePreTrainedModel):
     def __init__(self, config: GteConfig):
         super().__init__(config)
         self.config = config
-        self.num_labels = 1
-        self.model = GteModel(config, add_pooling_layer=True)
+        self.num_labels = config.num_labels
+        assert config.num_labels > 0, "num_labels should be greater than 0 for sequence classification"
+        self.model = GteModel(config, add_pooling_layer=config.add_pooling_layer)
 
         self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
         self.loss_function = nn.MSELoss()
@@ -1010,7 +1011,10 @@ class GteForSequenceClassification(GtePreTrainedModel):
             output_attentions=output_attentions,
             output_hidden_states=output_hidden_states,
         )
-        hidden_states = transformer_outputs.pooler_output
+        if self.config.add_pooling_layer:
+            hidden_states = transformer_outputs.pooler_output
+        else:
+            hidden_states = transformer_outputs.last_hidden_state[:, 0]
 
         logits = self.score(hidden_states)