pip install transformers[sentencepiece]
pip install datasets
from datasets import load_dataset
raw_datasets = load_dataset("glue", "sst2")
raw_datasets
这里使用的是GLUE中SST2数据集,主要针对电影评论来做情感分类(二分类)。
整个流程是先分词,将文本转化为向量,输入到模型,得到输出后再进行后处理操作。
from transformers import AutoTokenizer
checkpoint = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
checkpoint就是预训练好的模型。每个模型都有自己专门的tokenizer,可以执行分词,padding等操作。
来看一下具体的用法。
inputs = tokenizer("This is the first sentence.")
inputs
这里的输出有三项,input_id标识句子中的每个token,token_type_id就是bert中的segment embedding,用来区分两句话的(0代表第一句,1代表第二句),attention_mask用来mask掉句子末尾padding的词,防止计算注意力时有偏差。
def tokenize_function(example):
return tokenizer(example["sentence"], truncation=True)
tokenized_datasets = raw_datasets.map(tokenize_function, batched=True,remove_columns=['idx','sentence'])
tokenized_datasets
定义了一个分词的函数,将数据集中的sentence作为参数传进去。
from transformers import DataCollatorWithPadding
data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
data_collator的作用是自动将同一个batch的句子padding成同一长度,而不是一次性padding整个数据集。
from transformers import TrainingArguments
from transformers import AutoModelForSequenceClassification
from datasets import load_metric
from transformers import Trainer
import numpy as np
training_args = TrainingArguments("test-trainer", evaluation_strategy="epoch")# 训练需要的参数,默认的
model = AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels=2)
def compute_metrics(eval_preds):
'''
计算评价指标
'''
metric = load_metric("glue", "sst2")
logits, labels = eval_preds
predictions = np.argmax(logits, axis=-1)
return metric.compute(predictions=predictions, references=labels)
trainer = Trainer(
model,
training_args,
train_dataset=tokenized_datasets["train"],
eval_dataset=tokenized_datasets["validation"],
data_collator=data_collator,
tokenizer=tokenizer,
compute_metrics=compute_metrics,
)
这段很清晰,导入Trainer,把训练需要的模型、数据、tokenizer等一次性传入即可。compute_metrics是评价指标,我们只需要把模型得到的结果输入到metric.compute,就能自动计算得到结果。由于是二分类,模型输出的最后一维是2,我们取较大的一项作为分类结果。
接下来就是训练啦。
trainer.train()