huggingface实现中文文本分类

目录

1  自定义数据集

2  分词 

 2.1  重写collate_fn方法¶

3  用BertModel加载预训练模型 

 4  模型试算

 5  定义下游任务¶

6  训练 

7  测试 


#导包
import torch
from datasets import load_from_disk  #用于加载本地磁盘的datasets文件

1  自定义数据集

#自定义数据集
#需要继承 torch.utils.data.Dataset,
#并且实现__init__(self)/__len__(self)/__getitem__(self,i)这些方法
class Dataset(torch.utils.data.Dataset):
    def __init__(self, split):
        #加载本地磁盘的datasets
        self.datasets = load_from_disk('../data/ChnSentiCorp')  #self.datasets是一个字典,包含训练、校验、测试的datatset
        self.dataset = self.datasets[split]  #使用split来区分获取的是训练、校验、测试的datatset中的哪一个
        
    def __len__(self):
        return len(self.dataset)
    
    def __getitem__(self, i):
        """让数据集像列表一样可以根据索引获取数据‘text’与“label”"""
        text = self.dataset[i]['text']
        label = self.dataset[i]['label']
        return text, label
    
dataset = Dataset(split='train')
dataset
<__main__.Dataset at 0x2afb31f03a0>
dataset.dataset
Dataset({
    features: ['text', 'label'],
    num_rows: 9600
})
len(dataset)

 9600

dataset[0]
('选择珠江花园的原因就是方便,有电动扶梯直接到达海边,周围餐馆、食廊、商场、超市、摊位一应俱全。酒店装修一般,但还算整洁。 泳池在大堂的屋顶,因此很小,不过女儿倒是喜欢。 包的早餐是西式的,还算丰富。 服务吗,一般',
 1)

2  分词 

#分词工具导包
from transformers import BertTokenizer
#加载字典和分词工具    huggingface自带的中文词典bert-base-chinese加载进来
tokenizer = BertTokenizer.from_pretrained(r'../data/bert-base-chinese/')
tokenizer
BertTokenizer(name_or_path='../data/bert-base-chinese/', vocab_size=21128, model_max_length=512, is_fast=False, padding_side='right', truncation_side='right', special_tokens={'unk_token': '[UNK]', 'sep_token': '[SEP]', 'pad_token': '[PAD]', 'cls_token': '[CLS]', 'mask_token': '[MASK]'}, clean_up_tokenization_spaces=True),  added_tokens_decoder={
	0: AddedToken("[PAD]", rstrip=False, lstrip