jupyter快速实现单标签及多标签多分类的文本分类BERT模型

最新推荐文章于 2024-06-03 16:41:20 发布

xiyou__

最新推荐文章于 2024-06-03 16:41:20 发布

阅读量869

点赞数 9

分类专栏：模型复现文章标签： jupyter 分类 bert nlp

本文链接：https://blog.csdn.net/xiyou__/article/details/134603753

版权

模型复现专栏收录该内容

5 篇文章 1 订阅

订阅专栏

jupyter实现pytorch版BERT（单标签分类版）

nlp-notebooks/Text classification with BERT in PyTorch.ipynb

通过改写上述代码，实现多标签分类

参考解决方案，我选择的解决方案是继承BertForSequenceClassification并改写，即将上述代码的ln [9] 改为以下内容：

from transformers.modeling_bert import BertForSequenceClassification
from transformers.modeling_outputs import SequenceClassifierOutput

class BertForMultilabelSequenceClassification(BertForSequenceClassification):
   def __init__(self, config):
     super().__init__(config)

   def forward(self,
       input_ids=None,
       attention_mask=None,
       token_type_ids=None,
       position_ids=None,
       head_mask=None,
       inputs_embeds=None,
       labels=None,
       output_attentions=None,
       output_hidden_states=None,
       return_dict=None):
       return_dict = return_dict if return_dict is not None else self.config.use_return_dict

       outputs = self.bert(input_ids,
           attention_mask=attention_mask,
           token_type_ids=token_type_ids,
           position_ids=position_ids,
           head_mask=head_mask,
           inputs_embeds=inputs_embeds,
           output_attentions=output_attentions,
           output_hidden_states=output_hidden_states,
           return_dict=return_dict)

       pooled_output = outputs[1]
       pooled_output = self.dropout(pooled_output)
       logits = self.classifier(pooled_output)

       loss = None
       if labels is not None:
           loss_fct = torch.nn.BCEWithLogitsLoss()
           loss = loss_fct(logits.view(-1, self.num_labels), 
                           labels.float().view(-1, self.num_labels))

       if not return_dict:
           output = (logits,) + outputs[2:]
           return ((loss,) + output) if loss is not None else output

       return SequenceClassifierOutput(loss=loss,
           logits=logits,
           hidden_states=outputs.hidden_states,
           attentions=outputs.attentions)
           
model = BertForMultilabelSequenceClassification.from_pretrained(BERT_MODEL, num_labels = len(label2idx))
model.to(device)