[models] - fix: model contract issues

bert_squeeze/assistants/configs/train_adapter.yaml

@@ -5,6 +5,7 @@ general:
   get_mismatched: true
   evaluate_during_training: true
   labels: [ 0, 1 ]
+  num_labels: 2
   output_dir: outputs
   save_steps: 500
   validation_every_n_epoch: 1
@@ -34,7 +35,7 @@ model:
   training_config: ${train}
   task_name:
   adapter_config_name: "seq_bn"
-  labels: [ "0", "1" ]
+  labels: ${general.labels}
   scorer:
     _target_: bert_squeeze.utils.scorers.sequence_classification_scorer.BaseSequenceClassificationScorer
     labels: ${general.labels}

bert_squeeze/assistants/configs/train_bert.yaml

@@ -5,6 +5,7 @@ general:
   get_mismatched: true
   evaluate_during_training: true
   labels: [ 0, 1 ]
+  num_labels: 2
   output_dir: outputs
   save_steps: 500
   validation_every_n_epoch: 1
@@ -30,7 +31,7 @@ train:
 
 model:
   _target_: bert_squeeze.models.lt_bert.LtSequenceClassificationCustomBert
-  num_labels: 2
+  num_labels: ${general.num_labels}
   pretrained_model: "bert-base-cased"
   training_config: ${train}
   scorer:

bert_squeeze/assistants/configs/train_fastbert.yaml

@@ -5,6 +5,7 @@ general:
   get_mismatched: true
   evaluate_during_training: true
   labels: [ 0, 1 ]
+  num_labels: 2
   output_dir: outputs
   save_steps: 500
   validation_every_n_epoch: 1
@@ -35,7 +36,7 @@ model:
   _target_: bert_squeeze.models.lt_fastbert.LtFastBert
   training_config: ${train}
   pretrained_model: "bert-base-cased"
-  num_labels: 2
+  num_labels: ${general.num_labels}
   scorer_type: "fast"
   scorer:
     _target_: bert_squeeze.utils.scorers.sequence_classification_scorer.FastBertSequenceClassificationScorer
@@ -51,4 +52,4 @@ data:
     label_col: label
     truncate_mode: head
   tokenizer_name: ${model.pretrained_model}
-  max_length: 256
+  max_length: 256

bert_squeeze/assistants/configs/train_theseus_bert.yaml

@@ -5,6 +5,7 @@ general:
   get_mismatched: true
   evaluate_during_training: true
   labels: [ 0, 1 ]
+  num_labels: 2
   output_dir: outputs
   save_steps: 500
   validation_every_n_epoch: 1
@@ -32,7 +33,7 @@ model:
   _target_: bert_squeeze.models.lt_theseus_bert.LtTheseusBert
   training_config: ${train}
   pretrained_model: "bert-base-cased"
-  num_labels: 2
+  num_labels: ${general.num_labels}
   replacement_scheduler:
     type: "linear"
     base_replacing_rate: 0.3

bert_squeeze/assistants/train_assistant.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from copy import deepcopy
 from importlib import resources
 from typing import Dict, List, Optional
@@ -72,8 +74,11 @@ def __init__(
                 f" following: {CONFIG_MAPPER.keys()}"
             )
 
