chendl/compositional_test

1import unicodedata2from dataclasses import dataclass3from typing import Optional, Union4 5import numpy as np6 7from transformers.data.data_collator import DataCollatorMixin8from transformers.file_utils import PaddingStrategy9from transformers.tokenization_utils_base import PreTrainedTokenizerBase10 11 12def padding_tensor(sequences, padding_value, padding_side, sequence_length):13    if isinstance(padding_value, tuple):14        out_tensor = np.full((len(sequences), sequence_length, 2), padding_value)15    else:16        out_tensor = np.full((len(sequences), sequence_length), padding_value)17 18    for i, tensor in enumerate(sequences):19        if padding_side == "right":20            if isinstance(padding_value, tuple):21                out_tensor[i, : len(tensor[:sequence_length]), :2] = tensor[:sequence_length]22            else:23                out_tensor[i, : len(tensor[:sequence_length])] = tensor[:sequence_length]24        else:25            if isinstance(padding_value, tuple):26                out_tensor[i, len(tensor[:sequence_length]) - 1 :, :2] = tensor[:sequence_length]27            else:28                out_tensor[i, len(tensor[:sequence_length]) - 1 :] = tensor[:sequence_length]29 30    return out_tensor.tolist()31 32 33def is_punctuation(char):34    cp = ord(char)35    if (cp >= 33 and cp <= 47) or (cp >= 58 and cp <= 64) or (cp >= 91 and cp <= 96) or (cp >= 123 and cp <= 126):36        return True37    cat = unicodedata.category(char)38    if cat.startswith("P"):39        return True40    return False41 42 43@dataclass44class DataCollatorForLukeTokenClassification(DataCollatorMixin):45    """46    Data collator that will dynamically pad the inputs received, as well as the labels.47 48    Args:49        tokenizer ([`PreTrainedTokenizer`] or [`PreTrainedTokenizerFast`]):50            The tokenizer used for encoding the data.51        padding (`bool`, `str` or [`~file_utils.PaddingStrategy`], *optional*, defaults to `True`):52            Select a strategy to pad the returned sequences (according to the model's padding side and padding index)53            among:54 55            - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single56              sequence if provided).57            - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the58              maximum acceptable input length for the model if that argument is not provided.59            - `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of60              different lengths).61        max_length (`int`, *optional*):62            Maximum length of the returned list and optionally padding length (see above).63        pad_to_multiple_of (`int`, *optional*):64            If set will pad the sequence to a multiple of the provided value.65 66            This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability >=67            7.5 (Volta).68        label_pad_token_id (`int`, *optional*, defaults to -100):69            The id to use when padding the labels (-100 will be automatically ignore by PyTorch loss functions).70        return_tensors (`str`):71            The type of Tensor to return. Allowable values are "np", "pt" and "tf".72    """73 74    tokenizer: PreTrainedTokenizerBase75    padding: Union[bool, str, PaddingStrategy] = True76    max_length: Optional[int] = None77    pad_to_multiple_of: Optional[int] = None78    label_pad_token_id: int = -10079    return_tensors: str = "pt"80 81    def torch_call(self, features):82        import torch83 84        label_name = "label" if "label" in features[0].keys() else "labels"85        labels = [feature[label_name] for feature in features] if label_name in features[0].keys() else None86        batch = self.tokenizer.pad(87            features,88            padding=self.padding,89            max_length=self.max_length,90            pad_to_multiple_of=self.pad_to_multiple_of,91            # Conversion to tensors will fail if we have labels as they are not of the same length yet.92            return_tensors="pt" if labels is None else None,93        )94 95        if labels is None:96            return batch97 98        sequence_length = torch.tensor(batch["entity_ids"]).shape[1]99        padding_side = self.tokenizer.padding_side100        if padding_side == "right":101            batch[label_name] = [102                list(label) + [self.label_pad_token_id] * (sequence_length - len(label)) for label in labels103            ]104        else:105            batch[label_name] = [106                [self.label_pad_token_id] * (sequence_length - len(label)) + list(label) for label in labels107            ]108 109        ner_tags = [feature["ner_tags"] for feature in features]110        batch["ner_tags"] = padding_tensor(ner_tags, -1, padding_side, sequence_length)111        original_entity_spans = [feature["original_entity_spans"] for feature in features]112        batch["original_entity_spans"] = padding_tensor(original_entity_spans, (-1, -1), padding_side, sequence_length)113        batch = {k: torch.tensor(v, dtype=torch.int64) for k, v in batch.items()}114 115        return batch116