pyhighlights 0.1.0__py3-none-any.whl

pyhighlights/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

pyhighlights/components/__init__.py

@@ -0,0 +1,19 @@
+from pyhighlights.components.data import (
+    HighlightCollator,
+    HighlightDataset,
+    HighlightExample,
+    HighlightTokenizer,
+    HuggingFaceTokenizer,
+    TokenizedExample,
+    VocabularyTokenizer,
+)
+
+__all__ = [
+    "HighlightCollator",
+    "HighlightDataset",
+    "HighlightExample",
+    "HighlightTokenizer",
+    "HuggingFaceTokenizer",
+    "TokenizedExample",
+    "VocabularyTokenizer",
+]

pyhighlights/components/data.py

@@ -0,0 +1,211 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from numbers import Integral
+from typing import Iterable, Mapping, Protocol, Sequence
+
+import torch as th
+from torch.utils.data import Dataset
+
+from pyhighlights.components.models import InputData
+
+
+@dataclass(frozen=True)
+class HighlightExample:
+    sample_id: int
+    tokens: Sequence[str]
+    label: int
+    highlights: Sequence[int] | None = None
+
+    def __post_init__(self):
+        tokens = tuple(self.tokens)
+        highlights = None if self.highlights is None else tuple(self.highlights)
+        if not isinstance(self.sample_id, Integral) or not isinstance(
+            self.label, Integral
+        ):
+            raise TypeError("sample_id and label must be integers")
+        if any(not isinstance(token, str) for token in tokens):
+            raise TypeError("tokens must contain strings")
+        if highlights is not None:
+            if len(highlights) != len(tokens):
+                raise ValueError("highlights must align with tokens")
+            if any(value not in (0, 1) for value in highlights):
+                raise ValueError("highlights must contain only 0 or 1")
+        object.__setattr__(self, "tokens", tokens)
+        object.__setattr__(self, "highlights", highlights)
+
+
+@dataclass(frozen=True)
+class TokenizedExample:
+    input_ids: Sequence[int]
+    word_ids: Sequence[int | None]
+
+    def __post_init__(self):
+        input_ids = tuple(self.input_ids)
+        word_ids = tuple(self.word_ids)
+        if len(input_ids) != len(word_ids):
+            raise ValueError("input_ids and word_ids must have equal length")
+        if any(not isinstance(value, Integral) or value < 0 for value in input_ids):
+            raise ValueError("input_ids must contain non-negative integers")
+        if any(
+            value is not None and (not isinstance(value, Integral) or value < 0)
+            for value in word_ids
+        ):
+            raise ValueError("word_ids must contain non-negative integers or None")
+        object.__setattr__(self, "input_ids", input_ids)
+        object.__setattr__(self, "word_ids", word_ids)
+
+
+class HighlightTokenizer(Protocol):
+    pad_token_id: int
+
+    def encode(
+        self, tokens: Sequence[str], max_length: int | None = None
+    ) -> TokenizedExample: ...
+
+
+class VocabularyTokenizer:
+    """One-token-to-one-id encoder for GRU backbones."""
+
+    def __init__(
+        self,
+        vocabulary: Mapping[str, int],
+        unknown_token_id: int = 0,
+        pad_token_id: int = 0,
+    ):
+        if min([unknown_token_id, pad_token_id, *vocabulary.values()]) < 0:
+            raise ValueError("token ids must be non-negative")
+        self.vocabulary = dict(vocabulary)
+        self.unknown_token_id = unknown_token_id
+        self.pad_token_id = pad_token_id
+
+    def encode(
+        self, tokens: Sequence[str], max_length: int | None = None
+    ) -> TokenizedExample:
+        tokens = tokens if max_length is None else tokens[:max_length]
+        return TokenizedExample(
+            input_ids=[
+                self.vocabulary.get(token, self.unknown_token_id) for token in tokens
+            ],
+            word_ids=list(range(len(tokens))),
+        )
+
+
+class HuggingFaceTokenizer:
+    """Fast-tokenizer adapter preserving source-token alignment."""
+
+    def __init__(self, pretrained_model_card: str, **tokenizer_kwargs):
+        try:
+            from transformers import AutoTokenizer
+        except ImportError as error:
+            raise ImportError(
+                "HuggingFaceTokenizer requires pyhighlights[transformers]"
+            ) from error
+
+        if tokenizer_kwargs.pop("use_fast", True) is not True:
+            raise ValueError("HuggingFaceTokenizer requires a fast tokenizer")
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            pretrained_model_card, use_fast=True, **tokenizer_kwargs
+        )
+        if not getattr(self.tokenizer, "is_fast", False):
+            raise ValueError("HuggingFaceTokenizer requires a fast tokenizer")
+        if self.tokenizer.pad_token_id is None:
+            raise ValueError("tokenizer must define pad_token_id")
+        self.pad_token_id = self.tokenizer.pad_token_id
+
+    def encode(
+        self, tokens: Sequence[str], max_length: int | None = None
+    ) -> TokenizedExample:
+        kwargs = {
+            "is_split_into_words": True,
+            "add_special_tokens": False,
+            "return_attention_mask": False,
+        }
+        if max_length is not None:
+            kwargs.update(truncation=True, max_length=max_length)
+        encoded = self.tokenizer(list(tokens), **kwargs)
+        try:
+            word_ids = encoded.word_ids()
+        except (AttributeError, ValueError) as error:
+            raise ValueError(
+                "tokenizer must provide word_ids; use a fast tokenizer"
+            ) from error
+        return TokenizedExample(
+            input_ids=encoded["input_ids"],
+            word_ids=word_ids,
+        )
+
+
+class HighlightDataset(Dataset):
+    def __init__(self, examples: Iterable[HighlightExample]):
+        self.examples = list(examples)
+
+    def __len__(self) -> int:
+        return len(self.examples)
+
+    def __getitem__(self, index: int) -> HighlightExample:
+        return self.examples[index]
+
+
+class HighlightCollator:
+    """Tokenize, align highlights to subtokens, and dynamically pad a batch."""
+
+    def __init__(
+        self,
+        tokenizer: HighlightTokenizer,
+        max_length: int | None = None,
+    ):
+        if max_length is not None and max_length < 1:
+            raise ValueError("max_length must be positive")
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+
+    def __call__(self, examples: Sequence[HighlightExample]) -> InputData:
+        if not examples:
+            raise ValueError("cannot collate an empty batch")
+        encoded = [
+            self.tokenizer.encode(example.tokens, self.max_length)
+            for example in examples
+        ]
+        if self.max_length is not None:
+            encoded = [
+                TokenizedExample(
+                    item.input_ids[: self.max_length],
+                    item.word_ids[: self.max_length],
+                )
+                for item in encoded
+            ]
+        width = max(max(len(item.input_ids) for item in encoded), 1)
+
+        features = []
+        masks = []
+        highlights = []
+        for example, item in zip(examples, encoded):
+            input_ids = list(item.input_ids[:width])
+            word_ids = list(item.word_ids[:width])
+            for word_id in word_ids:
+                if word_id is not None and not 0 <= word_id < len(example.tokens):
+                    raise ValueError("word_id is outside source-token range")
+
+            padding = width - len(input_ids)
+            features.append(input_ids + [self.tokenizer.pad_token_id] * padding)
+            masks.append(
+                [word_id is not None for word_id in word_ids] + [False] * padding
+            )
+            highlights.append(
+                [
+                    -1
+                    if word_id is None or example.highlights is None
+                    else example.highlights[word_id]
+                    for word_id in word_ids
+                ]
+                + [-1] * padding
+            )
+
+        return InputData(
+            features=th.tensor(features, dtype=th.long),
+            mask=th.tensor(masks, dtype=th.float32),
+            sample_ids=th.tensor([example.sample_id for example in examples]),
+            y_true=th.tensor([example.label for example in examples]),
+            highlight_true=th.tensor(highlights, dtype=th.long),
+        )

