File size: 8,865 Bytes
1fbb475
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
import torch
from torch import nn
from typing import Optional, Union, Tuple

from transformers.modeling_outputs import SequenceClassifierOutput
from transformers.models.t5.modeling_t5 import T5Config, T5ClassificationHead, T5EncoderModel
from transformers.models.mt5.modeling_mt5 import MT5Config, MT5ClassificationHead, MT5EncoderModel
from transformers import AutoModelForSequenceClassification, AutoTokenizer, AutoConfig


def clean_t5_config(config: Union[T5Config, MT5Config], model_type: str):
    assert model_type in ['t5', 'mt5']
    setattr(config, 'pooling_mode', 'mean')
    setattr(config, 'model_type', model_type)
    setattr(config, 'use_cache', False)
    setattr(config, 'is_encoder_decoder', False)
    setattr(config, 'num_decoder_layers', None)
    delattr(config, 'task_specific_params') if hasattr(config, 'task_specific_params') else None


class T5EncoderForSequenceClassification(T5EncoderModel):
    """
    T5 encoder for sequence classification tasks.

    :param config: The T5 configuration object.
    """
    def __init__(self, config: T5Config):
        super().__init__(config)
        self.pool_layer = PoolLayer(config.pooling_mode)
        self.classification_head = T5ClassificationHead(config)

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple[torch.Tensor], SequenceClassifierOutput]:
        """
        Forward pass of the T5 encoder for sequence classification tasks.

        :param input_ids: The input token IDs.
        :param attention_mask: The attention mask.
        :param head_mask: The head mask.
        :param inputs_embeds: The input embeddings.
        :param labels: The target labels.
        :param output_attentions: Whether to output attentions.
        :param output_hidden_states: Whether to output hidden states.
        :param return_dict: Whether to return a dictionary.
        :returns: The logits for the classification task or a dictionary containing the outputs.
        """
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        loss = None

        outputs = self.encoder(
            input_ids=input_ids,
            attention_mask=attention_mask,
            inputs_embeds=inputs_embeds,
            head_mask=head_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        sequence_output = self.pool_layer(outputs.last_hidden_state, attention_mask)
        logits = self.classification_head(sequence_output)

        if not return_dict:
            output = (logits,) + outputs[2:]
            return ((loss,) + output) if loss is not None else output
        
        return SequenceClassifierOutput(
            loss=loss,
            logits=logits,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )


class MT5EncoderForSequenceClassification(MT5EncoderModel):
    """
    mT5 encoder for sequence classification tasks.

    :param config: The mT5 configuration object.
    """
    def __init__(self, config: MT5Config):
        super().__init__(config)
        self.pool_layer = PoolLayer(config.pooling_mode)
        self.classification_head = MT5ClassificationHead(config)

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple[torch.Tensor], SequenceClassifierOutput]:
        """
        Forward pass of the mT5 encoder for sequence classification tasks.

        :param input_ids: The input token IDs.
        :param attention_mask: The attention mask.
        :param head_mask: The head mask.
        :param inputs_embeds: The input embeddings.
        :param labels: The target labels.
        :param output_attentions: Whether to output attentions.
        :param output_hidden_states: Whether to output hidden states.
        :param return_dict: Whether to return a dictionary.
        :returns: The logits for the classification task or a dictionary containing the outputs.
        """
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        loss = None

        outputs = self.encoder(
            input_ids=input_ids,
            attention_mask=attention_mask,
            inputs_embeds=inputs_embeds,
            head_mask=head_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        sequence_output = self.pool_layer(outputs.last_hidden_state, attention_mask)
        logits = self.classification_head(sequence_output)

        if not return_dict:
            output = (logits,) + outputs[2:]
            return ((loss,) + output) if loss is not None else output
        
        return SequenceClassifierOutput(
            loss=loss,
            logits=logits,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )


class PoolLayer(nn.Module):
    """
    Pooling layer on top of the commputed token embeddings. 

    :param pooling_mode: The pooling strategy to use.
    """
    def __init__(self, pooling_mode: str):
        super().__init__()
        assert pooling_mode in ['first', 'mean', 'max'], f"ERROR: Unknown pooling strategy '{pooling_mode}'"
        self.pooling_mode = pooling_mode

    def forward(self, token_embeddings: torch.Tensor, attention_masks: torch.Tensor) -> torch.Tensor:
        """
        Compute the passage vector by pooling the token embeddings.

        :param token_embeddings: A 3D tensor of size [batch_size, seq_len, d_model].
        :param attention_masks: A 2D tensor of size [batch_size, seq_len].
        :returns: A 2D tensor of size [batch_size, d_model].
        """
        if self.pooling_mode == 'first':
            text_vectors = token_embeddings[:, 0, :]
        elif self.pooling_mode == 'max':
            # Set all values of the [PAD] embeddings to large negative values (so that they are never considered as maximum for a channel).
            attention_masks_expanded = attention_masks.unsqueeze(-1).expand(token_embeddings.size())
            token_embeddings[attention_masks_expanded == 0] = -1e+9 if token_embeddings.dtype == torch.float32 else -1e+4
            # Compute the maxima along the 'seq_length' dimension (-> Tensor[batch_size, d_model]).
            text_vectors = torch.max(token_embeddings, dim=1).values
        else:
            # Set all values of the [PAD] embeddings to zeros (so that they are not taken into account in the sum for a channel).
            attention_masks_expanded = attention_masks.unsqueeze(-1).expand(token_embeddings.size())
            token_embeddings[attention_masks_expanded == 0] = 0.0
            # Compute the means by first summing along the 'seq_length' dimension (-> Tensor[batch_size, d_model]).
            sum_embeddings = torch.sum(token_embeddings, dim=1)
            # Then, divide all values of a passage vector by the original passage length.
            sum_mask = attention_masks_expanded.sum(dim=1) # -> Tensor[batch_size, d_model] where each value is the length of the corresponding passage.
            sum_mask = torch.clamp(sum_mask, min=1e-7) # Make sure not to have zeros by lower bounding all elements to 1e-7.
            text_vectors = sum_embeddings / sum_mask # Divide each dimension by the sequence length.
        return text_vectors


if __name__ == "__main__":
    config = AutoConfig.from_pretrained(model_name)
    if isinstance(config, T5Config):
        clean_t5_config(self.config, model_type='t5')
        model = T5EncoderForSequenceClassification.from_pretrained(model_name, config=config)
    elif isinstance(config, MT5Config):
        clean_t5_config(self.config, model_type='t5')
        model = MT5EncoderForSequenceClassification.from_pretrained(model_name, config=config)
    else:
        model = AutoModelForSequenceClassification.from_pretrained(model_name, config=config)
    tokenizer = AutoTokenizer.from_pretrained(model_name, **tokenizer_args)