-        config_path = resources.files("bert_squeeze").joinpath(
-            "assistants/configs", config_name
+        config_path = (
+            resources.files("bert_squeeze")
+            .joinpath("assistants")
+            .joinpath("configs")
+            .joinpath(config_name)
         )
         with resources.as_file(config_path) as resolved_path:
             conf = OmegaConf.load(resolved_path)
@@ -102,6 +107,14 @@ def __init__(
                     overrides if base is None else deep_update(base, overrides)
                 )
 
+        labels = conf["general"].get("labels")
+        if labels is not None:
+            num_labels = len(labels)
+            configured_num_labels = conf["general"].get("num_labels")
+            if configured_num_labels is not None and configured_num_labels != num_labels:
+                raise ValueError("general.num_labels must match the number of labels.")
+            conf["general"]["num_labels"] = num_labels
+
         self.name = name
         self.general = conf["general"]
         self.train = conf["train"]
@@ -112,8 +125,8 @@ def __init__(
 
         self._model: Optional[pl.LightningModule] = None
         self._data: Optional[pl.LightningDataModule] = None
-        self._logger = None
-        self._callbacks = None
+        self._logger: Optional[Logger] = None
+        self._callbacks: Optional[List[Callback]] = None
 
     @property
     def model(self) -> pl.LightningModule:
@@ -162,11 +175,11 @@ def callbacks(self) -> List[Callback]:
         """"""
         if self._callbacks is None:
             if self._callbacks_conf is not None:
-                self.callbacks = [
+                self._callbacks = [
                     instantiate(callback) for callback in self._callbacks_conf
                 ]
             else:
-                self.callbacks = []
+                self._callbacks = []
         return self._callbacks
 
     @callbacks.setter

bert_squeeze/models/custom_transformers/deebert.py

@@ -1,7 +1,9 @@
 # This is heavily inspired by the following repo:
 # https://github.com/castorini/DeeBERT
+from __future__ import annotations
+
 from abc import ABC
-from typing import List, Union
+from typing import List, Optional, Tuple, Union
 
 import torch
 import torch.nn as nn
@@ -66,15 +68,17 @@ class DeeBertEncoder(nn.Module):
     def __init__(self, config: PretrainedConfig, inference: bool):
         super(DeeBertEncoder, self).__init__()
         self.config = config
-        self.layer = nn.ModuleList([BertLayer(config)] * config.num_hidden_layers)
-        self.ramp = nn.ModuleList([OffRamp(config)] * config.num_hidden_layers)
+        self.layer = nn.ModuleList(
+            [BertLayer(config) for _ in range(config.num_hidden_layers)]
+        )
+        self.ramp = nn.ModuleList(
+            [OffRamp(config) for _ in range(config.num_hidden_layers)]
+        )
 
-        self.early_exit_entropy = [
-            -1,
-        ] * config.num_hidden_layers
+        self.early_exit_entropy: List[float] = [-1.0] * config.num_hidden_layers
         self.inference = inference
 
-    def set_early_exit_entropy(self, x: Union[List[float], float]) -> None:
+    def set_early_exit_entropy(self, x: Union[List[float], float, int]) -> None:
         """
         Assigning an entropy threshold to every layer.
 
@@ -85,9 +89,9 @@ def set_early_exit_entropy(self, x: Union[List[float], float]) -> None:
         """
         if isinstance(x, float) or isinstance(x, int):
             for i in range(self.config.num_hidden_layers):
-                self.early_exit_entropy[i] = x
+                self.early_exit_entropy[i] = float(x)
         elif isinstance(x, list):
-            self.early_exit_entropy = x
+            self.early_exit_entropy = list(x)
         else:
             raise TypeError(
                 f"Expected 'x' to be of type 'float' or 'list' but got :'{type(x)}'"
@@ -117,11 +121,11 @@ def forward(
         output_hidden_states: bool = False,
     ) -> DeeBertEncoderOutput:
         """"""
-        all_hidden_states = tuple() if output_hidden_states else None
-        all_attentions = tuple() if output_attentions else None
+        all_hidden_states: Tuple[torch.Tensor, ...] = tuple()
+        all_attentions: Tuple[torch.Tensor, ...] = tuple()
 
         if not self.inference:
-            all_ramps = tuple()
+            all_ramps: Tuple[RampOutput, ...] = tuple()
 
             for i, layer_module in enumerate(self.layer):
                 if output_hidden_states:
@@ -149,16 +153,18 @@ def forward(
 
             return DeeBertEncoderOutput(
                 last_hidden_state=hidden_states,
-                hidden_states=all_hidden_states,
-                attentions=all_attentions,
+                hidden_states=all_hidden_states if output_hidden_states else None,
+                attentions=all_attentions if output_attentions else None,
                 ramps_exit=all_ramps,
                 exit_layer=i,
             )
         else:
-            all_ramps = [
-                0,
-            ] * hidden_states.shape[0]
-            positions = torch.arange(start=0, end=hidden_states.shape[0]).long()
+            batch_ramps: List[Optional[RampOutput]] = [None] * hidden_states.shape[0]
+            positions = torch.arange(
+                start=0,
+                end=hidden_states.shape[0],
+                device=hidden_states.device,
+            ).long()
 
             for i, layer_module in enumerate(self.layer):
                 layer_outputs = layer_module(
@@ -174,21 +180,35 @@ def forward(
 
                 if i == len(self.layer) - 1:
                     for idx, pos in enumerate(positions):
-                        all_ramps[pos] = ramp_exit[idx]
+                        batch_ramps[int(pos)] = ramp_exit[idx]
                 else:
                     enough_info = ramp_exit.entropy < self.early_exit_entropy[i]
                     right_pos = positions[enough_info]
 
                     for idx, pos in enumerate(right_pos):
-                        all_ramps[pos] = ramp_exit[idx]
+                        batch_ramps[int(pos)] = ramp_exit[idx]
 
                     hidden_states = hidden_states[~enough_info]
                     attention_mask = attention_mask[~enough_info]
                     positions = positions[~enough_info]
 
                     if positions.nelement() == 0:
-                        return DeeBertEncoderOutput(ramps_exit=all_ramps, exit_layer=i)
-            return DeeBertEncoderOutput(ramps_exit=all_ramps, exit_layer=i)
+                        return DeeBertEncoderOutput(
+                            ramps_exit=self._completed_ramps(batch_ramps),
+                            exit_layer=i,
+                        )
+            return DeeBertEncoderOutput(
+                ramps_exit=self._completed_ramps(batch_ramps),
+                exit_layer=i,
+            )
+
+    @staticmethod
+    def _completed_ramps(
+        ramps: List[Optional[RampOutput]],
+    ) -> Tuple[RampOutput, ...]:
+        if any(ramp is None for ramp in ramps):
+            raise RuntimeError("DeeBERT did not produce an output for every sample.")
+        return tuple(ramp for ramp in ramps if ramp is not None)
 
 
 class DeeBertModel(BertPreTrainedModel, ABC):

bert_squeeze/models/custom_transformers/fastbert.py

@@ -2,6 +2,8 @@
 # The main difference relies on the fact that I'm trying to use HuggingFace's
 # 'transformers' components as much as possible.
 
+from __future__ import annotations
+
 from typing import List, Tuple, Union
 
 import torch
@@ -129,7 +131,9 @@ def forward(
         if inference:
             # positions will keep track of the original position of each element in the
             # batch when elements will be removed
-            final_probs = torch.zeros((hidden_states[0].shape[0], 2), device=device)
+            final_probs = hidden_states[0].new_zeros(
+                (hidden_states[0].shape[0], self.config.num_labels)
+            )
             positions = torch.arange(
                 start=0, end=hidden_states[0].shape[0], device=device
             ).long()