pyhighlights/components/loaders.py

@@ -0,0 +1,21 @@
+"""Compatibility import path for data utilities."""
+
+from pyhighlights.components.data import (
+    HighlightCollator,
+    HighlightDataset,
+    HighlightExample,
+    HighlightTokenizer,
+    HuggingFaceTokenizer,
+    TokenizedExample,
+    VocabularyTokenizer,
+)
+
+__all__ = [
+    "HighlightCollator",
+    "HighlightDataset",
+    "HighlightExample",
+    "HighlightTokenizer",
+    "HuggingFaceTokenizer",
+    "TokenizedExample",
+    "VocabularyTokenizer",
+]

pyhighlights/components/models/__init__.py

@@ -0,0 +1,3 @@
+from pyhighlights.components.models.data import InputData, OutputData, SPPOutput
+
+__all__ = ["InputData", "OutputData", "SPPOutput"]

pyhighlights/components/models/base.py

@@ -0,0 +1,178 @@
+from __future__ import annotations
+
+import abc
+from typing import Dict, List, Literal, Tuple
+
+import lightning as L
+import torch as th
+from cinnamon.registry import RegistrationKey, Registry
+from torchmetrics import Metric, MetricCollection
+
+from pyhighlights.components.models.data import (
+    InputData,
+    ModelData,
+    OutputData,
+    SPPOutput,
+)
+from pyhighlights.utility.losses import Loss, build_losses
+from pyhighlights.utility.metrics import build_torchmetrics
+
+Split = Literal["train", "val", "test"]
+
+# Keep data containers importable from this module for compatibility.
+__all__ = ["InputData", "Model", "ModelData", "OutputData", "SPPOutput", "Split"]
+
+
+class Model(L.LightningModule, abc.ABC):
+    def __init__(
+        self,
+        name: str,
+        losses: List[RegistrationKey[Loss]],
+        optimizer: RegistrationKey[th.optim.Optimizer],
+        train_metrics: Dict[str, RegistrationKey[Metric]] | None = None,
+        val_metrics: Dict[str, RegistrationKey[Metric]] | None = None,
+        test_metrics: Dict[str, RegistrationKey[Metric]] | None = None,
+    ):
+        super().__init__()
+
+        self.save_hyperparameters(ignore=self.ignore_hyperparameters())
+        self.name = name
+        self.optimizer = optimizer
+
+        self.train_metrics = self._build_metrics(train_metrics)
+        self.val_metrics = self._build_metrics(val_metrics)
+        self.test_metrics = self._build_metrics(test_metrics)
+        self.losses = th.nn.ModuleList(build_losses(keys=losses))
+
+        self.store_predictions = False
+        self.predictions = []
+        self.forward_mapping = {
+            "train": self.training_forward,
+            "val": self.validation_forward,
+            "test": self.test_forward,
+        }
+
+    @staticmethod
+    def _build_metrics(
+        keys: Dict[str, RegistrationKey[Metric]] | None,
+    ) -> MetricCollection | None:
+        return build_torchmetrics(keys) if keys is not None else None
+
+    def ignore_hyperparameters(self) -> List[str]:
+        return []
+
+    def enable_storing_predictions(self):
+        self.store_predictions = True
+
+    def disable_storing_predictions(self):
+        self.store_predictions = False
+
+    def flush_predictions(self):
+        self.predictions.clear()
+
+    def update_metrics(
+        self, split: Split, input_data: InputData, output_data: OutputData
+    ):
+        metrics: MetricCollection | None = getattr(self, f"{split}_metrics")
+        if metrics is not None:
+            metrics.update(output_data.class_logits, input_data.y_true)
+
+    def compute_metrics(self, split: Split):
+        metrics: MetricCollection | None = getattr(self, f"{split}_metrics")
+        if metrics is None:
+            return
+
+        for key, value in metrics.compute().items():
+            self.log(f"{split}_{key}", value, prog_bar=True)
+        metrics.reset()
+
+    def on_train_epoch_end(self) -> None:
+        self.compute_metrics(split="train")
+
+    def on_validation_epoch_end(self) -> None:
+        self.compute_metrics(split="val")
+
+    def on_test_epoch_end(self) -> None:
+        self.compute_metrics(split="test")
+
+    def configure_optimizers(self):
+        return Registry.from_key(self.optimizer, params=self.parameters())
+
+    def log_metrics(
+        self,
+        split: Split,
+        total_loss: th.Tensor,
+        losses: Dict[str, th.Tensor],
+        batch_size: int,
+    ):
+        self.log(
+            name=f"{split}_loss",
+            value=total_loss,
+            on_step=False,
+            on_epoch=True,
+            prog_bar=True,
+            batch_size=batch_size,
+        )
+        for loss_name, loss_value in losses.items():
+            self.log(
+                name=f"{split}_{loss_name}",
+                value=loss_value,
+                on_step=False,
+                on_epoch=True,
+                prog_bar=True,
+                batch_size=batch_size,
+            )
+
+    def training_forward(self, batch: InputData) -> OutputData:
+        return self.forward(data=batch)
+
+    def validation_forward(self, batch: InputData) -> OutputData:
+        return self.training_forward(batch=batch)
+
+    def test_forward(self, batch: InputData) -> OutputData:
+        return self.training_forward(batch=batch)
+
+    def _step(self, batch: InputData, batch_idx: int, split: Split) -> th.Tensor:
+        output_data = self.forward_mapping[split](batch)
+        total_loss, losses = self.compute_loss(
+            input_data=batch, output_data=output_data
+        )
+
+        self.log_metrics(
+            split=split,
+            total_loss=total_loss,
+            losses=losses,
+            batch_size=batch.y_true.shape[0],
+        )
+        self.update_metrics(split=split, input_data=batch, output_data=output_data)
+
+        if self.store_predictions:
+            self.predictions.append({**batch.as_numpy(), **output_data.as_numpy()})
+
+        return total_loss
+
+    def training_step(self, batch: InputData, batch_idx: int):
+        return self._step(batch=batch, batch_idx=batch_idx, split="train")
+
+    def validation_step(self, batch: InputData, batch_idx: int):
+        return self._step(batch=batch, batch_idx=batch_idx, split="val")
+
+    def test_step(self, batch: InputData, batch_idx: int):
+        return self._step(batch=batch, batch_idx=batch_idx, split="test")
+
+    def compute_loss(
+        self,
+        input_data: InputData,
+        output_data: OutputData,
+    ) -> Tuple[th.Tensor, Dict[str, th.Tensor]]:
+        total_loss = output_data.class_logits.new_zeros(())
+        losses = {}
+
+        for loss in self.losses:
+            if not loss.enabled:
+                continue
+            loss_value = loss(input_data=input_data, output_data=output_data)
+            total_loss = total_loss + loss_value * loss.coefficient
+            losses[loss.name] = loss_value
+
+        return total_loss, losses

