File size: 15,963 Bytes
626eca0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
import collections
import contextlib
import logging
from typing import Any, Dict, Iterator, List

import torch
import transformers as tr
from lightning_fabric.utilities import move_data_to_device
from torch.utils.data import DataLoader, IterableDataset
from tqdm import tqdm

from relik.common.log import get_console_logger, get_logger
from relik.common.utils import get_callable_from_string
from relik.reader.data.relik_reader_sample import RelikReaderSample
from relik.reader.pytorch_modules.base import RelikReaderBase
from relik.reader.utils.special_symbols import get_special_symbols
from relik.retriever.pytorch_modules import PRECISION_MAP

console_logger = get_console_logger()
logger = get_logger(__name__, level=logging.INFO)


class RelikReaderForSpanExtraction(RelikReaderBase):
    """
    A class for the RelikReader model for span extraction.

    Args:
        transformer_model (:obj:`str` or :obj:`transformers.PreTrainedModel` or :obj:`None`, `optional`):
            The transformer model to use. If `None`, the default model is used.
        additional_special_symbols (:obj:`int`, `optional`, defaults to 0):
            The number of additional special symbols to add to the tokenizer.
        num_layers (:obj:`int`, `optional`):
            The number of layers to use. If `None`, all layers are used.
        activation (:obj:`str`, `optional`, defaults to "gelu"):
            The activation function to use.
        linears_hidden_size (:obj:`int`, `optional`, defaults to 512):
            The hidden size of the linears.
        use_last_k_layers (:obj:`int`, `optional`, defaults to 1):
            The number of last layers to use.
        training (:obj:`bool`, `optional`, defaults to False):
            Whether the model is in training mode.
        device (:obj:`str` or :obj:`torch.device` or :obj:`None`, `optional`):
            The device to use. If `None`, the default device is used.
        tokenizer (:obj:`str` or :obj:`transformers.PreTrainedTokenizer` or :obj:`None`, `optional`):
            The tokenizer to use. If `None`, the default tokenizer is used.
        dataset (:obj:`IterableDataset` or :obj:`str` or :obj:`None`, `optional`):
            The dataset to use. If `None`, the default dataset is used.
        dataset_kwargs (:obj:`Dict[str, Any]` or :obj:`None`, `optional`):
            The keyword arguments to pass to the dataset class.
        default_reader_class (:obj:`str` or :obj:`transformers.PreTrainedModel` or :obj:`None`, `optional`):
            The default reader class to use. If `None`, the default reader class is used.
        **kwargs:
            Keyword arguments.
    """

    default_reader_class: str = (
        "relik.reader.pytorch_modules.hf.modeling_relik.RelikReaderSpanModel"
    )
    default_data_class: str = "relik.reader.data.relik_reader_data.RelikDataset"

    def __init__(
        self,
        transformer_model: str | tr.PreTrainedModel | None = None,
        additional_special_symbols: int = 0,
        num_layers: int | None = None,
        activation: str = "gelu",
        linears_hidden_size: int | None = 512,
        use_last_k_layers: int = 1,
        training: bool = False,
        device: str | torch.device | None = None,
        tokenizer: str | tr.PreTrainedTokenizer | None = None,
        dataset: IterableDataset | str | None = None,
        dataset_kwargs: Dict[str, Any] | None = None,
        default_reader_class: tr.PreTrainedModel | str | None = None,
        **kwargs,
    ):
        super().__init__(
            transformer_model=transformer_model,
            additional_special_symbols=additional_special_symbols,
            num_layers=num_layers,
            activation=activation,
            linears_hidden_size=linears_hidden_size,
            use_last_k_layers=use_last_k_layers,
            training=training,
            device=device,
            tokenizer=tokenizer,
            dataset=dataset,
            default_reader_class=default_reader_class,
            **kwargs,
        )
        # and instantiate the dataset class
        self.dataset = dataset
        if self.dataset is None:
            default_data_kwargs = dict(
                dataset_path=None,
                materialize_samples=False,
                transformer_model=self.tokenizer,
                special_symbols=get_special_symbols(
                    self.relik_reader_model.config.additional_special_symbols
                ),
                for_inference=True,
            )
            # merge the default data kwargs with the ones passed to the model
            default_data_kwargs.update(dataset_kwargs or {})
            self.dataset = get_callable_from_string(self.default_data_class)(
                **default_data_kwargs
            )

    @torch.no_grad()
    @torch.inference_mode()
    def _read(
        self,
        samples: List[RelikReaderSample] | None = None,
        input_ids: torch.Tensor | None = None,
        attention_mask: torch.Tensor | None = None,
        token_type_ids: torch.Tensor | None = None,
        prediction_mask: torch.Tensor | None = None,
        special_symbols_mask: torch.Tensor | None = None,
        max_length: int = 1000,
        max_batch_size: int = 128,
        token_batch_size: int = 2048,
        precision: str = 32,
        annotation_type: str = "char",
        progress_bar: bool = False,
        *args: object,
        **kwargs: object,
    ) -> List[RelikReaderSample] | List[List[RelikReaderSample]]:
        """
        A wrapper around the forward method that returns the predicted labels for each sample.

        Args:
            samples (:obj:`List[RelikReaderSample]`, `optional`):
                The samples to read. If provided, `text` and `candidates` are ignored.
            input_ids (:obj:`torch.Tensor`, `optional`):
                The input ids of the text. If `samples` is provided, this is ignored.
            attention_mask (:obj:`torch.Tensor`, `optional`):
                The attention mask of the text. If `samples` is provided, this is ignored.
            token_type_ids (:obj:`torch.Tensor`, `optional`):
                The token type ids of the text. If `samples` is provided, this is ignored.
            prediction_mask (:obj:`torch.Tensor`, `optional`):
                The prediction mask of the text. If `samples` is provided, this is ignored.
            special_symbols_mask (:obj:`torch.Tensor`, `optional`):
                The special symbols mask of the text. If `samples` is provided, this is ignored.
            max_length (:obj:`int`, `optional`, defaults to 1000):
                The maximum length of the text.
            max_batch_size (:obj:`int`, `optional`, defaults to 128):
                The maximum batch size.
            token_batch_size (:obj:`int`, `optional`):
                The token batch size.
            progress_bar (:obj:`bool`, `optional`, defaults to False):
                Whether to show a progress bar.
            precision (:obj:`str`, `optional`, defaults to 32):
                The precision to use for the model.
            annotation_type (:obj:`str`, `optional`, defaults to "char"):
                The annotation type to use. It can be either "char", "token" or "word".
            *args:
                Positional arguments.
            **kwargs:
                Keyword arguments.

        Returns:
            :obj:`List[RelikReaderSample]` or :obj:`List[List[RelikReaderSample]]`:
                The predicted labels for each sample.
        """

        precision = precision or self.precision
        if samples is not None:

            def _read_iterator():
                def samples_it():
                    for i, sample in enumerate(samples):
                        assert sample._mixin_prediction_position is None
                        sample._mixin_prediction_position = i
                        yield sample

                next_prediction_position = 0
                position2predicted_sample = {}

                # instantiate dataset
                if self.dataset is None:
                    raise ValueError(
                        "You need to pass a dataset to the model in order to predict"
                    )
                self.dataset.samples = samples_it()
                self.dataset.model_max_length = max_length
                self.dataset.tokens_per_batch = token_batch_size
                self.dataset.max_batch_size = max_batch_size

                # instantiate dataloader
                iterator = DataLoader(
                    self.dataset, batch_size=None, num_workers=0, shuffle=False
                )
                if progress_bar:
                    iterator = tqdm(iterator, desc="Predicting with RelikReader")

                # fucking autocast only wants pure strings like 'cpu' or 'cuda'
                # we need to convert the model device to that
                device_type_for_autocast = str(self.device).split(":")[0]
                # autocast doesn't work with CPU and stuff different from bfloat16
                autocast_mngr = (
                    contextlib.nullcontext()
                    if device_type_for_autocast == "cpu"
                    else (
                        torch.autocast(
                            device_type=device_type_for_autocast,
                            dtype=PRECISION_MAP[precision],
                        )
                    )
                )

                with autocast_mngr:
                    for batch in iterator:
                        batch = move_data_to_device(batch, self.device)
                        batch_out = self._batch_predict(**batch)

                        for sample in batch_out:
                            if (
                                sample._mixin_prediction_position
                                >= next_prediction_position
                            ):
                                position2predicted_sample[
                                    sample._mixin_prediction_position
                                ] = sample

                        # yield
                        while next_prediction_position in position2predicted_sample:
                            yield position2predicted_sample[next_prediction_position]
                            del position2predicted_sample[next_prediction_position]
                            next_prediction_position += 1

            outputs = list(_read_iterator())
            for sample in outputs:
                self.dataset.merge_patches_predictions(sample)
                self.dataset.convert_tokens_to_char_annotations(sample)

        else:
            outputs = list(
                self._batch_predict(
                    input_ids,
                    attention_mask,
                    token_type_ids,
                    prediction_mask,
                    special_symbols_mask,
                    *args,
                    **kwargs,
                )
            )
        return outputs

    def _batch_predict(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor,
        token_type_ids: torch.Tensor | None = None,
        prediction_mask: torch.Tensor | None = None,
        special_symbols_mask: torch.Tensor | None = None,
        sample: List[RelikReaderSample] | None = None,
        top_k: int = 5,  # the amount of top-k most probable entities to predict
        *args,
        **kwargs,
    ) -> Iterator[RelikReaderSample]:
        """
        A wrapper around the forward method that returns the predicted labels for each sample.
        It also adds the predicted labels to the samples.

        Args:
            input_ids (:obj:`torch.Tensor`):
                The input ids of the text.
            attention_mask (:obj:`torch.Tensor`):
                The attention mask of the text.
            token_type_ids (:obj:`torch.Tensor`, `optional`):
                The token type ids of the text.
            prediction_mask (:obj:`torch.Tensor`, `optional`):
                The prediction mask of the text.
            special_symbols_mask (:obj:`torch.Tensor`, `optional`):
                The special symbols mask of the text.
            sample (:obj:`List[RelikReaderSample]`, `optional`):
                The samples to read. If provided, `text` and `candidates` are ignored.
            top_k (:obj:`int`, `optional`, defaults to 5):
                The amount of top-k most probable entities to predict.
            *args:
                Positional arguments.
            **kwargs:
                Keyword arguments.

        Returns:
            The predicted labels for each sample.
        """
        forward_output = self.forward(
            input_ids=input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            prediction_mask=prediction_mask,
            special_symbols_mask=special_symbols_mask,
        )

        ned_start_predictions = forward_output["ned_start_predictions"].cpu().numpy()
        ned_end_predictions = forward_output["ned_end_predictions"].cpu().numpy()
        ed_predictions = forward_output["ed_predictions"].cpu().numpy()
        ed_probabilities = forward_output["ed_probabilities"].cpu().numpy()

        batch_predictable_candidates = kwargs["predictable_candidates"]
        patch_offset = kwargs["patch_offset"]
        for ts, ne_sp, ne_ep, edp, edpr, pred_cands, po in zip(
            sample,
            ned_start_predictions,
            ned_end_predictions,
            ed_predictions,
            ed_probabilities,
            batch_predictable_candidates,
            patch_offset,
        ):
            ne_start_indices = [ti for ti, c in enumerate(ne_sp[1:]) if c > 0]
            ne_end_indices = [ti for ti, c in enumerate(ne_ep[1:]) if c > 0]

            final_class2predicted_spans = collections.defaultdict(list)
            spans2predicted_probabilities = dict()
            for start_token_index, end_token_index in zip(
                ne_start_indices, ne_end_indices
            ):
                # predicted candidate
                token_class = edp[start_token_index + 1] - 1
                predicted_candidate_title = pred_cands[token_class]
                final_class2predicted_spans[predicted_candidate_title].append(
                    [start_token_index, end_token_index]
                )

                # candidates probabilities
                classes_probabilities = edpr[start_token_index + 1]
                classes_probabilities_best_indices = classes_probabilities.argsort()[
                    ::-1
                ]
                titles_2_probs = []
                top_k = (
                    min(
                        top_k,
                        len(classes_probabilities_best_indices),
                    )
                    if top_k != -1
                    else len(classes_probabilities_best_indices)
                )
                for i in range(top_k):
                    titles_2_probs.append(
                        (
                            pred_cands[classes_probabilities_best_indices[i] - 1],
                            classes_probabilities[
                                classes_probabilities_best_indices[i]
                            ].item(),
                        )
                    )
                spans2predicted_probabilities[
                    (start_token_index, end_token_index)
                ] = titles_2_probs

            if "patches" not in ts._d:
                ts._d["patches"] = dict()

            ts._d["patches"][po] = dict()
            sample_patch = ts._d["patches"][po]

            sample_patch["predicted_window_labels"] = final_class2predicted_spans
            sample_patch["span_title_probabilities"] = spans2predicted_probabilities

            # additional info
            sample_patch["predictable_candidates"] = pred_cands

            yield ts