File size: 9,231 Bytes
d9115cb | 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 | from abc import ABC, abstractmethod
from typing import Optional
import torch
from src.data.loading.components.interfaces import (
LabelFunctionOutput,
)
class LabelFunction(ABC):
"""
An interface for the LabelFunction classes. The LabelFunction classes are used to transform the input sequence for training or inference and collecting the labels and label prediction locations.
"""
def __init__(self) -> None:
super().__init__()
@abstractmethod
def transform_label(self, sequence: torch.Tensor) -> LabelFunctionOutput:
"""
Function to transform the input sequence for training or inference and collecting the labels and label prediction locations.
The exact functionality needs to be implemented in the child class.
Parameters
----------
sequence: torch.Tensor of size (batch_size, sequence_length)
The matrix for input sequences.
Returns
----------
A LabelFunctionOutput object with the following attributes:
sequence: torch.Tensor of size (batch_size, new_sequence_length)
The original input sequence transformed for training or inference. The new_sequence_length can be the same as the original sequence length or different depending on the LabelFunction.
labels: torch.Tensor of size (number of labels,)
The labels across the sequences in the batch stacked together as labels for training.
label_location: torch.Tensor of size (number of masked tokens, 2)
The (row, col) location of the label_prediction corresponding to the labels in the batch stacked together.
"""
raise NotImplementedError("Need to implement in the child class.")
def get_input_attention_mask(
self, sequence: torch.Tensor, padding_token: int
) -> torch.Tensor:
"""
Function to get the input attention mask for the input sequence.
Defaults to returning the mask for the non-padding tokens but
can be overridden in the child class to return a different mask.
Parameters
----------
sequence: torch.Tensor of size (batch_size, sequence_length)
The matrix for input sequences.
padding_token: int
The index of token for padding.
Returns
----------
A tensor of size (batch_size, sequence_length) with 1s for non-padding tokens and 0s for padding tokens.
"""
return sequence != padding_token
class Identity(LabelFunction):
"""
LabelFunction class to return the original input for the non-masked values of the sequence. It's useful in situations where we don't to transform the input sequence.
"""
def transform_label(
self, sequence: torch.Tensor, padding_token: int = 0, masking_token: int = 0
) -> LabelFunctionOutput:
"""
Returns the original input for the non-masked values of the sequence.
Parameters
----------
sequence: torch.Tensor of size (batch_size, sequence_length)
The matrix for input sequences.
padding_token: int
The index of token for padding.
The padding_token is defined in the DataloaderConfig of LightningDataModule and is passed to the collate_fn which then passes it to the LabelFunction's transform_label function.
Thus, to change the padding_token, the user needs to change the padding_token in the DataloaderConfig.
Returns
----------
A LabelFunctionOutput object with the following attributes:
sequence: torch.Tensor of size (batch_size, sequence_length)
The original input sequence.
labels: torch.Tensor of size (number of non-padding tokens in batch,)
Collecting all the non-padding tokens in the batch stacked together as labels for training.
label_location: torch.Tensor of size (number of non-padding tokens in batch, 2)
The location of the label predictions, that is the (row, col) location of all non-padding tokens in the batch stacked together.
"""
content_mask = sequence != padding_token
labels = sequence[content_mask]
label_location = content_mask.nonzero()
return LabelFunctionOutput(
sequence=sequence, labels=labels, label_location=label_location
)
class NextKTokenMasking(LabelFunction):
"""
LabelFunction to create masking to use the last K tokens at the end of each sequence as labels.
Implements the `All Action` training objective from Pinnerformer (https://arxiv.org/pdf/2205.04507)
"""
def __init__(self, next_k: int = 5):
"""
Initialize the LabelFunction with the number of tokens to mask as labels.
Parameters
----------
next_k: int
Number of tokens K to mask as labels.
"""
self.next_k = next_k
def transform_label(
self, sequence: torch.Tensor, padding_token: int, masking_token: int
) -> LabelFunctionOutput:
"""
For each row of sequence, we save the original last next_k tokens as labels and replace them with 1 masking token and next_k - 1 padding tokens.
For all the next_k tokens, we use the first masked token as the label prediction.
Parameters
----------
sequence: torch.Tensor of size (batch_size, sequence_length)
The matrix for input sequences.
padding_token: int
The index of token for padding.
masking_token: int
The index of token for masked tokens.
The padding_token and the masking_token are defined in the DataloaderConfig of LightningDataModule and are passed to the collate_fn which then passes it to the LabelFunction's transform_label function.
Thus, to change the padding_token and the masking_token, the user needs to change them in the DataloaderConfig.
Returns
----------
A LabelFunctionOutput object with the following attributes:
sequence: torch.Tensor of size (batch_size, sequence_length),
The original input sequence transformed such that for each row, the last next_k non-padding tokens are replaced with 1 masking token and next_k - 1 padding tokens.
labels: torch.Tensor of size (batch_size * self.next_k,)
The original last next_k tokens of each row chosen as labels stacked together.
label_location: torch.Tensor of size (batch_size * self.next_k, 2)
The location of the label predictions, for each row the first masked token's location is returned as the label prediction location for all the next_k tokens.
"""
content_mask = sequence != padding_token
# check each sequence in batch is greater than next_k + 1
unpadded_seq_lengths = content_mask.sum(1) # shape: (batch_size,)
if torch.any(unpadded_seq_lengths < self.next_k + 1):
raise ValueError(
f"Sequence lengths: {unpadded_seq_lengths[unpadded_seq_lengths < self.next_k + 1]} should be greater than next_k + 1: {self.next_k + 1}"
)
# for each row, we want elements from unpadded_seq_lengths - next_k to unpadded_seq_lengths - 1 as labels
label_start_indices = unpadded_seq_lengths - self.next_k # shape: (batch_size,)
# for each row, we select [label_start_indices, label_start_indices + 1, ..., label_start_indices + next_k - 1]
label_col_offset = torch.arange(self.next_k) # shape: (next_k,)
label_col_indices = (
label_start_indices.unsqueeze(1) + label_col_offset
) # shape: (batch_size, next_k)
label_col_indices = label_col_indices.reshape(
-1
) # shape: (batch_size * next_k,)
# To get the row indices, we repeat each row index next_k times
row_orig_indices = torch.arange(sequence.size(0)) # shape: (batch_size,)
row_interleaved_indices = row_orig_indices.repeat_interleave(
self.next_k
) # shape: (batch_size * next_k,)
labels = sequence[row_interleaved_indices, label_col_indices]
# for prediction, for each row, we replace the last next_k tokens with
# 1 masking token at the label_start_indices, and next_k - 1 with padding tokens
sequence[row_interleaved_indices, label_col_indices] = padding_token
sequence[row_orig_indices, label_start_indices] = masking_token
# for each row, we use the label_start_index (which was masked) as label prediction for all the next_k labels
label_location = torch.stack(
(
row_interleaved_indices,
label_start_indices.repeat_interleave(self.next_k),
),
dim=1,
)
return LabelFunctionOutput(
sequence=sequence,
labels=labels,
label_location=label_location,
)
|