最近實在是有點忙,沒啥時間寫博客了。趁着周末水一文,把最近用 huggingface transformers 訓練文本分類模型時遇到的一個小問題說下。
背景
之前只聞 transformers 超厲害超好用,但是沒有實際用過。之前涉及到 bert 類模型都是直接手寫或是在別人的基礎上修改。但這次由於某些原因,需要快速訓練一個簡單的文本分類模型。其實這種場景應該挺多的,例如簡單的 POC 或是臨時測試某些模型。
我的需求很簡單:用我們自己的數據集,快速訓練一個文本分類模型,驗證想法。
我覺得如此簡單的一個需求,應該有模板代碼。但實際去搜的時候發現,官方文檔什么時候變得這么多這么龐大了?還多了個 Trainer
API?瞬間讓我想起了 Pytorch Lightning 那個坑人的同名 API。但可能是時間原因,找了一圈沒找到適用於自定義數據集的代碼,都是用的官方、預定義的數據集。
所以弄完后,我決定簡單寫一個文章,來說下這原本應該極其容易解決的事情。
數據
假設我們數據的格式如下:
0 第一個句子
1 第二個句子
0 第三個句子
即每一行都是 label sentence
的格式,中間空格分隔。並且我們已將數據集分成了 train.txt
和 val.txt
。
代碼
加載數據集
首先使用 datasets
加載數據集:
from datasets import load_dataset
dataset = load_dataset('text', data_files={'train': 'data/train_20w.txt', 'test': 'data/val_2w.txt'})
加載后的 dataset
是一個 DatasetDict
對象:
DatasetDict({
train: Dataset({
features: ['text'],
num_rows: 3
})
test: Dataset({
features: ['text'],
num_rows: 3
})
})
類似 tf.data
,此后我們需要對其進行 map
,對每一個句子進行 tokenize、padding、batch、shuffle:
def tokenize_function(examples):
labels = []
texts = []
for example in examples['text']:
split = example.split(' ', maxsplit=1)
labels.append(int(split[0]))
texts.append(split[1])
tokenized = tokenizer(texts, padding='max_length', truncation=True, max_length=32)
tokenized['labels'] = labels
return tokenized
tokenized_datasets = dataset.map(tokenize_function, batched=True)
train_dataset = tokenized_datasets["train"].shuffle(seed=42)
eval_dataset = tokenized_datasets["test"].shuffle(seed=42)
根據數據集格式不同,我們可以在 tokenize_function
中隨意自定義處理過程,以得到 text 和 labels。注意 batch_size
和 max_length
也是在此處指定。處理完我們便得到了可以輸入給模型的訓練集和測試集。
訓練
model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=2, cache_dir='data/pretrained')
training_args = TrainingArguments('ckpts', per_device_train_batch_size=256, num_train_epochs=5)
trainer = Trainer(
model=model,
args=training_args,
train_dataset=train_dataset,
eval_dataset=eval_dataset
)
trainer.train()
你可以根據情況修改訓練 batchsize per_device_train_batch_size
。
完整代碼
完整代碼見 GitHub。