[WIP] add TrainableDocClassifier

Author: LucasDedieu

edsnlp/metrics/doc_classif.py

@@ -0,0 +1,106 @@
+from typing import Any, Dict, Iterable, Optional, Tuple, Union
+
+from spacy.tokens import Doc
+from spacy.training import Example
+
+from edsnlp import registry
+from edsnlp.metrics import make_examples
+
+
+def doc_classification_metric(
+    examples: Union[Tuple[Iterable[Doc], Iterable[Doc]], Iterable[Example]],
+    label_attr: str = "label",
+    micro_key: str = "micro",
+    filter_expr: Optional[str] = None,
+) -> Dict[str, Any]:
+    """
+    Scores document-level classification (accuracy, precision, recall, F1).
+
+    Parameters
+    ----------
+    examples: Examples
+        The examples to score, either a tuple of (golds, preds) or a list of
+        spacy.training.Example objects
+    label_attr: str
+        The Doc._ attribute containing the label
+    micro_key: str
+        The key to use to store the micro-averaged results
+    filter_expr: str
+        The filter expression to use to filter the documents
+
+    Returns
+    -------
+    Dict[str, Any]
+    """
+    examples = make_examples(examples)
+    if filter_expr is not None:
+        filter_fn = eval(f"lambda doc: {filter_expr}")
+        examples = [eg for eg in examples if filter_fn(eg.reference)]
+
+    pred_labels = []
+    gold_labels = []
+    for eg in examples:
+        pred = getattr(eg.predicted._, label_attr, None)
+        gold = getattr(eg.reference._, label_attr, None)
+        pred_labels.append(pred)
+        gold_labels.append(gold)
+
+    print(pred_labels, gold_labels)
+
+    labels = set(gold_labels) | set(pred_labels)
+    results = {}
+    for label in labels:
+        pred_set = [i for i, p in enumerate(pred_labels) if p == label]
+        gold_set = [i for i, g in enumerate(gold_labels) if g == label]
+        tp = len(set(pred_set) & set(gold_set))
+        num_pred = len(pred_set)
+        num_gold = len(gold_set)
+        results[label] = {
+            "f": 2 * tp / max(1, num_pred + num_gold),
+            "p": 1 if tp == num_pred else (tp / num_pred) if num_pred else 0.0,
+            "r": 1 if tp == num_gold else (tp / num_gold) if num_gold else 0.0,
+            "tp": tp,
+            "support": num_gold,
+            "positives": num_pred,
+        }
+
+    tp = sum(1 for p, g in zip(pred_labels, gold_labels) if p == g)
+    num_pred = len(pred_labels)
+    num_gold = len(gold_labels)
+    results[micro_key] = {
+        "accuracy": tp / num_gold if num_gold else 0.0,
+        "f": 2 * tp / max(1, num_pred + num_gold),
+        "p": tp / num_pred if num_pred else 0.0,
+        "r": tp / num_gold if num_gold else 0.0,
+        "tp": tp,
+        "support": num_gold,
+        "positives": num_pred,
+    }
+    return results
+
+
+@registry.metrics.register("eds.doc_classification")
+class DocClassificationMetric:
+    def __init__(
+        self,
+        label_attr: str = "label",
+        micro_key: str = "micro",
+        filter_expr: Optional[str] = None,
+    ):
+        self.label_attr = label_attr
+        self.micro_key = micro_key
+        self.filter_expr = filter_expr
+
+    def __call__(self, *examples):
+        return doc_classification_metric(
+            examples,
+            label_attr=self.label_attr,
+            micro_key=self.micro_key,
+            filter_expr=self.filter_expr,
+        )
+
+
+__all__ = [
+    "doc_classification_metric",
+    "DocClassificationMetric",
+]

edsnlp/pipes/__init__.py

@@ -82,4 +82,6 @@
     from .trainable.embeddings.span_pooler.factory import create_component as span_pooler
     from .trainable.embeddings.transformer.factory import create_component as transformer
     from .trainable.embeddings.text_cnn.factory import create_component as text_cnn
+    from .trainable.embeddings.doc_pooler.factory import create_component as doc_pooler
+    from .trainable.doc_classifier.factory import create_component as doc_classifier
     from .misc.split import Split as split

edsnlp/pipes/trainable/doc_classifier/__init__.py

@@ -0,0 +1 @@
+from .factory import create_component

edsnlp/pipes/trainable/doc_classifier/doc_classifier.py

