Transformers基本组件（二）快速入门Datasets、Evaluate、Trainer

Transformers基本组件（二）快速入门Datasets、Evaluate、Trainer

1、基础组件Datasets

数据集部分的工作，一部分在于数据集的收集，另一部分在于数据集的处理。Datasets库的出现，一定程度上也使得这两部分的工作变得简单了许多。

1.1 加载数据集

加载在线数据集

from datasets import load_dataset

# 加载公开数据集只需要两步，导入Datasets包，然后加载想要的数据集即可
# 经过短暂的下载后（当然，大概率会出现443错误），便可以看到数据集已经被成功加载。
# 而且已经被划分为了训练集和验证集两部分，训练集5850条，验证集1679条。


# 官网中列出了公开数据集，点击具体的数据集后，可以查看详细信息。
datasets = load_dataset("madao33/new-title-chinese")
datasets

DatasetDict({
    train: Dataset({
        features: ['title', 'content'],
        num_rows: 5850
    })
    validation: Dataset({
        features: ['title', 'content'],
        num_rows: 1679
    })
})

加载子数据集

super_glue是一系列任务的合集，里面包含多个数据集，假设要加载boolq数据集，要如何做呢

# 直接加子数据集参数
boolq_dataset = load_dataset("super_glue", "boolq")
boolq_dataset

DatasetDict({
    train: Dataset({
        features: ['question', 'passage', 'idx', 'label'],
        num_rows: 9427
    })
    validation: Dataset({
        features: ['question', 'passage', 'idx', 'label'],
        num_rows: 3270
    })
    test: Dataset({
        features: ['question', 'passage', 'idx', 'label'],
        num_rows: 3245
    })
})

按照数据集划分进行加载

# 前面加载的数据集都是将全部数据集加载了，包括训练集、验证集、测试集。
# 我们也可以根据数据集的划分，选择要加载的数据集划分，只需要指定split参数。
# 假设我们要加载前面中文新闻数据集中的训练集，那么代码可以这样：

dataset = load_dataset("madao33/new-title-chinese", split="train")
dataset

Dataset({
    features: ['title', 'content'],
    num_rows: 5850
})

# 其他形式的加载

# 按照条数加载
dataset = load_dataset("madao33/new-title-chinese", split="train[10:100]")

# 按照比例加载
dataset = load_dataset("madao33/new-title-chinese", split="train[:50%]")

1.2 数据集常用操作

1.2.1 查看数据集

datasets = load_dataset("madao33/new-title-chinese")

print(datasets["train"][0])
print(datasets["train"]["title"][:5])
print(datasets["train"].column_names)
print(datasets["train"].features)

1.2.2 数据集划分

dataset = datasets["train"]

# 这里我们对原始的train数据集进行了划分，可以看到数据按照9:1的比例重新进行了划分。
dataset.train_test_split(test_size=0.1)

DatasetDict({
    train: Dataset({
        features: ['title', 'content'],
        num_rows: 5265
    })
    test: Dataset({
        features: ['title', 'content'],
        num_rows: 585
    })
})

1.2.3 数据集的选取与过滤

# 选取
datasets["train"].select([0, 1])

Dataset({
    features: ['title', 'content'],
    num_rows: 2
})

# 过滤
filter_dataset = datasets["train"].filter(lambda example: "中国" in example["title"])
filter_dataset["title"][:5]

['聚焦两会，世界探寻中国成功秘诀',
 '望海楼中国经济的信心来自哪里',
 '“中国奇迹”助力世界减贫跑出加速度',
 '和音瞩目历史交汇点上的中国',
 '中国风采感染世界']

1.2.4 数据映射

def add_prefix(example):
    example["title"] = 'Prefix: ' + example["title"]
    return example


prefix_dataset = datasets.map(add_prefix)
prefix_dataset["train"][:10]["title"]

