迁移学习xlm-roberta-base模型应用于分类任务

下载模型

加载模型

这里不使用官网推荐的方式

from transformers import AutoTokenizer, AutoModelForMaskedLM

tokenizer = AutoTokenizer.from_pretrained("xlm-roberta-base")

model = AutoModelForMaskedLM.from_pretrained("xlm-roberta-base")

使用普通的加载方式
https://huggingface.co/docs/transformers/model_doc/xlm-roberta#transformers.XLMRobertaTokenizer

from transformers import XLMRobertaTokenizer, XLMRobertaModel, BertConfig

设置Config类

class Config(object):
    def __init__(self, dataset):
    # 还有很多自己设定的config，我这里没写，每个人都不一样
    self.num_classes = n # 设置n分类
    self.hidden_size = 768
    self.model_path = "/***/xlm-roberta-base"
    self.tokenizer = XLMRobertaTokenizer.from_pretrained("xlm-roberta-base")
    self.bert_config = BertConfig.from_pretrained(self.model_path + '/config.json')
        

设置分类模型

class Model(nn.Module):
    def __init__(self, config):
        super(Model, self).__init__()
        self.bert = XLMRobertaModel.from_pretrained(config.model_path, config=config.bert_config)
        self.fc = nn.Linear(config.hidden_size, config.num_classes)


    def forward(self, x):
        context = x  # 输入的句子
        _, pooled = self.bert(context, attention_mask=mask,return_dict=False)
        # 如果电脑又GPU，要加上return_dict=False，否则不用加
        out = self.fc(pooled)
        return out

剩下的训练函数很简单，自己写吧