move classif head model logic to doc_classifier instead of doc_pooler

Author: LucasDedieu

edsnlp/pipes/trainable/doc_classifier/doc_classifier.py

@@ -3,8 +3,9 @@
 from typing import Any, Dict, Iterable, Optional, Sequence, Set, Union
 
 import torch
+import torch.nn as nn
 from spacy.tokens import Doc
-from typing_extensions import NotRequired, TypedDict
+from typing_extensions import Literal, NotRequired, TypedDict
 
 from edsnlp.core.pipeline import PipelineProtocol
 from edsnlp.core.torch_component import BatchInput, TorchComponent
@@ -36,6 +37,8 @@ class TrainableDocClassifier(
     TorchComponent[DocClassifierBatchOutput, DocClassifierBatchInput],
     BaseComponent,
 ):
+    """A trainable document classifier that uses embeddings to classify documents."""
+
     def __init__(
         self,
         nlp: Optional[PipelineProtocol] = None,
@@ -49,12 +52,21 @@ def __init__(
         loss_fn=None,
         labels: Optional[Sequence[str]] = None,
         class_weights: Optional[Union[Dict[str, float], str]] = None,
+        hidden_size: Optional[int] = None,
+        activation_mode: Literal["relu", "gelu", "silu"] = "relu",
+        dropout_rate: Optional[float] = 0.0,
+        layer_norm: Optional[bool] = False,
     ):
+        self.num_classes = num_classes
         self.label_attr: Attributes = label_attr
         self.label2id = label2id or {}
         self.id2label = id2label or {}
         self.labels = labels
         self.class_weights = class_weights
+        self.hidden_size = hidden_size
+        self.activation_mode = activation_mode
+        self.dropout_rate = dropout_rate
+        self.layer_norm = layer_norm
 
         super().__init__(nlp, name)
         self.embedding = embedding
@@ -66,9 +78,23 @@ def __init__(
             raise ValueError(
                 "The embedding component must have an 'output_size' attribute."
             )
-        embedding_size = self.embedding.output_size
-        if num_classes:
-            self.classifier = torch.nn.Linear(embedding_size, num_classes)
+        self.embedding_size = self.embedding.output_size
+        if self.num_classes:
+            self.build_classifier()
+
+    def build_classifier(self):
+        """Build classification head"""
+        if self.hidden_size:
+            self.hidden_layer = torch.nn.Linear(self.embedding_size, self.hidden_size)
+            self.activation = {"relu": nn.ReLU(), "gelu": nn.GELU(), "silu": nn.SiLU()}[
+                self.activation_mode
+            ]
+            if self.layer_norm:
+                self.norm = nn.LayerNorm(self.hidden_size)
+            self.dropout = nn.Dropout(self.dropout_rate)
+            self.classifier = torch.nn.Linear(self.hidden_size, self.num_classes)
+        else:
+            self.classifier = torch.nn.Linear(self.embedding_size, self.num_classes)
 
     def _compute_class_weights(self, freq_dict: Dict[str, int]) -> torch.Tensor:
         """
@@ -112,10 +138,9 @@ def post_init(self, gold_data: Iterable[Doc], exclude: Set[str]):
                 for i, label in enumerate(labels):
                     self.label2id[label] = i
                     self.id2label[i] = label
-                print("num classes:", len(self.label2id))
-                self.classifier = torch.nn.Linear(
-                    self.embedding.output_size, len(self.label2id)
-                )
+                self.num_classes = len(self.label2id)
+                print("num classes:", self.num_classes)
+                self.build_classifier()
 
         weight_tensor = None
         if self.class_weights is not None:
@@ -138,6 +163,7 @@ def preprocess(self, doc: Doc) -> Dict[str, Any]:
         return {"embedding": self.embedding.preprocess(doc)}
 
     def preprocess_supervised(self, doc: Doc) -> Dict[str, Any]:
+        """Preprocess document with target labels for training."""
         preps = self.preprocess(doc)
         label = getattr(doc._, self.label_attr, None)
         if label is None:
@@ -166,9 +192,14 @@ def forward(self, batch: DocClassifierBatchInput) -> DocClassifierBatchOutput:
         if targets provided.
         """
         pooled = self.embedding(batch["embedding"])
-        embeddings = pooled["embeddings"]
-
-        logits = self.classifier(embeddings)
+        x = pooled["embeddings"]
+        if self.hidden_size:
+            x = self.hidden_layer(x)
+            x = self.activation(x)
+            if self.layer_norm:
+                x = self.norm(x)
+            x = self.dropout(x)
+        logits = self.classifier(x)
 
         output: DocClassifierBatchOutput = {}
         if "targets" in batch:
@@ -181,6 +212,7 @@ def forward(self, batch: DocClassifierBatchInput) -> DocClassifierBatchOutput:
         return output
 
     def postprocess(self, docs, results, input):
+        """Postprocess predictions by assigning labels to documents."""
         labels = results["labels"]
         if isinstance(labels, torch.Tensor):
             labels = labels.tolist()

edsnlp/pipes/trainable/embeddings/doc_pooler/doc_pooler.py

@@ -52,20 +52,11 @@ def __init__(
         *,
         embedding: WordEmbeddingComponent,
         pooling_mode: Literal["max", "sum", "mean", "cls"] = "mean",
-        hidden_size: Optional[int] = None,
     ):
         super().__init__(nlp, name)
         self.embedding = embedding
         self.pooling_mode = pooling_mode
-        self.output_size = embedding.output_size if hidden_size is None else hidden_size
-        self.projector = (
-            torch.nn.Linear(self.embedding.output_size, hidden_size)
-            if hidden_size is not None
-            else torch.nn.Identity()
-        )
-
-    def feed_forward(self, doc_embed: torch.Tensor) -> torch.Tensor:
-        return self.projector(doc_embed)
+        self.output_size = embedding.output_size
 
     def preprocess(self, doc: Doc, **kwargs) -> Dict[str, Any]:
         embedding_out = self.embedding.preprocess(doc, **kwargs)
@@ -85,21 +76,26 @@ def collate(self, batch: Dict[str, Any]) -> DocPoolerBatchInput:
         }
 
     def forward(self, batch: DocPoolerBatchInput) -> DocPoolerBatchOutput:
-        device = next(self.parameters()).device
-
         embeds = self.embedding(batch["embedding"])["embeddings"]
         device = embeds.device
 
+        if self.pooling_mode == "cls":
+            pooled = self.embedding(batch["embedding"])["cls"].to(device)
+            return {"embeddings": pooled}
+
+        mask = embeds.mask
+        mask_expanded = mask.unsqueeze(-1)
+        masked_embeds = embeds * mask_expanded
+        sum_embeds = masked_embeds.sum(dim=1)
         if self.pooling_mode == "mean":
-            pooled = embeds.mean(dim=1)
+            valid_counts = mask.sum(dim=1, keepdim=True).clamp(min=1)
+            pooled = sum_embeds / valid_counts
         elif self.pooling_mode == "max":
-            pooled = embeds.max(dim=1).values
+            masked_embeds = embeds.masked_fill(~mask_expanded, float("-inf"))
+            pooled, _ = masked_embeds.max(dim=1)
         elif self.pooling_mode == "sum":
-            pooled = embeds.sum(dim=1) / embeds.size(1)
-        elif self.pooling_mode == "cls":
-            pooled = self.embedding(batch["embedding"])["cls"].to(device)
+            pooled = sum_embeds
         else:
             raise ValueError(f"Unknown pooling mode: {self.pooling_mode}")
 
-        pooled = self.feed_forward(pooled)
         return {"embeddings": pooled}