['Prefix: 望海楼美国打“台湾牌”是危险的赌博',
 'Prefix: 大力推进高校治理能力建设',
 'Prefix: 坚持事业为上选贤任能',
 'Prefix: “大朋友”的话儿记心头',
 'Prefix: 用好可持续发展这把“金钥匙”',
 'Prefix: 跨越雄关，我们走在大路上',
 'Prefix: 脱贫奇迹彰显政治优势',
 'Prefix: 拱卫亿万人共同的绿色梦想',
 'Prefix: 为党育人、为国育才',
 'Prefix: 净化网络语言']

当然，真实情况下，我们不会只做这样的处理。

在这里，一般会与Tokenizer做结合实现完整的数据处理函数，将其直接转换为模型想要的输入。

from transformers import AutoTokenizer


tokenizer = AutoTokenizer.from_pretrained("bert-base-chinese")

def preprocess_function(example, tokenizer=tokenizer):
    model_inputs = tokenizer(example["content"], max_length=512, truncation=True)
    labels = tokenizer(example["title"], max_length=32, truncation=True)
    # label就是title编码的结果
    model_inputs["labels"] = labels["input_ids"]
    return model_inputs

# 处理完成后，在数据的字段中便多了四个字段，这些字段便是模型能够处理的字段。
processed_datasets = datasets.map(preprocess_function)
processed_datasets

DatasetDict({
    train: Dataset({
        features: ['title', 'content', 'input_ids', 'token_type_ids', 'attention_mask', 'labels'],
        num_rows: 5850
    })
    validation: Dataset({
        features: ['title', 'content', 'input_ids', 'token_type_ids', 'attention_mask', 'labels'],
        num_rows: 1679
    })
})

map函数还支持批量处理数据，只需要指定batched参数值为True。

# 二者效果是一样的，但是使用批量数据处理效率要高的多。
processed_datasets = datasets.map(preprocess_function, batched=True)
processed_datasets

当然，我们一般只需要处理后的字段

processed_datasets = datasets.map(preprocess_function, batched=True, remove_columns=datasets["train"].column_names)


# 可以发现原始字段已经删除，只保留了处理后的字段
processed_datasets

DatasetDict({
    train: Dataset({
        features: ['input_ids', 'token_type_ids', 'attention_mask', 'labels'],
        num_rows: 5850
    })
    validation: Dataset({
        features: ['input_ids', 'token_type_ids', 'attention_mask', 'labels'],
        num_rows: 1679
    })
})

1.2.5 数据保存和加载

# 当我们处理完数据之后，我们可以将其保存到本地磁盘，需要用时下次直接加载即可，无需二次处理。
processed_datasets.save_to_disk("./processed_data")

processed_datasets = load_from_disk("./processed_data")
processed_datasets

DatasetDict({
    train: Dataset({
        features: ['input_ids', 'token_type_ids', 'attention_mask', 'labels'],
        num_rows: 5850
    })
    validation: Dataset({
        features: ['input_ids', 'token_type_ids', 'attention_mask', 'labels'],
        num_rows: 1679
    })
})

1.3 加载本地数据集

多数情况下，公开数据集并不能满足我们的需求，需要加载自行准备的数据集。下面来介绍如何加载本地的数据集。

1.3.1 直接加载文件

# 本地文件支持csv、json文件的直接加载
# 1、加载csv文件
dataset = load_dataset("csv", data_files="./ChnSentiCorp_htl_all.csv", split="train")

# 或者
from datasets import Dataset
dataset = Dataset.from_csv("./ChnSentiCorp_htl_all.csv")

# 2、加载json文件
load_dataset("json", data_files=["./cmrc2018_trial.json"], field="data")

1.3.2 加载文件夹中的文件

# 加载整个文件夹的内容
dataset = load_dataset("csv", data_dir='./all_data/', split='train')


# 加载文件夹中特定的文件
dataset = load_dataset("csv", data_files=["./all_data/ChnSentiCorp_htl_all.csv", "./all_data/ChnSentiCorp_htl_all_2.csv"], split='train')

dataset

Dataset({
    features: ['label', 'review'],
    num_rows: 15532
})

1.3.3 转换加载数据集

import pandas as pd

# 1、从pandas中加载
data = pd.read_csv("./ChnSentiCorp_htl_all.csv")

dataset = Dataset.from_pandas(data)

