| 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 |