ver1.0

configuration_bert_concat.py

@@ -0,0 +1,7 @@
+from transformers import BertConfig
+
+class BertConcatConfig(BertConfig):
+    def __init__(self, bert_model_name='klue/bert-base', num_labels=2, **kwargs):
+        super().__init__(**kwargs)
+        self.bert_model_name = bert_model_name
+        self.num_labels = num_labels

modeling_bert_concat.py

@@ -0,0 +1,51 @@
+from transformers import PreTrainedModel, AutoModel
+from transformers.modeling_outputs import SequenceClassifierOutput
+import torch
+import torch.nn as nn
+
+class BertConcatClassifier(PreTrainedModel):
+    config_class = BertConcatConfig
+    
+    def __init__(self, config):
+        super().__init__(config)
+        self.bert = AutoModel.from_pretrained(config.bert_model_name, output_hidden_states=True)
+        self.num_labels = config.num_labels
+
+        # Classification layers
+        self.conv = nn.Conv1d(in_channels=3, out_channels=1, kernel_size=1)  # 3x768 -> 1x768
+        self.relu = nn.ReLU()
+        self.classifier = nn.Linear(768, self.num_labels)
+
+    def forward(self, input_ids, attention_mask=None, labels=None):
+        outputs = self.bert(input_ids, attention_mask=attention_mask)
+        hidden_states = outputs.hidden_states
+
+        # Concatenate the vectors as per custom model design
+        last_cls_vector = hidden_states[-1][:, 0, :]
+        fourth_last_cls_vector = hidden_states[-4][:, 0, :]
+        mean_pooled_vector = hidden_states[-1].mean(dim=1)
+
+        concatenated_vector = torch.cat(
+            (last_cls_vector.unsqueeze(1),
+             fourth_last_cls_vector.unsqueeze(1),
+             mean_pooled_vector.unsqueeze(1)),
+            dim=1
+        )
+
+        # Apply convolution and linear layers
+        conv_output = self.conv(concatenated_vector).squeeze(2)
+        relu_output = self.relu(conv_output)
+        logits = self.classifier(relu_output)
+        logits = logits.squeeze(1)
+
+        loss = None
+        if labels is not None:
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+
+        return SequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions
+        )