new config and new MCQBert classes

config.json

@@ -1,8 +1,12 @@
 {
   "_name_or_path": "epfl-ml4ed/MCQBert",
   "architectures": [
-    "BertModel"
+    "MCQBert"
   ],
+  "auto_map": {
+    "AutoConfig": "configuration_mcqbert.MCQBertConfig",
+    "AutoModelForCausalLM": "modeling_mcqbert.MCQBert"
+  },
   "attention_probs_dropout_prob": 0.1,
   "classifier_dropout": null,
   "cls_hidden_size": 256,
@@ -10,6 +14,7 @@
   "hidden_dropout_prob": 0.1,
   "hidden_size": 768,
   "initializer_range": 0.02,
+  "integration_strategy": null,
   "intermediate_size": 3072,
   "layer_norm_eps": 1e-12,
   "max_position_embeddings": 512,

configuration_mcqbert.py

@@ -0,0 +1,10 @@
+from transformers import BertConfig
+
+class MCQBertConfig(BertConfig):
+    model_type = "mcqbert"
+
+    def __init__(self, integration_strategy=None, student_embedding_size=4096, cls_hidden_size=256, **kwargs):
+        super().__init__(**kwargs)
+        self.integration_strategy = integration_strategy
+        self.student_embedding_size = student_embedding_size
+        self.cls_hidden_size = cls_hidden_size

modeling_mcqbert.py

@@ -0,0 +1,39 @@
+class MCQBert(BertPreTrainedModel):
+    def __init__(self, config: MCQBertConfig):
+        super().__init__(config)
+        self.student_embedding_layer = torch.nn.Linear(config.student_embedding_size, config.hidden_size)
+        
+        cls_input_dim_multiplier = 2 if config.integration_strategy == "cat" else 1
+        cls_input_dim = self.config.hidden_size * cls_input_dim_multiplier
+        
+        self.classifier = torch.nn.Sequential(
+            torch.nn.Linear(cls_input_dim, config.cls_hidden_size),
+            torch.nn.ReLU(),
+            torch.nn.Linear(config.cls_hidden_size, 1)
+        )
+        
+    def forward(self, input_ids, student_embeddings=None):
+        if self.config.integration_strategy is None:
+            # don't consider embeddings is no integration strategy (MCQBert)
+            student_embeddings = torch.zeros(self.config.student_embedding_layer)
+
+            input_embeddings = self.embeddings(input_ids)
+            combined_embeddings = input_embeddings + self.student_embedding_layer(student_embeddings).unsqueeze(1).repeat(1, input_embeddings.size(1), 1)
+            output = super().forward(inputs_embeds = combined_embeddings)
+            return self.classifier(output.last_hidden_state[:, 0, :])
+
+        elif self.config.integration_strategy == "cat":
+            # MCQStudentBertCat
+            output = super().forward(input_ids)
+            output_with_student_embedding = torch.cat((output.last_hidden_state[:, 0, :], self.student_embedding_layer(student_embeddings)), dim = 1)
+            return self.classifier(output_with_student_embedding)
+
+        elif self.config.integration_strategy == "sum":
+            # MCQStudentBertSum
+            input_embeddings = self.embeddings(input_ids)
+            combined_embeddings = input_embeddings + self.student_embedding_layer(student_embeddings).unsqueeze(1).repeat(1, input_embeddings.size(1), 1)
+            output = super().forward(inputs_embeds = combined_embeddings)
+            return self.classifier(output.last_hidden_state[:, 0, :])
+
+        else:
+            raise ValueError(f"{self.config.integration_strategy} is not a known integration_strategy")