| 1329 | """ |
| 1330 | |
| 1331 | def __init__(self, config, num_labels=2): |
| 1332 | super(BertForTokenClassification, self).__init__(config) |
| 1333 | self.num_labels = num_labels |
| 1334 | self.bert = BertModel(config) |
| 1335 | self.dropout = nn.Dropout(config.hidden_dropout_prob) |
| 1336 | self.classifier = nn.Linear(config.hidden_size, num_labels) |
| 1337 | # self.classifier = mpu.RowParallelLinear( |
| 1338 | # input_size=config.hidden_size, |
| 1339 | # output_size=num_labels, |
| 1340 | # bias=True, |
| 1341 | # input_is_parallel=True, |
| 1342 | # stride=1, |
| 1343 | # init_method=normal_init_method(mean=0.0, |
| 1344 | # std=config.initializer_range)) |
| 1345 | self.apply(self.init_bert_weights) |
| 1346 | |
| 1347 | def forward(self, input_ids, token_type_ids=None, attention_mask=None, labels=None, checkpoint_activations=False): |
| 1348 | sequence_output, _ = self.bert(input_ids, token_type_ids, attention_mask, output_all_encoded_layers=False, |