pyhighlights 0.1.0__tar.gz

pyhighlights-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Federico Ruggeri
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

pyhighlights-0.1.0/PKG-INFO

@@ -0,0 +1,104 @@
+Metadata-Version: 2.4
+Name: pyhighlights
+Version: 0.1.0
+Summary: A simple library implementing Select-Then-Predict (SPP) models.
+Author-email: Federico Ruggeri <federico.ruggeri6@unibo.it>
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/federicoruggeri/pyhighlights
+Project-URL: Documentation, https://federicoruggeri.github.io/pyhighlights/
+Project-URL: Source, https://github.com/federicoruggeri/pyhighlights
+Project-URL: Issues, https://github.com/federicoruggeri/pyhighlights/issues
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Operating System :: OS Independent
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: cinnamon-core<3,>=2.0.0
+Requires-Dist: lightning>=2.0.0
+Requires-Dist: torch>=2.0.0
+Requires-Dist: torchmetrics>=1.0.0
+Provides-Extra: transformers
+Requires-Dist: transformers>=4.0.0; extra == "transformers"
+Provides-Extra: dev
+Requires-Dist: nox>=2024.3.2; extra == "dev"
+Requires-Dist: pytest>=7.0.0; extra == "dev"
+Requires-Dist: pytest-cov>=5.0; extra == "dev"
+Requires-Dist: ruff>=0.15.20; extra == "dev"
+Provides-Extra: docs
+Requires-Dist: sphinx>=7.0; extra == "docs"
+Requires-Dist: sphinx-autodoc-typehints>=2.0; extra == "docs"
+Requires-Dist: sphinx-rtd-theme>=2.0; extra == "docs"
+Dynamic: license-file
+
+# pyhighlights
+
+[![Tests](https://github.com/federicoruggeri/pyhighlights/actions/workflows/ci.yml/badge.svg)](https://github.com/federicoruggeri/pyhighlights/actions/workflows/ci.yml)
+[![Documentation](https://github.com/federicoruggeri/pyhighlights/actions/workflows/docs.yml/badge.svg)](https://federicoruggeri.github.io/pyhighlights/)
+[![PyPI](https://img.shields.io/pypi/v/pyhighlights)](https://pypi.org/project/pyhighlights/)
+
+Research library for highlight-based explainable AI models.
+
+[Documentation](https://federicoruggeri.github.io/pyhighlights/) · [Contributing](CONTRIBUTING.md)
+
+## Installation
+
+```bash
+pip install pyhighlights
+```
+
+Use `pip install "pyhighlights[transformers]"` for Transformer backends.
+
+## GRU folded rationalization
+
+```python
+from pathlib import Path
+
+import pyhighlights
+from cinnamon.registry import Registry
+from pyhighlights.configurations.spp import (
+    GRU_FR,
+    GRU_GRAT,
+    GRU_MCD,
+    GRU_MGR,
+    TRANSFORMER_FR,
+)
+
+Registry.build(directory=Path(pyhighlights.__file__).parent)
+fr = Registry.from_key(GRU_FR)
+mgr = Registry.from_key(GRU_MGR)
+mcd = Registry.from_key(GRU_MCD)
+grat = Registry.from_key(GRU_GRAT)
+transformer_fr = Registry.from_key(TRANSFORMER_FR)  # needs pyhighlights[transformers]
+```
+
+GRU and Transformer implementations conform to same backbone interface; model
+classes contain rationalization logic only.
+
+## Highlight data
+
+```python
+from torch.utils.data import DataLoader
+from pyhighlights.components import (
+    HighlightCollator,
+    HighlightDataset,
+    HighlightExample,
+    VocabularyTokenizer,
+)
+
+examples = HighlightDataset([
+    HighlightExample(0, ["great", "stay"], label=1, highlights=[1, 0]),
+    HighlightExample(1, ["bad"], label=0),  # unlabeled highlights become -1
+])
+collator = HighlightCollator(VocabularyTokenizer({"great": 1, "stay": 2, "bad": 3}))
+batch = next(iter(DataLoader(examples, batch_size=2, collate_fn=collator)))
+```
+
+`HuggingFaceTokenizer` expands word highlights across subtokens and requires
+`pyhighlights[transformers]`.

pyhighlights-0.1.0/README.md

@@ -0,0 +1,65 @@
+# pyhighlights
+
+[![Tests](https://github.com/federicoruggeri/pyhighlights/actions/workflows/ci.yml/badge.svg)](https://github.com/federicoruggeri/pyhighlights/actions/workflows/ci.yml)
+[![Documentation](https://github.com/federicoruggeri/pyhighlights/actions/workflows/docs.yml/badge.svg)](https://federicoruggeri.github.io/pyhighlights/)
+[![PyPI](https://img.shields.io/pypi/v/pyhighlights)](https://pypi.org/project/pyhighlights/)
+
+Research library for highlight-based explainable AI models.
+
+[Documentation](https://federicoruggeri.github.io/pyhighlights/) · [Contributing](CONTRIBUTING.md)
+
+## Installation
+
+```bash
+pip install pyhighlights
+```
+
+Use `pip install "pyhighlights[transformers]"` for Transformer backends.
+
+## GRU folded rationalization
+
+```python
+from pathlib import Path
+
+import pyhighlights
+from cinnamon.registry import Registry
+from pyhighlights.configurations.spp import (
+    GRU_FR,
+    GRU_GRAT,
+    GRU_MCD,
+    GRU_MGR,
+    TRANSFORMER_FR,
+)
+
+Registry.build(directory=Path(pyhighlights.__file__).parent)
+fr = Registry.from_key(GRU_FR)
+mgr = Registry.from_key(GRU_MGR)
+mcd = Registry.from_key(GRU_MCD)
+grat = Registry.from_key(GRU_GRAT)
+transformer_fr = Registry.from_key(TRANSFORMER_FR)  # needs pyhighlights[transformers]
+```
+
+GRU and Transformer implementations conform to same backbone interface; model
+classes contain rationalization logic only.
+
+## Highlight data
+
+```python
+from torch.utils.data import DataLoader
+from pyhighlights.components import (
+    HighlightCollator,
+    HighlightDataset,
+    HighlightExample,
+    VocabularyTokenizer,
+)
+
+examples = HighlightDataset([
+    HighlightExample(0, ["great", "stay"], label=1, highlights=[1, 0]),
+    HighlightExample(1, ["bad"], label=0),  # unlabeled highlights become -1
+])
+collator = HighlightCollator(VocabularyTokenizer({"great": 1, "stay": 2, "bad": 3}))
+batch = next(iter(DataLoader(examples, batch_size=2, collate_fn=collator)))
+```
+
+`HuggingFaceTokenizer` expands word highlights across subtokens and requires
+`pyhighlights[transformers]`.

pyhighlights-0.1.0/pyhighlights/__init__.py

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

pyhighlights-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-0.1.0/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-0.1.0/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-0.1.0/pyhighlights/components/models/__init__.py

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

pyhighlights-0.1.0/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