# 2、从list中加载

# List格式的数据需要内嵌{}，明确数据字段
data = [{"text": "abc"}, {"text": "def"}]
# data = ["abc", "def"]
Dataset.from_list(data)

1.3.4 利用脚本加载

data中数据如下

      "paragraphs": [
        {
          "id": "TRIAL_800", 
          "context": "基于《跑跑卡丁车》与《泡泡堂》上所开发的游戏，由韩国Nexon开发与发行。中国大陆由盛大游戏运营，这是Nexon时隔6年再次授予盛大网络其游戏运营权。台湾由游戏橘子运营。......", 
          "qas": [
            {
              "question": "生命数耗完即算为什么？", 
              "id": "TRIAL_800_QUERY_0", 
              "answers": [
                {
                  "text": "踢爆", 
                  "answer_start": 127
                }
              ]
            }
          ]
        }
      ], 
      "id": "TRIAL_800", 
      "title": "泡泡战士"

# 直接解析josn不能满足要求，只解析出'paragraphs', 'id', 'title'字段
load_dataset("json", data_files="./cmrc2018_trial.json", field="data")

DatasetDict({
    train: Dataset({
        features: ['paragraphs', 'id', 'title'],
        num_rows: 256
    })
})

使用脚本进行加载

load_dataset("./load_cmrc.py", split="train")

Dataset({
    features: ['id', 'context', 'question', 'answers'],
    num_rows: 1002
})

脚本如下

import json
import datasets
from datasets import DownloadManager, DatasetInfo

'''
_info方法定义数据集结构
_split_generators方法指定不同数据划分要加载的数据
_generate_examples方法实现具体的读取数据代码
'''
class CMRC2018TRIAL(datasets.GeneratorBasedBuilder):

    def _info(self) -> DatasetInfo:
        """
            info方法，定义数据集的信息，这里要对数据的字段进行定义
        :return:
        """
        return datasets.DatasetInfo(
            description="CMRC2018 trial",
            features=datasets.Features({
                    "id": datasets.Value("string"),
                    "context": datasets.Value("string"),
                    "question": datasets.Value("string"),
                    "answers": datasets.features.Sequence(
                        {
                            "text": datasets.Value("string"),
                            "answer_start": datasets.Value("int32"),
                        }
                    )
                }),
        )

    def _split_generators(self, dl_manager: DownloadManager):
        """
            返回datasets.SplitGenerator
            涉及两个参数：name和gen_kwargs
            name: 指定数据集的划分
            gen_kwargs: 指定要读取的文件的路径，与_generate_examples的入参数一致
        :param dl_manager:
        :return: [ datasets.SplitGenerator ]
        """
        return [datasets.SplitGenerator(name=datasets.Split.TRAIN, gen_kwargs={"filepath": "./cmrc2018_trial.json"})]

    def _generate_examples(self, filepath):
        """
            生成具体的样本，使用yield
            需要额外指定key，id从0开始自增就可以
        :param filepath:
        :return:
        """
        # Yields (key, example) tuples from the dataset
        with open(filepath, encoding="utf-8") as f:
            data = json.load(f)
            for example in data["data"]:
                for paragraph in example["paragraphs"]:
                    context = paragraph["context"].strip()
                    for qa in paragraph["qas"]:
                        question = qa["question"].strip()
                        id_ = qa["id"]

                        answer_starts = [answer["answer_start"] for answer in qa["answers"]]
                        answers = [answer["text"].strip() for answer in qa["answers"]]

                        yield id_, {
                            "context": context,
                            "question": question,
                            "id": id_,
                            "answers": {
                                "answer_start": answer_starts,
                                "text": answers,
                            },
                        }

1.4 DataCollator

from transformers import  DataCollatorWithPadding

dataset = load_dataset("csv", data_files="./ChnSentiCorp_htl_all.csv", split='train')
dataset = dataset.filter(lambda x: x["review"] is not None)
dataset

Dataset({
    features: ['label', 'review'],
    num_rows: 7765
})

