(self, encoder, config, tokenizer, args)
| 204 | |
| 205 | class DecoderClassifier(nn.Module): |
| 206 | def __init__(self, encoder, config, tokenizer, args): |
| 207 | super(DecoderClassifier, self).__init__() |
| 208 | self.encoder = encoder |
| 209 | self.config=config |
| 210 | self.tokenizer=tokenizer |
| 211 | self.args=args |
| 212 | self.classifier = nn.Linear(config.hidden_size, 2) |
| 213 | |
| 214 | def forward(self, input_ids=None, labels=None, weight=None): |
| 215 | attention_mask = input_ids.ne(self.tokenizer.pad_token_id) |