pyhighlights 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (40) hide show
  1. pyhighlights-0.1.0/LICENSE +21 -0
  2. pyhighlights-0.1.0/PKG-INFO +104 -0
  3. pyhighlights-0.1.0/README.md +65 -0
  4. pyhighlights-0.1.0/pyhighlights/__init__.py +1 -0
  5. pyhighlights-0.1.0/pyhighlights/components/__init__.py +19 -0
  6. pyhighlights-0.1.0/pyhighlights/components/data.py +211 -0
  7. pyhighlights-0.1.0/pyhighlights/components/loaders.py +21 -0
  8. pyhighlights-0.1.0/pyhighlights/components/models/__init__.py +3 -0
  9. pyhighlights-0.1.0/pyhighlights/components/models/base.py +178 -0
  10. pyhighlights-0.1.0/pyhighlights/components/models/data.py +70 -0
  11. pyhighlights-0.1.0/pyhighlights/components/models/spp/__init__.py +41 -0
  12. pyhighlights-0.1.0/pyhighlights/components/models/spp/base.py +205 -0
  13. pyhighlights-0.1.0/pyhighlights/components/models/spp/fr.py +18 -0
  14. pyhighlights-0.1.0/pyhighlights/components/models/spp/grat.py +219 -0
  15. pyhighlights-0.1.0/pyhighlights/components/models/spp/implementations.py +141 -0
  16. pyhighlights-0.1.0/pyhighlights/components/models/spp/mcd.py +217 -0
  17. pyhighlights-0.1.0/pyhighlights/components/models/spp/mgr.py +78 -0
  18. pyhighlights-0.1.0/pyhighlights/configurations/__init__.py +0 -0
  19. pyhighlights-0.1.0/pyhighlights/configurations/spp.py +398 -0
  20. pyhighlights-0.1.0/pyhighlights/metrics.py +178 -0
  21. pyhighlights-0.1.0/pyhighlights/models/__init__.py +0 -0
  22. pyhighlights-0.1.0/pyhighlights/models/layers.py +54 -0
  23. pyhighlights-0.1.0/pyhighlights/utility/__init__.py +0 -0
  24. pyhighlights-0.1.0/pyhighlights/utility/losses.py +151 -0
  25. pyhighlights-0.1.0/pyhighlights/utility/metrics.py +14 -0
  26. pyhighlights-0.1.0/pyhighlights.egg-info/PKG-INFO +104 -0
  27. pyhighlights-0.1.0/pyhighlights.egg-info/SOURCES.txt +38 -0
  28. pyhighlights-0.1.0/pyhighlights.egg-info/dependency_links.txt +1 -0
  29. pyhighlights-0.1.0/pyhighlights.egg-info/requires.txt +18 -0
  30. pyhighlights-0.1.0/pyhighlights.egg-info/top_level.txt +1 -0
  31. pyhighlights-0.1.0/pyproject.toml +74 -0
  32. pyhighlights-0.1.0/setup.cfg +4 -0
  33. pyhighlights-0.1.0/tests/test_configurations.py +185 -0
  34. pyhighlights-0.1.0/tests/test_data.py +146 -0
  35. pyhighlights-0.1.0/tests/test_fr.py +108 -0
  36. pyhighlights-0.1.0/tests/test_grat.py +89 -0
  37. pyhighlights-0.1.0/tests/test_imports.py +25 -0
  38. pyhighlights-0.1.0/tests/test_losses.py +69 -0
  39. pyhighlights-0.1.0/tests/test_metrics.py +34 -0
  40. pyhighlights-0.1.0/tests/test_transformer_configurations.py +102 -0
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Federico Ruggeri
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,104 @@
1
+ Metadata-Version: 2.4
2
+ Name: pyhighlights
3
+ Version: 0.1.0
4
+ Summary: A simple library implementing Select-Then-Predict (SPP) models.
5
+ Author-email: Federico Ruggeri <federico.ruggeri6@unibo.it>
6
+ License-Expression: MIT
7
+ Project-URL: Homepage, https://github.com/federicoruggeri/pyhighlights
8
+ Project-URL: Documentation, https://federicoruggeri.github.io/pyhighlights/
9
+ Project-URL: Source, https://github.com/federicoruggeri/pyhighlights
10
+ Project-URL: Issues, https://github.com/federicoruggeri/pyhighlights/issues
11
+ Classifier: Development Status :: 3 - Alpha
12
+ Classifier: Intended Audience :: Science/Research
13
+ Classifier: Programming Language :: Python :: 3
14
+ Classifier: Programming Language :: Python :: 3.10
15
+ Classifier: Programming Language :: Python :: 3.11
16
+ Classifier: Programming Language :: Python :: 3.12
17
+ Classifier: Programming Language :: Python :: 3.13
18
+ Classifier: Operating System :: OS Independent
19
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
20
+ Requires-Python: >=3.10
21
+ Description-Content-Type: text/markdown
22
+ License-File: LICENSE
23
+ Requires-Dist: cinnamon-core<3,>=2.0.0
24
+ Requires-Dist: lightning>=2.0.0
25
+ Requires-Dist: torch>=2.0.0
26
+ Requires-Dist: torchmetrics>=1.0.0
27
+ Provides-Extra: transformers
28
+ Requires-Dist: transformers>=4.0.0; extra == "transformers"
29
+ Provides-Extra: dev
30
+ Requires-Dist: nox>=2024.3.2; extra == "dev"
31
+ Requires-Dist: pytest>=7.0.0; extra == "dev"
32
+ Requires-Dist: pytest-cov>=5.0; extra == "dev"
33
+ Requires-Dist: ruff>=0.15.20; extra == "dev"
34
+ Provides-Extra: docs
35
+ Requires-Dist: sphinx>=7.0; extra == "docs"
36
+ Requires-Dist: sphinx-autodoc-typehints>=2.0; extra == "docs"
37
+ Requires-Dist: sphinx-rtd-theme>=2.0; extra == "docs"
38
+ Dynamic: license-file
39
+
40
+ # pyhighlights
41
+
42
+ [![Tests](https://github.com/federicoruggeri/pyhighlights/actions/workflows/ci.yml/badge.svg)](https://github.com/federicoruggeri/pyhighlights/actions/workflows/ci.yml)
43
+ [![Documentation](https://github.com/federicoruggeri/pyhighlights/actions/workflows/docs.yml/badge.svg)](https://federicoruggeri.github.io/pyhighlights/)
44
+ [![PyPI](https://img.shields.io/pypi/v/pyhighlights)](https://pypi.org/project/pyhighlights/)
45
+
46
+ Research library for highlight-based explainable AI models.
47
+
48
+ [Documentation](https://federicoruggeri.github.io/pyhighlights/) · [Contributing](CONTRIBUTING.md)
49
+
50
+ ## Installation
51
+
52
+ ```bash
53
+ pip install pyhighlights
54
+ ```
55
+
56
+ Use `pip install "pyhighlights[transformers]"` for Transformer backends.
57
+
58
+ ## GRU folded rationalization
59
+
60
+ ```python
61
+ from pathlib import Path
62
+
63
+ import pyhighlights
64
+ from cinnamon.registry import Registry
65
+ from pyhighlights.configurations.spp import (
66
+ GRU_FR,
67
+ GRU_GRAT,
68
+ GRU_MCD,
69
+ GRU_MGR,
70
+ TRANSFORMER_FR,
71
+ )
72
+
73
+ Registry.build(directory=Path(pyhighlights.__file__).parent)
74
+ fr = Registry.from_key(GRU_FR)
75
+ mgr = Registry.from_key(GRU_MGR)
76
+ mcd = Registry.from_key(GRU_MCD)
77
+ grat = Registry.from_key(GRU_GRAT)
78
+ transformer_fr = Registry.from_key(TRANSFORMER_FR) # needs pyhighlights[transformers]
79
+ ```
80
+
81
+ GRU and Transformer implementations conform to same backbone interface; model
82
+ classes contain rationalization logic only.
83
+
84
+ ## Highlight data
85
+
86
+ ```python
87
+ from torch.utils.data import DataLoader
88
+ from pyhighlights.components import (
89
+ HighlightCollator,
90
+ HighlightDataset,
91
+ HighlightExample,
92
+ VocabularyTokenizer,
93
+ )
94
+
95
+ examples = HighlightDataset([
96
+ HighlightExample(0, ["great", "stay"], label=1, highlights=[1, 0]),
97
+ HighlightExample(1, ["bad"], label=0), # unlabeled highlights become -1
98
+ ])
99
+ collator = HighlightCollator(VocabularyTokenizer({"great": 1, "stay": 2, "bad": 3}))
100
+ batch = next(iter(DataLoader(examples, batch_size=2, collate_fn=collator)))
101
+ ```
102
+
103
+ `HuggingFaceTokenizer` expands word highlights across subtokens and requires
104
+ `pyhighlights[transformers]`.
@@ -0,0 +1,65 @@
1
+ # pyhighlights
2
+
3
+ [![Tests](https://github.com/federicoruggeri/pyhighlights/actions/workflows/ci.yml/badge.svg)](https://github.com/federicoruggeri/pyhighlights/actions/workflows/ci.yml)
4
+ [![Documentation](https://github.com/federicoruggeri/pyhighlights/actions/workflows/docs.yml/badge.svg)](https://federicoruggeri.github.io/pyhighlights/)
5
+ [![PyPI](https://img.shields.io/pypi/v/pyhighlights)](https://pypi.org/project/pyhighlights/)
6
+
7
+ Research library for highlight-based explainable AI models.
8
+
9
+ [Documentation](https://federicoruggeri.github.io/pyhighlights/) · [Contributing](CONTRIBUTING.md)
10
+
11
+ ## Installation
12
+
13
+ ```bash
14
+ pip install pyhighlights
15
+ ```
16
+
17
+ Use `pip install "pyhighlights[transformers]"` for Transformer backends.
18
+
19
+ ## GRU folded rationalization
20
+
21
+ ```python
22
+ from pathlib import Path
23
+
24
+ import pyhighlights
25
+ from cinnamon.registry import Registry
26
+ from pyhighlights.configurations.spp import (
27
+ GRU_FR,
28
+ GRU_GRAT,
29
+ GRU_MCD,
30
+ GRU_MGR,
31
+ TRANSFORMER_FR,
32
+ )
33
+
34
+ Registry.build(directory=Path(pyhighlights.__file__).parent)
35
+ fr = Registry.from_key(GRU_FR)
36
+ mgr = Registry.from_key(GRU_MGR)
37
+ mcd = Registry.from_key(GRU_MCD)
38
+ grat = Registry.from_key(GRU_GRAT)
39
+ transformer_fr = Registry.from_key(TRANSFORMER_FR) # needs pyhighlights[transformers]
40
+ ```
41
+
42
+ GRU and Transformer implementations conform to same backbone interface; model
43
+ classes contain rationalization logic only.
44
+
45
+ ## Highlight data
46
+
47
+ ```python
48
+ from torch.utils.data import DataLoader
49
+ from pyhighlights.components import (
50
+ HighlightCollator,
51
+ HighlightDataset,
52
+ HighlightExample,
53
+ VocabularyTokenizer,
54
+ )
55
+
56
+ examples = HighlightDataset([
57
+ HighlightExample(0, ["great", "stay"], label=1, highlights=[1, 0]),
58
+ HighlightExample(1, ["bad"], label=0), # unlabeled highlights become -1
59
+ ])
60
+ collator = HighlightCollator(VocabularyTokenizer({"great": 1, "stay": 2, "bad": 3}))
61
+ batch = next(iter(DataLoader(examples, batch_size=2, collate_fn=collator)))
62
+ ```
63
+
64
+ `HuggingFaceTokenizer` expands word highlights across subtokens and requires
65
+ `pyhighlights[transformers]`.
@@ -0,0 +1 @@
1
+ __version__ = "0.1.0"
@@ -0,0 +1,19 @@
1
+ from pyhighlights.components.data import (
2
+ HighlightCollator,
3
+ HighlightDataset,
4
+ HighlightExample,
5
+ HighlightTokenizer,
6
+ HuggingFaceTokenizer,
7
+ TokenizedExample,
8
+ VocabularyTokenizer,
9
+ )
10
+
11
+ __all__ = [
12
+ "HighlightCollator",
13
+ "HighlightDataset",
14
+ "HighlightExample",
15
+ "HighlightTokenizer",
16
+ "HuggingFaceTokenizer",
17
+ "TokenizedExample",
18
+ "VocabularyTokenizer",
19
+ ]
@@ -0,0 +1,211 @@
1
+ from __future__ import annotations
2
+
3
+ from dataclasses import dataclass
4
+ from numbers import Integral
5
+ from typing import Iterable, Mapping, Protocol, Sequence
6
+
7
+ import torch as th
8
+ from torch.utils.data import Dataset
9
+
10
+ from pyhighlights.components.models import InputData
11
+
12
+
13
+ @dataclass(frozen=True)
14
+ class HighlightExample:
15
+ sample_id: int
16
+ tokens: Sequence[str]
17
+ label: int
18
+ highlights: Sequence[int] | None = None
19
+
20
+ def __post_init__(self):
21
+ tokens = tuple(self.tokens)
22
+ highlights = None if self.highlights is None else tuple(self.highlights)
23
+ if not isinstance(self.sample_id, Integral) or not isinstance(
24
+ self.label, Integral
25
+ ):
26
+ raise TypeError("sample_id and label must be integers")
27
+ if any(not isinstance(token, str) for token in tokens):
28
+ raise TypeError("tokens must contain strings")
29
+ if highlights is not None:
30
+ if len(highlights) != len(tokens):
31
+ raise ValueError("highlights must align with tokens")
32
+ if any(value not in (0, 1) for value in highlights):
33
+ raise ValueError("highlights must contain only 0 or 1")
34
+ object.__setattr__(self, "tokens", tokens)
35
+ object.__setattr__(self, "highlights", highlights)
36
+
37
+
38
+ @dataclass(frozen=True)
39
+ class TokenizedExample:
40
+ input_ids: Sequence[int]
41
+ word_ids: Sequence[int | None]
42
+
43
+ def __post_init__(self):
44
+ input_ids = tuple(self.input_ids)
45
+ word_ids = tuple(self.word_ids)
46
+ if len(input_ids) != len(word_ids):
47
+ raise ValueError("input_ids and word_ids must have equal length")
48
+ if any(not isinstance(value, Integral) or value < 0 for value in input_ids):
49
+ raise ValueError("input_ids must contain non-negative integers")
50
+ if any(
51
+ value is not None and (not isinstance(value, Integral) or value < 0)
52
+ for value in word_ids
53
+ ):
54
+ raise ValueError("word_ids must contain non-negative integers or None")
55
+ object.__setattr__(self, "input_ids", input_ids)
56
+ object.__setattr__(self, "word_ids", word_ids)
57
+
58
+
59
+ class HighlightTokenizer(Protocol):
60
+ pad_token_id: int
61
+
62
+ def encode(
63
+ self, tokens: Sequence[str], max_length: int | None = None
64
+ ) -> TokenizedExample: ...
65
+
66
+
67
+ class VocabularyTokenizer:
68
+ """One-token-to-one-id encoder for GRU backbones."""
69
+
70
+ def __init__(
71
+ self,
72
+ vocabulary: Mapping[str, int],
73
+ unknown_token_id: int = 0,
74
+ pad_token_id: int = 0,
75
+ ):
76
+ if min([unknown_token_id, pad_token_id, *vocabulary.values()]) < 0:
77
+ raise ValueError("token ids must be non-negative")
78
+ self.vocabulary = dict(vocabulary)
79
+ self.unknown_token_id = unknown_token_id
80
+ self.pad_token_id = pad_token_id
81
+
82
+ def encode(
83
+ self, tokens: Sequence[str], max_length: int | None = None
84
+ ) -> TokenizedExample:
85
+ tokens = tokens if max_length is None else tokens[:max_length]
86
+ return TokenizedExample(
87
+ input_ids=[
88
+ self.vocabulary.get(token, self.unknown_token_id) for token in tokens
89
+ ],
90
+ word_ids=list(range(len(tokens))),
91
+ )
92
+
93
+
94
+ class HuggingFaceTokenizer:
95
+ """Fast-tokenizer adapter preserving source-token alignment."""
96
+
97
+ def __init__(self, pretrained_model_card: str, **tokenizer_kwargs):
98
+ try:
99
+ from transformers import AutoTokenizer
100
+ except ImportError as error:
101
+ raise ImportError(
102
+ "HuggingFaceTokenizer requires pyhighlights[transformers]"
103
+ ) from error
104
+
105
+ if tokenizer_kwargs.pop("use_fast", True) is not True:
106
+ raise ValueError("HuggingFaceTokenizer requires a fast tokenizer")
107
+ self.tokenizer = AutoTokenizer.from_pretrained(
108
+ pretrained_model_card, use_fast=True, **tokenizer_kwargs
109
+ )
110
+ if not getattr(self.tokenizer, "is_fast", False):
111
+ raise ValueError("HuggingFaceTokenizer requires a fast tokenizer")
112
+ if self.tokenizer.pad_token_id is None:
113
+ raise ValueError("tokenizer must define pad_token_id")
114
+ self.pad_token_id = self.tokenizer.pad_token_id
115
+
116
+ def encode(
117
+ self, tokens: Sequence[str], max_length: int | None = None
118
+ ) -> TokenizedExample:
119
+ kwargs = {
120
+ "is_split_into_words": True,
121
+ "add_special_tokens": False,
122
+ "return_attention_mask": False,
123
+ }
124
+ if max_length is not None:
125
+ kwargs.update(truncation=True, max_length=max_length)
126
+ encoded = self.tokenizer(list(tokens), **kwargs)
127
+ try:
128
+ word_ids = encoded.word_ids()
129
+ except (AttributeError, ValueError) as error:
130
+ raise ValueError(
131
+ "tokenizer must provide word_ids; use a fast tokenizer"
132
+ ) from error
133
+ return TokenizedExample(
134
+ input_ids=encoded["input_ids"],
135
+ word_ids=word_ids,
136
+ )
137
+
138
+
139
+ class HighlightDataset(Dataset):
140
+ def __init__(self, examples: Iterable[HighlightExample]):
141
+ self.examples = list(examples)
142
+
143
+ def __len__(self) -> int:
144
+ return len(self.examples)
145
+
146
+ def __getitem__(self, index: int) -> HighlightExample:
147
+ return self.examples[index]
148
+
149
+
150
+ class HighlightCollator:
151
+ """Tokenize, align highlights to subtokens, and dynamically pad a batch."""
152
+
153
+ def __init__(
154
+ self,
155
+ tokenizer: HighlightTokenizer,
156
+ max_length: int | None = None,
157
+ ):
158
+ if max_length is not None and max_length < 1:
159
+ raise ValueError("max_length must be positive")
160
+ self.tokenizer = tokenizer
161
+ self.max_length = max_length
162
+
163
+ def __call__(self, examples: Sequence[HighlightExample]) -> InputData:
164
+ if not examples:
165
+ raise ValueError("cannot collate an empty batch")
166
+ encoded = [
167
+ self.tokenizer.encode(example.tokens, self.max_length)
168
+ for example in examples
169
+ ]
170
+ if self.max_length is not None:
171
+ encoded = [
172
+ TokenizedExample(
173
+ item.input_ids[: self.max_length],
174
+ item.word_ids[: self.max_length],
175
+ )
176
+ for item in encoded
177
+ ]
178
+ width = max(max(len(item.input_ids) for item in encoded), 1)
179
+
180
+ features = []
181
+ masks = []
182
+ highlights = []
183
+ for example, item in zip(examples, encoded):
184
+ input_ids = list(item.input_ids[:width])
185
+ word_ids = list(item.word_ids[:width])
186
+ for word_id in word_ids:
187
+ if word_id is not None and not 0 <= word_id < len(example.tokens):
188
+ raise ValueError("word_id is outside source-token range")
189
+
190
+ padding = width - len(input_ids)
191
+ features.append(input_ids + [self.tokenizer.pad_token_id] * padding)
192
+ masks.append(
193
+ [word_id is not None for word_id in word_ids] + [False] * padding
194
+ )
195
+ highlights.append(
196
+ [
197
+ -1
198
+ if word_id is None or example.highlights is None
199
+ else example.highlights[word_id]
200
+ for word_id in word_ids
201
+ ]
202
+ + [-1] * padding
203
+ )
204
+
205
+ return InputData(
206
+ features=th.tensor(features, dtype=th.long),
207
+ mask=th.tensor(masks, dtype=th.float32),
208
+ sample_ids=th.tensor([example.sample_id for example in examples]),
209
+ y_true=th.tensor([example.label for example in examples]),
210
+ highlight_true=th.tensor(highlights, dtype=th.long),
211
+ )
@@ -0,0 +1,21 @@
1
+ """Compatibility import path for data utilities."""
2
+
3
+ from pyhighlights.components.data import (
4
+ HighlightCollator,
5
+ HighlightDataset,
6
+ HighlightExample,
7
+ HighlightTokenizer,
8
+ HuggingFaceTokenizer,
9
+ TokenizedExample,
10
+ VocabularyTokenizer,
11
+ )
12
+
13
+ __all__ = [
14
+ "HighlightCollator",
15
+ "HighlightDataset",
16
+ "HighlightExample",
17
+ "HighlightTokenizer",
18
+ "HuggingFaceTokenizer",
19
+ "TokenizedExample",
20
+ "VocabularyTokenizer",
21
+ ]
@@ -0,0 +1,3 @@
1
+ from pyhighlights.components.models.data import InputData, OutputData, SPPOutput
2
+
3
+ __all__ = ["InputData", "OutputData", "SPPOutput"]
@@ -0,0 +1,178 @@
1
+ from __future__ import annotations
2
+
3
+ import abc
4
+ from typing import Dict, List, Literal, Tuple
5
+
6
+ import lightning as L
7
+ import torch as th
8
+ from cinnamon.registry import RegistrationKey, Registry
9
+ from torchmetrics import Metric, MetricCollection
10
+
11
+ from pyhighlights.components.models.data import (
12
+ InputData,
13
+ ModelData,
14
+ OutputData,
15
+ SPPOutput,
16
+ )
17
+ from pyhighlights.utility.losses import Loss, build_losses
18
+ from pyhighlights.utility.metrics import build_torchmetrics
19
+
20
+ Split = Literal["train", "val", "test"]
21
+
22
+ # Keep data containers importable from this module for compatibility.
23
+ __all__ = ["InputData", "Model", "ModelData", "OutputData", "SPPOutput", "Split"]
24
+
25
+
26
+ class Model(L.LightningModule, abc.ABC):
27
+ def __init__(
28
+ self,
29
+ name: str,
30
+ losses: List[RegistrationKey[Loss]],
31
+ optimizer: RegistrationKey[th.optim.Optimizer],
32
+ train_metrics: Dict[str, RegistrationKey[Metric]] | None = None,
33
+ val_metrics: Dict[str, RegistrationKey[Metric]] | None = None,
34
+ test_metrics: Dict[str, RegistrationKey[Metric]] | None = None,
35
+ ):
36
+ super().__init__()
37
+
38
+ self.save_hyperparameters(ignore=self.ignore_hyperparameters())
39
+ self.name = name
40
+ self.optimizer = optimizer
41
+
42
+ self.train_metrics = self._build_metrics(train_metrics)
43
+ self.val_metrics = self._build_metrics(val_metrics)
44
+ self.test_metrics = self._build_metrics(test_metrics)
45
+ self.losses = th.nn.ModuleList(build_losses(keys=losses))
46
+
47
+ self.store_predictions = False
48
+ self.predictions = []
49
+ self.forward_mapping = {
50
+ "train": self.training_forward,
51
+ "val": self.validation_forward,
52
+ "test": self.test_forward,
53
+ }
54
+
55
+ @staticmethod
56
+ def _build_metrics(
57
+ keys: Dict[str, RegistrationKey[Metric]] | None,
58
+ ) -> MetricCollection | None:
59
+ return build_torchmetrics(keys) if keys is not None else None
60
+
61
+ def ignore_hyperparameters(self) -> List[str]:
62
+ return []
63
+
64
+ def enable_storing_predictions(self):
65
+ self.store_predictions = True
66
+
67
+ def disable_storing_predictions(self):
68
+ self.store_predictions = False
69
+
70
+ def flush_predictions(self):
71
+ self.predictions.clear()
72
+
73
+ def update_metrics(
74
+ self, split: Split, input_data: InputData, output_data: OutputData
75
+ ):
76
+ metrics: MetricCollection | None = getattr(self, f"{split}_metrics")
77
+ if metrics is not None:
78
+ metrics.update(output_data.class_logits, input_data.y_true)
79
+
80
+ def compute_metrics(self, split: Split):
81
+ metrics: MetricCollection | None = getattr(self, f"{split}_metrics")
82
+ if metrics is None:
83
+ return
84
+
85
+ for key, value in metrics.compute().items():
86
+ self.log(f"{split}_{key}", value, prog_bar=True)
87
+ metrics.reset()
88
+
89
+ def on_train_epoch_end(self) -> None:
90
+ self.compute_metrics(split="train")
91
+
92
+ def on_validation_epoch_end(self) -> None:
93
+ self.compute_metrics(split="val")
94
+
95
+ def on_test_epoch_end(self) -> None:
96
+ self.compute_metrics(split="test")
97
+
98
+ def configure_optimizers(self):
99
+ return Registry.from_key(self.optimizer, params=self.parameters())
100
+
101
+ def log_metrics(
102
+ self,
103
+ split: Split,
104
+ total_loss: th.Tensor,
105
+ losses: Dict[str, th.Tensor],
106
+ batch_size: int,
107
+ ):
108
+ self.log(
109
+ name=f"{split}_loss",
110
+ value=total_loss,
111
+ on_step=False,
112
+ on_epoch=True,
113
+ prog_bar=True,
114
+ batch_size=batch_size,
115
+ )
116
+ for loss_name, loss_value in losses.items():
117
+ self.log(
118
+ name=f"{split}_{loss_name}",
119
+ value=loss_value,
120
+ on_step=False,
121
+ on_epoch=True,
122
+ prog_bar=True,
123
+ batch_size=batch_size,
124
+ )
125
+
126
+ def training_forward(self, batch: InputData) -> OutputData:
127
+ return self.forward(data=batch)
128
+
129
+ def validation_forward(self, batch: InputData) -> OutputData:
130
+ return self.training_forward(batch=batch)
131
+
132
+ def test_forward(self, batch: InputData) -> OutputData:
133
+ return self.training_forward(batch=batch)
134
+
135
+ def _step(self, batch: InputData, batch_idx: int, split: Split) -> th.Tensor:
136
+ output_data = self.forward_mapping[split](batch)
137
+ total_loss, losses = self.compute_loss(
138
+ input_data=batch, output_data=output_data
139
+ )
140
+
141
+ self.log_metrics(
142
+ split=split,
143
+ total_loss=total_loss,
144
+ losses=losses,
145
+ batch_size=batch.y_true.shape[0],
146
+ )
147
+ self.update_metrics(split=split, input_data=batch, output_data=output_data)
148
+
149
+ if self.store_predictions:
150
+ self.predictions.append({**batch.as_numpy(), **output_data.as_numpy()})
151
+
152
+ return total_loss
153
+
154
+ def training_step(self, batch: InputData, batch_idx: int):
155
+ return self._step(batch=batch, batch_idx=batch_idx, split="train")
156
+
157
+ def validation_step(self, batch: InputData, batch_idx: int):
158
+ return self._step(batch=batch, batch_idx=batch_idx, split="val")
159
+
160
+ def test_step(self, batch: InputData, batch_idx: int):
161
+ return self._step(batch=batch, batch_idx=batch_idx, split="test")
162
+
163
+ def compute_loss(
164
+ self,
165
+ input_data: InputData,
166
+ output_data: OutputData,
167
+ ) -> Tuple[th.Tensor, Dict[str, th.Tensor]]:
168
+ total_loss = output_data.class_logits.new_zeros(())
169
+ losses = {}
170
+
171
+ for loss in self.losses:
172
+ if not loss.enabled:
173
+ continue
174
+ loss_value = loss(input_data=input_data, output_data=output_data)
175
+ total_loss = total_loss + loss_value * loss.coefficient
176
+ losses[loss.name] = loss_value
177
+
178
+ return total_loss, losses