pyhighlights/components/models/data.py

@@ -0,0 +1,70 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, fields
+from typing import Generator, TypeVar
+
+import torch as th
+
+D = TypeVar("D", bound="ModelData")
+
+
+@dataclass
+class ModelData:
+    def as_numpy(self):
+        return {
+            field.name: value.detach().cpu().numpy()
+            for field in fields(self)
+            if isinstance((value := getattr(self, field.name)), th.Tensor)
+        }
+
+    def unbind(self: D, dim: int = 0) -> Generator[D, None, None]:
+        tensor_fields = {
+            field.name: th.unbind(value, dim=dim)
+            for field in fields(self)
+            if isinstance((value := getattr(self, field.name)), th.Tensor)
+        }
+        sizes = {len(values) for values in tensor_fields.values()}
+        if not sizes:
+            raise RuntimeError("No tensors found to unbind")
+        if len(sizes) != 1:
+            raise RuntimeError("Cannot unbind tensors of different sizes")
+
+        for index in range(sizes.pop()):
+            yield type(self)(
+                **{
+                    field.name: tensor_fields[field.name][index]
+                    if field.name in tensor_fields
+                    else getattr(self, field.name)
+                    for field in fields(self)
+                }
+            )
+
+
+@dataclass
+class InputData(ModelData):
+    features: th.Tensor
+    mask: th.Tensor
+    sample_ids: th.Tensor
+    y_true: th.Tensor
+    highlight_true: th.Tensor
+
+
+@dataclass
+class OutputData(ModelData):
+    class_logits: th.Tensor
+
+    @property
+    def y_pred(self) -> th.Tensor:
+        """Compatibility alias for pre-0.2 code."""
+        return self.class_logits
+
+
+@dataclass
+class SPPOutput(OutputData):
+    highlight_logits: th.Tensor
+    highlight_mask: th.Tensor
+
+    @property
+    def highlight_pred(self) -> th.Tensor:
+        """Compatibility alias for pre-0.2 code."""
+        return self.highlight_mask

