发现一个很好用的预训练模型网站: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
小伙伴们,赶快使用起来吧。