| 125 | return loss,prob,labels |
| 126 | |
| 127 | class ModelWithCLRForClassification(ModelWithCLR): |
| 128 | def __init__(self, encoder, config, tokenizer, args): |
| 129 | super().__init__(encoder, config, tokenizer, args, clr_mask=False) |
| 130 | |
| 131 | def forward(self, input_ids=None, labels=None): |
| 132 | attention_mask = input_ids.ne(1) |
| 133 | outputs = self.encoder(input_ids,attention_mask) |
| 134 | cls_output = self.cls_head(outputs[0]) |
| 135 | logits=cls_output |
| 136 | prob=torch.sigmoid(logits) |
| 137 | return prob |
| 138 | |
| 139 | |
| 140 | class DefectModel(nn.Module): |