@@ -0,0 +1,125 @@
+import os
+import pickle
+from typing import Any, Dict, Iterable, Optional, Sequence, Set, Union
+
+import torch
+from spacy.tokens import Doc
+from typing_extensions import NotRequired, TypedDict
+
+from edsnlp.core.pipeline import PipelineProtocol
+from edsnlp.core.torch_component import BatchInput, TorchComponent
+from edsnlp.pipes.base import BaseComponent
+from edsnlp.pipes.trainable.embeddings.typing import (
+    WordContextualizerComponent,
+    WordEmbeddingComponent,
+)
+from edsnlp.utils.bindings import Attributes
+
+DocClassifierBatchInput = TypedDict(
+    "DocClassifierBatchInput",
+    {
+        "embedding": BatchInput,
+        "targets": NotRequired[torch.Tensor],
+    },
+)
+
+DocClassifierBatchOutput = TypedDict(
+    "DocClassifierBatchOutput",
+    {
+        "loss": Optional[torch.Tensor],
+        "labels": Optional[torch.Tensor],
+    },
+)
+
+
+class TrainableDocClassifier(
+    TorchComponent[DocClassifierBatchOutput, DocClassifierBatchInput],
+    BaseComponent,
+):
+    def __init__(
+        self,
+        nlp: Optional[PipelineProtocol] = None,
+        name: str = "doc_classifier",
+        *,
+        embedding: Union[WordEmbeddingComponent, WordContextualizerComponent],
+        num_classes: int,
+        label_attr: str = "label",
+        loss_fn=None,
+    ):
+        self.label_attr: Attributes = label_attr
+        super().__init__(nlp, name)
+        self.embedding = embedding
+        self.loss_fn = loss_fn or torch.nn.CrossEntropyLoss()
+
+        if not hasattr(self.embedding, "output_size"):
+            raise ValueError(
+                "The embedding component must have an 'output_size' attribute."
+            )
+        embedding_size = self.embedding.output_size
+        self.classifier = torch.nn.Linear(embedding_size, num_classes)
+
+    def set_extensions(self) -> None:
+        super().set_extensions()
+        if not Doc.has_extension(self.label_attr):
+            Doc.set_extension(self.label_attr, default={})
+
+    def post_init(self, gold_data: Iterable[Doc], exclude: Set[str]):
+        super().post_init(gold_data, exclude=exclude)
+
+    def preprocess(self, doc: Doc) -> Dict[str, Any]:
+        return {"embedding": self.embedding.preprocess(doc)}
+
+    def preprocess_supervised(self, doc: Doc) -> Dict[str, Any]:
+        preps = self.preprocess(doc)
+        label = getattr(doc._, self.label_attr, None)
+        if label is None:
+            raise ValueError(
+                f"Document does not have a gold label in 'doc._.{self.label_attr}'"
+            )
+        return {
+            **preps,
+            "targets": torch.tensor(label, dtype=torch.long),
+        }
+
+    def collate(self, batch: Dict[str, Sequence[Any]]) -> DocClassifierBatchInput:
+        embeddings = self.embedding.collate(batch["embedding"])
+        batch_input: DocClassifierBatchInput = {"embedding": embeddings}
+        if "targets" in batch:
+            batch_input["targets"] = torch.stack(batch["targets"])
+        return batch_input
+
+    def forward(self, batch: DocClassifierBatchInput) -> DocClassifierBatchOutput:
+        pooled = self.embedding(batch["embedding"])
+        embeddings = pooled["embeddings"]
+
+        logits = self.classifier(embeddings)
+
+        output: DocClassifierBatchOutput = {}
+        if "targets" in batch:
+            loss = self.loss_fn(logits, batch["targets"])
+            output["loss"] = loss
+            output["labels"] = None
+        else:
+            output["loss"] = None
+            output["labels"] = torch.argmax(logits, dim=-1)
+        return output
+
+    def postprocess(self, docs, results, input):
+        labels = results["labels"]
+        if isinstance(labels, torch.Tensor):
+            labels = labels.tolist()
+        for doc, label in zip(docs, labels):
+            setattr(doc._, self.label_attr, label)
+            # doc._.label = label
+        return docs
+
+    def to_disk(self, path, *, exclude=set()):
+        repr_id = object.__repr__(self)
+        if repr_id in exclude:
+            return
+        exclude.add(repr_id)
+        os.makedirs(path, exist_ok=True)
+        data_path = path / "label_attr.pkl"
+        with open(data_path, "wb") as f:
+            pickle.dump({"label_attr": self.label_attr}, f)
+        return super().to_disk(path, exclude=exclude)

edsnlp/pipes/trainable/doc_classifier/factory.py

