前往小程序,Get更优阅读体验!
立即前往
首页
学习
活动
专区
工具
TVP
发布
社区首页 >专栏 >使用 Transformers 在你自己的数据集上训练文本分类模型

使用 Transformers 在你自己的数据集上训练文本分类模型

作者头像
Alan Lee
发布2021-12-07 15:50:57
2.3K0
发布2021-12-07 15:50:57
举报
文章被收录于专栏:Small Code

最近实在是有点忙,没啥时间写博客了。趁着周末水一文,把最近用 huggingface transformers 训练文本分类模型时遇到的一个小问题说下。

背景

之前只闻 transformers 超厉害超好用,但是没有实际用过。之前涉及到 bert 类模型都是直接手写或是在别人的基础上修改。但这次由于某些原因,需要快速训练一个简单的文本分类模型。其实这种场景应该挺多的,例如简单的 POC 或是临时测试某些模型。

我的需求很简单:用我们自己的数据集,快速训练一个文本分类模型,验证想法。

我觉得如此简单的一个需求,应该有模板代码。但实际去搜的时候发现,官方文档什么时候变得这么多这么庞大了?还多了个 Trainer API?瞬间让我想起了 Pytorch Lightning 那个坑人的同名 API。但可能是时间原因,找了一圈没找到适用于自定义数据集的代码,都是用的官方、预定义的数据集。

所以弄完后,我决定简单写一个文章,来说下这原本应该极其容易解决的事情。

数据

假设我们数据的格式如下:

代码语言:javascript
复制
0 第一个句子
1 第二个句子
0 第三个句子

即每一行都是 label sentence 的格式,中间空格分隔。并且我们已将数据集分成了 train.txtval.txt

代码

加载数据集

首先使用 datasets 加载数据集:

代码语言:javascript
复制
from datasets import load_dataset
dataset = load_dataset('text', data_files={'train': 'data/train_20w.txt', 'test': 'data/val_2w.txt'})

加载后的 dataset 是一个 DatasetDict 对象:

代码语言:javascript
复制
DatasetDict({
    train: Dataset({
        features: ['text'],
        num_rows: 3
    })
    test: Dataset({
        features: ['text'],
        num_rows: 3
    })
})

类似 tf.data ,此后我们需要对其进行 map ,对每一个句子进行 tokenize、padding、batch、shuffle:

代码语言:javascript
复制
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_sizemax_length 也是在此处指定。处理完我们便得到了可以输入给模型的训练集和测试集。

训练

代码语言:javascript
复制
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

END

本文参与 腾讯云自媒体同步曝光计划,分享自作者个人站点/博客。
原始发表:2021/11/13 ,如有侵权请联系 cloudcommunity@tencent.com 删除

本文分享自 作者个人站点/博客 前往查看

如有侵权,请联系 cloudcommunity@tencent.com 删除。

本文参与 腾讯云自媒体同步曝光计划  ,欢迎热爱写作的你一起参与!

评论
登录后参与评论
0 条评论
热度
最新
推荐阅读
目录
  • 背景
  • 数据
  • 代码
    • 加载数据集
      • 训练
      • 完整代码
      • END
      领券
      问题归档专栏文章快讯文章归档关键词归档开发者手册归档开发者手册 Section 归档