def process_function(examples):
    # 这里我们指定max_length为128，超过此长度，进行截断
    # 注意：这里，我们不指定填充,我们利用 DataCollator进行填充
    tokenized_examples = tokenizer(examples["review"], max_length=128, truncation=True)
    tokenized_examples["labels"] = examples["label"]
    return tokenized_examples



tokenized_dataset = dataset.map(process_function, batched=True, remove_columns=dataset.column_names)

# DataCollatorWithPadding：把一批样本补齐到【这批样本最长句子的长度】而非整个数据集的最大长度
# 这样能够加快补齐速度
collator = DataCollatorWithPadding(tokenizer=tokenizer)


from torch.utils.data import DataLoader
dl = DataLoader(tokenized_dataset, batch_size=4, collate_fn=collator, shuffle=True)

num = 0
for batch in dl:
    # 可以看到每个批次的样本长度不是一样的
    print(batch["input_ids"].size())
    num += 1
    if num > 10:
        break

torch.Size([4, 117])
torch.Size([4, 128])
torch.Size([4, 128])
torch.Size([4, 128])
torch.Size([4, 128])
torch.Size([4, 100])
torch.Size([4, 82])
torch.Size([4, 128])
torch.Size([4, 128])
torch.Size([4, 128])
torch.Size([4, 91])

1.5 利用Datasets优化情感分类

1.5.1 加载数据集

from transformers import AutoTokenizer, AutoModelForSequenceClassification
from datasets import load_dataset

# 1、加载数据集
dataset = load_dataset("csv", data_files="./ChnSentiCorp_htl_all.csv", split="train")
dataset = dataset.filter(lambda x: x["review"] is not None)

# 2、划分数据集
datasets = dataset.train_test_split(test_size=0.1)
datasets

DatasetDict({
    train: Dataset({
        features: ['label', 'review'],
        num_rows: 6988
    })
    test: Dataset({
        features: ['label', 'review'],
        num_rows: 777
    })
})

1.5.2 创建Dataloader

import torch
from torch.utils.data import DataLoader
from transformers import DataCollatorWithPadding

# 1、加载离线模型
model_path = '/root/autodl-fs/models/rbt3'
tokenizer = AutoTokenizer.from_pretrained(model_path)

def process_function(examples):
    tokenized_examples = tokenizer(examples["review"], max_length=128, truncation=True)
    tokenized_examples["labels"] = examples["label"]
    return tokenized_examples

# 2、词元化
tokenized_datasets = datasets.map(process_function, batched=True, remove_columns=datasets["train"].column_names)



# 3、创建Dataloader
trainset, validset = tokenized_datasets["train"], tokenized_datasets["test"]
trainloader = DataLoader(trainset, batch_size=32, shuffle=True, collate_fn=DataCollatorWithPadding(tokenizer))
validloader = DataLoader(validset, batch_size=64, shuffle=False, collate_fn=DataCollatorWithPadding(tokenizer))

1.5.3 创建模型及优化器

from torch.optim import Adam

# 1、创建模型
model = AutoModelForSequenceClassification.from_pretrained(model_path)

if torch.cuda.is_available():
    model = model.cuda()
# 2、创建优化器
optimizer = Adam(model.parameters(), lr=2e-5)    

1.5.4 训练与验证

def evaluate():
    model.eval()
    acc_num = 0
    with torch.inference_mode():
        for batch in validloader:
            if torch.cuda.is_available():
                batch = {k: v.cuda() for k, v in batch.items()}
            output = model(**batch)
            pred = torch.argmax(output.logits, dim=-1)
            acc_num += (pred.long() == batch["labels"].long()).float().sum()
    return acc_num / len(validset)

def train(epoch=3, log_step=100):
    global_step = 0
    for ep in range(epoch):
        model.train()
        for batch in trainloader:
            if torch.cuda.is_available():
                batch = {k: v.cuda() for k, v in batch.items()}
            optimizer.zero_grad()
            output = model(**batch)
            output.loss.backward()
            optimizer.step()
            if global_step % log_step == 0:
                print(f"ep: {ep}, global_step: {global_step}, loss: {output.loss.item()}")
            global_step += 1
        acc = evaluate()
        print(f"ep: {ep}, acc: {acc}")

