HuggingFace实战（一）

视频链接：HuggingFace简明教程,BERT中文模型实战示例.NLP预训练模型,Transformers类库,datasets类库快速入门._哔哩哔哩_bilibili

文本分类

import torch
from datasets import load_from_disk
from transformers import BertTokenizer, BertModel
from torch.utils.data import Dataset, DataLoader
from torch.optim import AdamW
import os
from transformers import logging
 
# 设置transformers模块的日志等级，减少不必要的警告，对训练过程无影响，请忽略
logging.set_verbosity_error()

# 从本地加载数据集
dataset = load_from_disk('../datasets/ChnSentiCorp/data')

#定义数据集
class Dataset(Dataset):
    def __init__(self, dataset):
        self.dataset = dataset

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, i):
        text = self.dataset[i]['text']
        label = self.dataset[i]['label']

        return text, label


train = Dataset(dataset['train'])
validation = Dataset(dataset['validation']) # 不能叫test，会与test()重名，报错

#加载字典和分词工具
token = BertTokenizer.from_pretrained('bert-base-chinese')

# 定义批处理函数
def collate_fn(data):
    sents = [i[0] for i in data]
    labels = [i[1] for i in data]

    #编码
    data = token.batch_encode_plus(batch_text_or_text_pairs=sents,
                                   truncation=True,
                                   padding='max_length',
                                   max_length=500,
                                   return_tensors='pt',
                                   return_length=True)

    #input_ids:编码之后的数字
    #attention_mask:是补零的位置是0,其他位置是1
    input_ids = data['input_ids']
    attention_mask = data['attention_mask']
    token_type_ids = data['token_type_ids']
    labels = torch.LongTensor(labels)

    #print(data['length'], data['length'].max())

    return input_ids, attention_mask, token_type_ids, labels


#数据加载器
loader = DataLoader(dataset=train,
                    batch_size=16,
                    collate_fn=collate_fn,
                    shuffle=True,
                    drop_last=True)

for i, (input_ids, attention_mask, token_type_ids, labels) in enumerate(loader):
    break

# print(len(loader)) # 600
# print(input_ids.shape, attention_mask.shape, token_type_ids.shape, labels) # torch.Size([16, 500]) torch.Size([16, 500]) torch.Size([16, 500]) tensor([0, 1, 0, 0, 1, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0, 1])

#加载预训练模型
pretrained = BertModel.from_pretrained('bert-base-chinese')

#不训练,不需要计算梯度
for param in pretrained.parameters():
    param.requires_grad_(False)

#模型试算
# out = pretrained(input_ids=input_ids,attention_mask=attention_mask, token_type_ids=token_type_ids)

# print(out.last_hidden_state.shape) # torch.Size([16, 500, 768])

#定义下游任务模型
class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = torch.nn.Linear(768, 2) # 预训练模型的隐层输出维度是768

    def forward(self, input_ids, attention_mask, token_type_ids):
        with torch.no_grad(): # 不改变预训练模型参数
            out = pretrained(input_ids=input_ids,
                       attention_mask=attention_mask,
                       token_type_ids=token_type_ids)

        out = self.fc(out.last_hidden_state[:, 0]) # 解决分类问题只需要使用到[CLS]，即第一个字符

        out = out.softmax(dim=