(self, input_ids, position_ids, attention_mask)
| 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) |
nothing calls this directly
no outgoing calls
no test coverage detected