train()

ep: 0, global_step: 0, loss: 0.7134996056556702
ep: 0, global_step: 100, loss: 0.2385680228471756
ep: 0, global_step: 200, loss: 0.3343896269798279
ep: 0, acc: 0.8764479160308838
ep: 1, global_step: 300, loss: 0.15788553655147552
ep: 1, global_step: 400, loss: 0.339910626411438
ep: 1, acc: 0.8867439031600952
ep: 2, global_step: 500, loss: 0.24706700444221497
ep: 2, global_step: 600, loss: 0.24163328111171722
ep: 2, acc: 0.8828828930854797

1.5.5 模型预测

sen = "我觉得这家酒店不错，饭很好吃！"
id2_label = {0: "差评！", 1: "好评！"}
model.eval()
with torch.inference_mode():
    inputs = tokenizer(sen, return_tensors="pt")
    inputs = {k: v.cuda() for k, v in inputs.items()}
    logits = model(**inputs).logits
    pred = torch.argmax(logits, dim=-1)
    print(f"输入：{sen}\n模型预测结果:{id2_label.get(pred.item())}")

输入：我觉得这家酒店不错，饭很好吃！
模型预测结果:好评！

from transformers import pipeline

model.config.id2label = id2_label
pipe = pipeline("text-classification", model=model, tokenizer=tokenizer, device=0)
pipe(sen)

[{'label': '好评！', 'score': 0.9937521815299988}]

2、基础组件Evaluate

• 在训练模型的时候，我们往往会将数据划分为训练集和验证集，会根据模型在验证集的性能来评价模型的好坏。

• 任务不同，对应的评价指标也不一样，常见的指标包括Accuracy、F1、Rouge等。

• Evaluate包的出现，一定程度上简化了上述工作，我们可以通过像加载数据集一样加载对应的评估函数，进行模型评估。

2.1 评估函数常用操作

2.1.1 查看评估函数

import evaluate

# 查看支持的评估函数
evaluate.list_evaluation_modules(include_community=False, with_details=True)

# 在线加载评估函数
accuracy = evaluate.load("accuracy")

# 查看函数说明
print(accuracy.description)

print(accuracy.inputs_description)

2.1.2 评估指标计算

# 全局计算
accuracy = evaluate.load("accuracy")
results = accuracy.compute(references=[0, 1, 2, 0, 1, 2], predictions=[0, 1, 1, 2, 1, 0])
results # {'accuracy': 0.5}

# 迭代计算
accuracy = evaluate.load("accuracy")
for refs, preds in zip([[0,1],[0,1]], [[1,0],[0,1]]):
    accuracy.add_batch(references=refs, predictions=preds)
accuracy.compute() # {'accuracy': 0.5}

# 多个指标计算
clf_metrics = evaluate.combine(["accuracy", "f1", "recall", "precision"])
clf_metrics.compute(predictions=[0, 1, 0], references=[0, 1, 1])

{'accuracy': 0.6666666666666666,
 'f1': 0.6666666666666666,
 'recall': 0.5,
 'precision': 1.0
 }

2.2 利用Evaluate优化情感分类

加载数据集、创建Dataloader、创建模型及优化器和1.5一致。

2.2.1 训练与验证

import evaluate

# clf_metrics = evaluate.combine(["accuracy", "f1"])

"""
如果直接使用指标名称“accuracy”等，
由于网络问题，无法顺利下载
但是这个错误不会显示，因此导入进去会卡死


解决方法：


从官网下载metrics文件夹，放置在本地的路径，进行离线加载

"""
accuracy_path = '/root/autodl-tmp/transformers-code/metrics/accuracy'
f1_path = '/root/autodl-tmp/transformers-code/metrics/f1'
clf_metrics = evaluate.combine([accuracy_path, f1_path])

# 示例
clf_metrics.compute(references=[0,1,0,1], predictions=[1,0,0,1])

# {'accuracy': 0.5, 'f1': 0.5}

