Attention pooling

changelog.md

@@ -1,5 +1,12 @@
 # Changelog
 
+## Unreleased
+
+### Added
+
+- New `attention` pooling mode in `eds.span_pooler`
+- New `word_pooling_mode=False` in `eds.transformer` to allow returning the worpiece embeddings directly, instead of the mean-pooled word embeddings. At the moment, this only works with `eds.span_pooler` which can pool over wordpieces or words seamlessly.
+
 ## v0.18.0 (2025-09-02)
 
 📢 EDS-NLP will drop support for Python 3.7, 3.8 and 3.9 support in the next major release (v0.19.0), in October 2025. Please upgrade to Python 3.10 or later.
@@ -13,6 +20,7 @@
 - New `eds.explode` pipe that splits one document into multiple documents, one per span yielded by its `span_getter` parameter, each new document containing exactly that single span.
 - New `Training a span classifier` tutorial, and reorganized deep-learning docs
 - `ScheduledOptimizer` now warns when a parameter selector does not match any parameter.
+- New `attention` pooling mode in `eds.span_pooler`
 
 ### Fixed
 

docs/tutorials/index.md

@@ -4,6 +4,7 @@ We provide step-by-step guides to get you started. We cover the following use-ca
 
 ### Base tutorials
 
+<!-- --8<-- [start:tutorials] -->
 <!-- --8<-- [start:classic-tutorials] -->
 
 === card {: href=/tutorials/spacy101 }
@@ -85,6 +86,8 @@ We provide step-by-step guides to get you started. We cover the following use-ca
     ---
     Quickly visualize the results of your pipeline as annotations or tables.
 
+<!-- --8<-- [end:classic-tutorials] -->
+
 ### Deep learning tutorials
 
 We also provide tutorials on how to train deep-learning models with EDS-NLP. These tutorials cover the training API, hyperparameter tuning, and more.
@@ -123,8 +126,5 @@ We also provide tutorials on how to train deep-learning models with EDS-NLP. The
     ---
     Learn how to tune hyperparameters of a model with `edsnlp.tune`.
 
-
 <!-- --8<-- [end:deep-learning-tutorials] -->
-
-
 <!-- --8<-- [end:tutorials] -->

edsnlp/core/torch_component.py

@@ -339,7 +339,14 @@ def compute_training_metrics(
         This is useful to compute averages when doing multi-gpu training or mini-batch
         accumulation since full denominators are not known during the forward pass.
         """
-        return batch_output
+        return (
+            {
+                **batch_output,
+                "loss": batch_output["loss"] / count,
+            }
+            if "loss" in batch_output
+            else batch_output
+        )
 
     def module_forward(self, *args, **kwargs):  # pragma: no cover
         """
@@ -348,6 +355,31 @@ def module_forward(self, *args, **kwargs):  # pragma: no cover
         """
         return torch.nn.Module.__call__(self, *args, **kwargs)
 
+    def preprocess_batch(self, docs: Sequence[Doc], supervision=False, **kwargs):
+        """
+        Convenience method to preprocess a batch of documents.
+        Features corresponding to the same path are grouped together in a list,
+        under the same key.
+
+        Parameters
+        ----------
+        docs: Sequence[Doc]
+            Batch of documents
+        supervision: bool
+            Whether to extract supervision features or not
+
+        Returns
+        -------
+        Dict[str, Sequence[Any]]
+            The batch of features
+        """
+        batch = [
+            (self.preprocess_supervised(d) if supervision else self.preprocess(d))
+            for d in docs
+        ]
+        batch = decompress_dict(list(batch_compress_dict(batch)))
+        return batch
+
     def prepare_batch(
         self,
         docs: Sequence[Doc],
@@ -372,11 +404,7 @@ def prepare_batch(
         -------
         Dict[str, Sequence[Any]]
         """
-        batch = [
-            (self.preprocess_supervised(doc) if supervision else self.preprocess(doc))
-            for doc in docs
-        ]
-        batch = decompress_dict(list(batch_compress_dict(batch)))
+        batch = self.preprocess_batch(docs, supervision=supervision)
         batch = self.collate(batch)
         batch = self.batch_to_device(batch, device=device)
         return batch