|
|
from dataclasses import dataclass |
|
|
from typing import Union, Optional |
|
|
|
|
|
from easy_tpp.preprocess.event_tokenizer import EventTokenizer |
|
|
from easy_tpp.utils import PaddingStrategy, TruncationStrategy |
|
|
|
|
|
|
|
|
@dataclass |
|
|
class TPPDataCollator: |
|
|
""" |
|
|
Data collator that will dynamically pad the inputs of event sequences. |
|
|
|
|
|
Args: |
|
|
tokenizer ([`EventTokenizer`]): |
|
|
The tokenizer used for encoding the data. |
|
|
padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `True`): |
|
|
Select a strategy to pad the returned sequences (according to the model's padding side and padding index) |
|
|
among: |
|
|
|
|
|
- `True` or `'longest'` (default): Pad to the longest sequence in the batch (or no padding if only a single |
|
|
sequence is provided). |
|
|
- `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum |
|
|
acceptable input length for the model if that argument is not provided. |
|
|
- `False` or `'do_not_pad'`: No padding (i.e., can output a batch with sequences of different lengths). |
|
|
max_length (`int`, *optional*): |
|
|
Maximum length of the returned list and optionally padding length (see above). |
|
|
return_tensors (`str`): |
|
|
The type of Tensor to return. Allowable values are "np", "pt" and "tf". |
|
|
""" |
|
|
|
|
|
tokenizer: EventTokenizer |
|
|
padding: Union[bool, str, PaddingStrategy] = True |
|
|
max_length: Optional[int] = None |
|
|
truncation: Union[bool, str, TruncationStrategy] = False |
|
|
return_tensors: str = "pt" |
|
|
|
|
|
def __call__(self, features, return_tensors=None): |
|
|
if return_tensors is None: |
|
|
return_tensors = self.return_tensors |
|
|
|
|
|
batch = self.tokenizer.pad( |
|
|
features, |
|
|
padding=self.padding, |
|
|
max_length=self.max_length, |
|
|
truncation=self.truncation, |
|
|
return_tensors=return_tensors, |
|
|
) |
|
|
|
|
|
return batch |
|
|
|
