feat: create SequenceEmbedder as a new component

- moved LabelAttention as one SentenceEmbedder type
- adapt tests

Author: meilame-tayebjee

tests/conftest.py

@@ -63,6 +63,7 @@ def mock_tokenizer():
     tokenizer = Mock()
     tokenizer.vocab_size = 1000
     tokenizer.padding_idx = 1
+    tokenizer.output_vectorized = False
     tokenizer.tokenize = Mock(
         return_value={
             "input_ids": np.array([[1, 2, 3], [4, 5, 6]]),

tests/test_pipeline.py

@@ -11,8 +11,10 @@
     CategoricalVariableNet,
     ClassificationHead,
     LabelAttentionConfig,
-    TextEmbedder,
-    TextEmbedderConfig,
+    SentenceEmbedder,
+    SentenceEmbedderConfig,
+    TokenEmbedder,
+    TokenEmbedderConfig,
 )
 from torchTextClassifiers.tokenizers import NGramTokenizer
 from torchTextClassifiers.value_encoder import DictEncoder, ValueEncoder
@@ -122,24 +124,29 @@ def run_full_pipeline(
         sequence_len=sequence_len,
     )
 
-    # Create text embedder
-    text_embedder_config = TextEmbedderConfig(
+    # Create token embedder
+    token_embedder_config = TokenEmbedderConfig(
         vocab_size=vocab_size,
         embedding_dim=model_params["embedding_dim"],
         padding_idx=padding_idx,
         attention_config=attention_config,
+    )
+    token_embedder = TokenEmbedder(token_embedder_config=token_embedder_config)
+
+    # Create sentence embedder
+    sentence_embedder_config = SentenceEmbedderConfig(
         label_attention_config=(
             LabelAttentionConfig(
                 n_head=attention_config.n_head,
                 num_classes=num_classes,
+                embedding_dim=model_params["embedding_dim"],
             )
             if label_attention_enabled
             else None
         ),
+        aggregation_method=None if label_attention_enabled else "mean",
     )
-
-    text_embedder = TextEmbedder(text_embedder_config=text_embedder_config)
-    text_embedder.init_weights()
+    sentence_embedder = SentenceEmbedder(sentence_embedder_config=sentence_embedder_config)
 
     # Create categorical variable net (vocab sizes from fitted encoder)
     categorical_var_net = CategoricalVariableNet(
@@ -156,7 +163,8 @@ def run_full_pipeline(
 
     # Create model
     model = TextClassificationModel(
-        text_embedder=text_embedder,
+        token_embedder=token_embedder,
+        sentence_embedder=sentence_embedder,
         categorical_variable_net=categorical_var_net,
         classification_head=classification_head,
     )

torchTextClassifiers/model/components/text_embedder.py

@@ -46,6 +46,9 @@ class TokenEmbedder(nn.Module):
 
     """
 
+    cos: torch.Tensor
+    sin: torch.Tensor
+
     def __init__(self, token_embedder_config: TokenEmbedderConfig):
         super().__init__()
 
@@ -112,10 +115,7 @@ def init_weights(self):
             for block in self.transformer.h:
                 torch.nn.init.zeros_(block.mlp.c_proj.weight)
                 torch.nn.init.zeros_(block.attn.c_proj.weight)
-            # init the rotary embeddings
-            head_dim = self.attention_config.n_embd // self.attention_config.n_head
-            cos, sin = self._precompute_rotary_embeddings(self.rotary_seq_len, head_dim)
-            self.cos, self.sin = cos, sin
+
         # Cast the embeddings from fp32 to bf16: optim can tolerate it and it saves memory: both in the model and the activations
         if self.embedding_layer.weight.device.type == "cuda":
             self.embedding_layer.to(dtype=torch.bfloat16)
@@ -205,9 +205,7 @@ def __init__(self, label_attention_config: LabelAttentionConfig):
         super().__init__()
 
         if label_attention_config is None:
-            raise ValueError(
-                "label_attention_config must be provided to use LabelAttention."
-            )
+            raise ValueError("label_attention_config must be provided to use LabelAttention.")
 
         self.label_attention_config = label_attention_config
         self.num_classes = label_attention_config.num_classes
@@ -311,7 +309,7 @@ def forward(
 
 class SentenceEmbedder(nn.Module):
     def __init__(self, sentence_embedder_config: SentenceEmbedderConfig):
-
+        super().__init__()
         """
         A module to aggregate token embeddings.
 
@@ -322,7 +320,7 @@ def __init__(self, sentence_embedder_config: SentenceEmbedderConfig):
         - aggregation_method=None: in that case you need to provide a label attention
         """
 
-        self.config
+        self.config = sentence_embedder_config
         self.label_attention_config = sentence_embedder_config.label_attention_config
         self.aggregation_method = sentence_embedder_config.aggregation_method
 

torchTextClassifiers/model/model.py

@@ -15,8 +15,8 @@
     CategoricalForwardType,
     CategoricalVariableNet,
     ClassificationHead,
+    SentenceEmbedder,
     TokenEmbedder,
-    SentenceEmbedder
 )
 from torchTextClassifiers.model.components.attention import norm
 
@@ -53,20 +53,22 @@ def __init__(
             classification_head (ClassificationHead): The classification head module.
             token_embedder (Optional[TextEmbedder]): The text embedding module.
                 If not provided, assumes that input text is already embedded (as tensors) and directly passed to the classification head.
-            sentence_embedder: 
+            sentence_embedder:
             categorical_variable_net (Optional[CategoricalVariableNet]): The categorical variable network module.
                 If not provided, assumes no categorical variables are used.
         """
         super().__init__()
 
         self.token_embedder = token_embedder
+        self.sentence_embedder = sentence_embedder
 
         if self.token_embedder is not None:
             self.token_embedder.init_weights()
             if self.sentence_embedder is None:
-                raise ValueError("You have provided a TokenEmbedder but no SentenceEmbedder: please provide one.")
-            else:
-                self.sentence_embedder = sentence_embedder
+                raise ValueError(
+                    "You have provided a TokenEmbedder but no SentenceEmbedder: please provide one."
+                )
+
         self.categorical_variable_net = categorical_variable_net
         if not self.categorical_variable_net:
             logger.info("🔹 No categorical variable network provided; using only text embeddings.")
@@ -76,7 +78,6 @@ def __init__(
         self._validate_component_connections()
 
         torch.nn.init.zeros_(self.classification_head.net.weight)
-            
 
     def _validate_component_connections(self):
         def _check_text_categorical_connection(self, token_embedder, cat_var_net):
@@ -158,7 +159,9 @@ def forward(
                 attention_mask=attention_mask,
             )
             x_token = token_embed_output["token_embeddings"]
-            sentence_embedding_output = self.sentence_embedder(x_token, attention_mask, return_label_attention_matrix=return_label_attention_matrix)
+            sentence_embedding_output = self.sentence_embedder(
+                x_token, attention_mask, return_label_attention_matrix=return_label_attention_matrix
+            )
             x_text = sentence_embedding_output["sentence_embedding"]
             if return_label_attention_matrix:
                 label_attention_matrix = sentence_embedding_output["label_attention_matrix"]