def evaluate():
    model.eval()
    with torch.inference_mode():
        for batch in validloader:
            if torch.cuda.is_available():
                batch = {k: v.cuda() for k, v in batch.items()}
            output = model(**batch)
            pred = torch.argmax(output.logits, dim=-1)
            clf_metrics.add_batch(predictions=pred.long(), references=batch["labels"].long())
    return clf_metrics.compute()

def train(epoch=3, log_step=100):
    global_step = 0
    for ep in range(epoch):
        model.train()
        for batch in trainloader:
            if torch.cuda.is_available():
                batch = {k: v.cuda() for k, v in batch.items()}
            optimizer.zero_grad()
            output = model(**batch)
            output.loss.backward()
            optimizer.step()
            if global_step % log_step == 0:
                print(f"ep: {ep}, global_step: {global_step}, loss: {output.loss.item()}")
            global_step += 1
        clf = evaluate()
        print(f"ep: {ep}, {clf}")

# 进行训练
train()

ep: 0, global_step: 0, loss: 0.6701369285583496
ep: 0, global_step: 100, loss: 0.3432188630104065
ep: 0, global_step: 200, loss: 0.29397645592689514
ep: 0, {'accuracy': 0.9086229086229086, 'f1': 0.9356300997280146}
ep: 1, global_step: 300, loss: 0.2728956639766693
ep: 1, global_step: 400, loss: 0.23207390308380127
ep: 1, {'accuracy': 0.9137709137709138, 'f1': 0.9391462306993643}
ep: 2, global_step: 500, loss: 0.10437105596065521
ep: 2, global_step: 600, loss: 0.18092380464076996
ep: 2, {'accuracy': 0.8996138996138996, 'f1': 0.9275092936802974}

2.2.2 模型预测

sen = "我觉得这家酒店不错，饭很好吃！"
id2_label = {0: "差评！", 1: "好评！"}
model.eval()
with torch.inference_mode():
    inputs = tokenizer(sen, return_tensors="pt")
    inputs = {k: v.cuda() for k, v in inputs.items()}
    logits = model(**inputs).logits
    pred = torch.argmax(logits, dim=-1)
    print(f"输入：{sen}\n模型预测结果:{id2_label.get(pred.item())}")

输入：我觉得这家酒店不错，饭很好吃！
模型预测结果:好评！

from transformers import pipeline

model.config.id2label = id2_label
pipe = pipeline("text-classification", model=model, tokenizer=tokenizer, device=0)

pipe('这家酒店床上特别脏、隔音特别差、还特别贵')

[{'label': '差评！', 'score': 0.9784351587295532}]

3、基础组件Trainer

• Trainer模块是基础组件的最后一个模块，它封装了一套完整的在数据集上训练、评估与预测的流程。借助Trainer模块，可以快速启动训练。

• Trainer模块主要包含两部分的内容：TrainingArguments与Trainer，前者用于训练参数的设置，后者用于创建真正的训练器，进行训练、评估预测等实际操作。

• 此外，针对Seq2Seq训练任务，提供了专门的Seq2SeqTrainingArguments与Seq2SeqTrainer，整体与TrainingArguments和Trainer类似，但是提供了专门用于生成的部分参数。

3.1 TrainingArguments

TrainingArguments中可以配置整个训练过程中使用的参数，默认版本是包含90个参数，涉及模型存储、模型优化、训练日志、GPU使用、模型精度、分布式训练等多方面的配置内容，等后面用到再介绍。

Seq2SeqTrainingArguments中除了上述的内容还包括生成部分的参数设置，如是否要进行生成、最大长度等共94个参数。

3.2 Trainer

Trainer中配置具体的训练用到的内容，包括模型、训练参数、训练集、验证集、分词器、评估函数等内容。

• 当指定完上述对应参数，便可以通过调用train方法进行模型训练；

• 训练完成后可以通过调用evaluate方法对模型进行评估；

• 得到满意的模型后，最后调用predict方法对数据集进行预测。

from transformers import TrainingArguments, Trainer

# 创建TrainingArguments
training_args = TrainingArguments(...)
# 创建Trainer
trainer = Trainer(..., args=training_args, ...)


