hugging face 预训练模型

发现一个很好用的预训练模型网站:https://huggingface.co,里面有超级多的预训练模型,常见的bert,robert,gpt,electra,xlnet等。

使用方法也很简单,https://huggingface.co/transformers/index.html。

以bert为例:

只需要下载https://huggingface.co中的预训练模型及配套的config,vocab等文件

加载时只需要两步

1.from transformers.models.bert.modeling_bert import BertForSequenceClassification,BertForTokenClassification
2.BertForTokenClassification.from_pretrained(pretrained_model_name_or_path=pretrain_model_path,config =self.model_config)

构建一个BERT-NER模型也很方便

class Bert_NER_Model(nn.Module):
    def __init__(self,pretrain_model_path):
        super(Bert_NER_Model, self).__init__()
        self.model_config = BertConfig.from_pretrained(pretrain_model_path)
        # 设定类别数
        self.model_config.num_labels = 63
        self.bert_model = BertForTokenClassification.from_pretrained(pretrained_model_name_or_path=pretrain_model_path,config =self.model_config)
   def forward(self,token_ids,type_token_ids,attention_mask_ids,labels):
       pretrain_features = self.bert_model(token_ids,type_token_ids,attention_mask_ids).logits
       loss = F.cross_entropy(pretrain_features,labels.flatten())
       predict = torch.argmax(pretrain_features,dim=1)
       return loss,predict

小伙伴们,赶快使用起来吧。

 

你可能感兴趣的:(AI,python,深度学习)