| 146 | |
| 147 | |
| 148 | class GLMForSequenceClassification(torch.nn.Module): |
| 149 | def __init__(self, language_model, hidden_size, hidden_dropout, pool_token, num_class=1): |
| 150 | super().__init__() |
| 151 | self.pool_token = pool_token |
| 152 | self.model = language_model |
| 153 | self.num_class = num_class |
| 154 | # Multi-choice head. |
| 155 | self.pool_layer = torch.nn.Linear(hidden_size, hidden_size) |
| 156 | self.multichoice_dropout = torch.nn.Dropout(hidden_dropout) |
| 157 | self.multichoice_head = torch.nn.Linear(hidden_size, num_class) |
| 158 | |
| 159 | def forward(self, input_ids, position_ids, attention_mask): |
| 160 | num_choices = None |
| 161 | if len(input_ids.shape) == 3: |
| 162 | assert self.num_class == 1 |
| 163 | batch_size, num_choices = input_ids.shape[:2] |
| 164 | input_ids = input_ids.reshape(-1, input_ids.size(-1)) |
| 165 | attention_mask = attention_mask.reshape(-1, *attention_mask.size()[2:]) |
| 166 | position_ids = position_ids.reshape(-1, *position_ids.size()[2:]) |
| 167 | outputs, *mems = self.model(input_ids, position_ids, attention_mask) |
| 168 | if self.pool_token == 'start': |
| 169 | output = outputs[ |
| 170 | torch.arange(outputs.size(0), dtype=attention_mask.dtype, device=attention_mask.device), attention_mask] |
| 171 | elif self.pool_token == 'pad': |
| 172 | output = outputs[torch.arange(outputs.size(0), dtype=attention_mask.dtype, |
| 173 | device=attention_mask.device), attention_mask - 1] |
| 174 | elif self.pool_token == 'cls': |
| 175 | output = outputs[:, 0] |
| 176 | else: |
| 177 | raise NotImplementedError |
| 178 | output = torch.tanh(self.pool_layer(output)) |
| 179 | multichoice_output = self.multichoice_dropout(output) |
| 180 | logits = self.multichoice_head(multichoice_output) |
| 181 | if num_choices is not None: |
| 182 | logits = logits.view(-1, num_choices) |
| 183 | return (logits, *mems) |