# 模型训练
trainer.train()
# 模型评估
trainer.evaluate()
# 模型预测
trainer.predict()

需要特别注意的是，使用Trainer进行模型训练对模型的输入输出是有限制的，要求模型返回ModelOutput的元组或子类，同时如果提供了标签，模型要能返回loss结果，并且loss要作为ModelOutput元组的第一个值。

3.3 利用Trainer优化情感分类

3.1 数据预处理

from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments
from datasets import load_dataset

import warnings
warnings.filterwarnings("ignore")

# 1、加载数据集
dataset = load_dataset("csv", data_files="./ChnSentiCorp_htl_all.csv", split="train")
dataset = dataset.filter(lambda x: x["review"] is not None)

# 2、划分数据集
datasets = dataset.train_test_split(test_size=0.1)


# 3、词元化
# 加载离线模型
model_path = '/root/autodl-fs/models/rbt3'
tokenizer = AutoTokenizer.from_pretrained(model_path)

def process_function(examples):
    tokenized_examples = tokenizer(examples["review"], max_length=128, truncation=True)
    tokenized_examples["labels"] = examples["label"]
    return tokenized_examples

tokenized_datasets = datasets.map(process_function, batched=True, remove_columns=datasets["train"].column_names)

tokenized_datasets

DatasetDict({
    train: Dataset({
        features: ['input_ids', 'token_type_ids', 'attention_mask', 'labels'],
        num_rows: 6988
    })
    test: Dataset({
        features: ['input_ids', 'token_type_ids', 'attention_mask', 'labels'],
        num_rows: 777
    })
})

3.2 创建模型

model = AutoModelForSequenceClassification.from_pretrained(model_path)

3.3 创建模型评估函数

import evaluate

# 下面方式可能加载失败
# acc_metric = evaluate.load("accuracy")
# f1_metirc = evaluate.load("f1")

# 这里采用离线加载
accuracy_path = '/root/autodl-tmp/transformers-code/metrics/accuracy'
f1_path = '/root/autodl-tmp/transformers-code/metrics/f1'

acc_metric = evaluate.load(accuracy_path)
f1_metirc = evaluate.load(f1_path)

def eval_metric(eval_predict):
    predictions, labels = eval_predict
    predictions = predictions.argmax(axis=-1)
    acc = acc_metric.compute(predictions=predictions, references=labels)
    f1 = f1_metirc.compute(predictions=predictions, references=labels)
    acc.update(f1)
    return acc

3.4 创建TrainingArguments及Trainer

创建TrainingArguments

train_args = TrainingArguments(output_dir="./checkpoints",      # 输出文件夹
                               per_device_train_batch_size=64,  # 训练时的batch_size
                               per_device_eval_batch_size=128,  # 验证时的batch_size
                               logging_steps=10,                # log 打印的频率
                               evaluation_strategy="epoch",     # 评估策略
                               save_strategy="epoch",           # 保存策略
                               save_total_limit=3,              # 最大保存数
                               learning_rate=2e-5,              # 学习率
                               weight_decay=0.01,               # weight_decay
                               metric_for_best_model="f1",      # 设定评估指标
                               load_best_model_at_end=True)     # 训练完成后加载最优模型

创建Trainer

from transformers import DataCollatorWithPadding
trainer = Trainer(model=model,     # 预训练模型
                  args=train_args, # 训练参数
                  train_dataset=tokenized_datasets["train"], # 训练集
                  eval_dataset=tokenized_datasets["test"],   # 验证集
                  data_collator=DataCollatorWithPadding(tokenizer=tokenizer),# DataCollator，填充到一个批次中最大长度，加快填充的速度
                  compute_metrics=eval_metric  # 指标评估的方法
)

3.5 模型训练及预测

trainer.train()

from transformers import pipeline


id2_label = {0: "差评！", 1: "好评！"}
model.config.id2label = id2_label
pipe = pipeline("text-classification", model=model, tokenizer=tokenizer, device=0)

pipe('这家酒店隔音不错、位置离地铁站近、去西湖很方便。')

[{'label': '好评！', 'score': 0.9933173656463623}]