- 一、IMDb数据集
- 二、数据预处理
- 2.1 读取数据集
- 2.2 构建数据集
- 三、封装
- References
IMDb数据集[1]是一个情感分析数据集(二分类),训练集和测试集各有 25000 25000 25000 个样本(每一个样本都是一段影评),无论是训练集还是测试集,其中的正/负类(即积极/消极)样本个数均相同,为 12500 12500 12500 个。
数据集下载:IMDb-Large Movie Review Dataset。注意,这里提供的数据集和官网的相比,仅仅删除了一些不必要的文件,其他均未改动。
结构:
$ tree aclImdb -d aclImdb ├── test │ ├── neg │ └── pos └── train ├── neg └── pos 6 directories
因文件过多,这里仅展示了目录。其中每个 neg 或 pos 目录下均有 12500 12500 12500 个文件。
二、数据预处理导入本文所需要的所有包
import os import torch from torchtext.data import get_tokenizer from torchtext.vocab import build_vocab_from_iterator from torchtext import transforms as T from torch.utils.data import TensorDataset2.1 读取数据集
首先定义一个 read_imdb() 函数,它的作用是从给定的路径中读取训练集或测试集,并返回分词后的影评及其对应的标签。
这里我们采用 torchtext 的分词器进行分词:
def read_imdb(path='./aclImdb', is_train=True): reviews, labels = [], [] tokenizer = get_tokenizer('basic_english') for label in ['pos', 'neg']: folder_name = os.path.join(path, 'train' if is_train else 'test', label) for filename in os.listdir(folder_name): with open(os.path.join(folder_name, filename), mode='r', encoding='utf-8') as f: reviews.append(tokenizer(f.read())) labels.append(1 if label == 'pos' else 0) return reviews, labels
返回的 reviews 是一个二维列表,即 reviews[0] 代表一段(分词了的)影评,reviews[0][0] 代表这段影评中的一个词元。返回的 labels 是一个一维列表,其中的每个元素非 0 0 0 即 1 1 1。 0 0 0 代表是消极影评(负样本), 1 1 1 代表是积极影评(正样本)。
2.2 构建数据集经过简单的统计可以发现,分词后,训练集中最短影评的长度为 11 11 11,最长影评的长度为 2752 2752 2752,平均长度为 271 271 271;测试集中最短影评的长度为 8 8 8,最长影评的长度为 2623 2623 2623,平均长度为 265 265 265。
因为批量加载句子需要统一长度,所以这里我们选择
512
512
512 作为标准,长度超过
512
512
512 的进行截断,长度不到
512
512
512 的使用填充词元
对于截断和填充操作,这里选择调用 torchtext.transforms 中的相关API:
def build_dataset(reviews, labels, vocab, max_len=512): text_transform = T.Sequential( T.VocabTransform(vocab=vocab), T.Truncate(max_seq_len=max_len), T.ToTensor(padding_value=vocab['']), T.PadTransform(max_length=max_len, pad_value=vocab[' ']), ) dataset = TensorDataset(text_transform(reviews), torch.tensor(labels)) return dataset
对于词典,我们依然采用 torchtext 中提供的API,之后就可以设置一个 load_imdb() 函数用来加载数据集:
def load_imdb(): reviews_train, labels_train = read_imdb(is_train=True) reviews_test, labels_test = read_imdb(is_train=False) vocab = build_vocab_from_iterator(reviews_train, min_freq=3, specials=['', ' ', ' ', ' ']) vocab.set_default_index(vocab[' ']) train_data = build_dataset(reviews_train, labels_train, vocab) test_data = build_dataset(reviews_test, labels_test, vocab) return train_data, test_data, vocab
这里之所以加上
为了让项目更加结构化,我们将上述三个函数封装进 data_preprocess.py 文件中,之后需要加载数据集的时候,只需要简单导入:
from data_preprocess import load_imdb
References
[1] https://ai.stanford.edu/~amaas/data/sentiment/