pyhighlights/components/models/spp/__init__.py

@@ -0,0 +1,41 @@
+from pyhighlights.components.models.spp.base import (
+    SPPAggregator,
+    SPPBackbone,
+    SPPFirstAggregator,
+    SPPPredictor,
+    SPPSelector,
+)
+from pyhighlights.components.models.spp.fr import FR
+from pyhighlights.components.models.spp.grat import (
+    GRAT,
+    AttentionGuider,
+    GRATGuider,
+    GRATGuiderOutput,
+)
+from pyhighlights.components.models.spp.implementations import (
+    GRUBackbone,
+    MLPPredictor,
+    MLPSelector,
+    TransformerBackbone,
+)
+from pyhighlights.components.models.spp.mcd import MCD
+from pyhighlights.components.models.spp.mgr import MGR
+
+__all__ = [
+    "AttentionGuider",
+    "FR",
+    "GRAT",
+    "GRATGuider",
+    "GRATGuiderOutput",
+    "GRUBackbone",
+    "MCD",
+    "MGR",
+    "MLPPredictor",
+    "MLPSelector",
+    "SPPAggregator",
+    "SPPBackbone",
+    "SPPFirstAggregator",
+    "SPPPredictor",
+    "SPPSelector",
+    "TransformerBackbone",
+]