File size: 1,982 Bytes
f43af3c |
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 |
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
|