@@ -0,0 +1,9 @@
+from edsnlp import registry
+
+from .doc_classifier import TrainableDocClassifier
+
+create_component = registry.factory.register(
+    "eds.doc_classifier",
+    assigns=["doc._.predicted_class"],
+    deprecated=[],
+)(TrainableDocClassifier)

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

@@ -0,0 +1,105 @@
+from typing import Any, Dict, Optional
+
+import torch
+from spacy.tokens import Doc
+from typing_extensions import Literal, TypedDict
+
+from edsnlp.core.pipeline import Pipeline
+from edsnlp.core.torch_component import BatchInput
+from edsnlp.pipes.base import BaseComponent
+from edsnlp.pipes.trainable.embeddings.typing import WordEmbeddingComponent
+
+DocPoolerBatchInput = TypedDict(
+    "DocPoolerBatchInput",
+    {
+        "embedding": BatchInput,
+        "mask": torch.Tensor,  # shape: (batch_size, seq_len)
+        "stats": Dict[str, Any],
+    },
+)
+
+DocPoolerBatchOutput = TypedDict(
+    "DocPoolerBatchOutput",
+    {
+        "embeddings": torch.Tensor,  # shape: (batch_size, embedding_dim)
+    },
+)
+
+
+class DocPooler(WordEmbeddingComponent, BaseComponent):
+    """
+    Pools word embeddings over the entire document to produce
+    a single embedding per doc.
+
+    Parameters
+    ----------
+    nlp: Pipeline
+        The pipeline object
+    name: str
+        Name of the component
+    embedding : WordEmbeddingComponent
+        The word embedding component
+    pooling_mode: Literal["max", "sum", "mean"]
+        How word embeddings are aggregated into a single embedding per document.
+    hidden_size : Optional[int]
+        The size of the hidden layer. If None, no projection is done.
+    """
+
+    def __init__(
+        self,
+        nlp: Optional[Pipeline] = None,
+        name: str = "document_pooler",
+        *,
+        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)
+
+    def preprocess(self, doc: Doc, **kwargs) -> Dict[str, Any]:
+        embedding_out = self.embedding.preprocess(doc, **kwargs)
+        return {
+            "embedding": embedding_out,
+            "stats": {"doc_length": len(doc)},
+        }
+
+    def collate(self, batch: Dict[str, Any]) -> DocPoolerBatchInput:
+        embedding_batch = self.embedding.collate(batch["embedding"])
+        stats = batch["stats"]
+        return {
+            "embedding": embedding_batch,
+            "stats": {
+                "doc_length": sum(stats["doc_length"])
+            },  # <-- sum(...) pour aggréger les comptes par doc en un compte par batch
+        }
+
+    def forward(self, batch: DocPoolerBatchInput) -> DocPoolerBatchOutput:
+        device = next(self.parameters()).device
+
+        embeds = self.embedding(batch["embedding"])["embeddings"]
+        device = embeds.device
+
+        if self.pooling_mode == "mean":
+            pooled = embeds.mean(dim=1)
+        elif self.pooling_mode == "max":
+            pooled = embeds.max(dim=1).values
+        elif self.pooling_mode == "sum":
+            pooled = embeds.sum(dim=1)
+        elif self.pooling_mode == "cls":
+            pooled = self.embedding(batch["embedding"])["cls"].to(device)
+        else:
+            raise ValueError(f"Unknown pooling mode: {self.pooling_mode}")
+
+        pooled = self.feed_forward(pooled)
+        return {"embeddings": pooled}

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

@@ -0,0 +1,9 @@
+from edsnlp import registry
+
+from .doc_pooler import DocPooler
+
+create_component = registry.factory.register(
+    "eds.doc_pooler",
+    assigns=[],
+    deprecated=[],
+)(DocPooler)

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

@@ -505,7 +505,7 @@ def forward(self, batch: TransformerBatchInput) -> TransformerBatchOutput:
                 if "out of memory" in str(e) and trial_idx <= 2:
                     print(
                         f"Out of memory: tried to fit {max_windows} "
-                        f"in {free_mem / (1024 ** 3)} (try n°{trial_idx}/2)"
+                        f"in {free_mem / (1024**3)} (try n°{trial_idx}/2)"
                     )
                     torch.cuda.empty_cache()
                     self._mem_per_unit = (free_mem / max_windows) * 1.5
@@ -535,6 +535,7 @@ def forward(self, batch: TransformerBatchInput) -> TransformerBatchOutput:
         word_embeddings[batch["empty_word_indices"]] = self.empty_word_embedding
         return {
             "embeddings": word_embeddings.refold("context", "word"),
+            "cls": wordpiece_embeddings[:, 0, :],
         }
 
     @staticmethod