情感分析系列（一）——IMDb数据集及其预处理

一、IMDb数据集

IMDb数据集[1]是一个情感分析数据集（二分类），训练集和测试集各有 $25000$ 个样本（每一个样本都是一段影评），无论是训练集还是测试集，其中的正/负类（即积极/消极）样本个数均相同，为 $12500$ 个。

数据集下载：IMDb-Large Movie Review Dataset。注意，这里提供的数据集和官网的相比，仅仅删除了一些不必要的文件，其他均未改动。

结构：

$ tree aclImdb -d
aclImdb
├── test
│   ├── neg
│   └── pos
└── train
    ├── neg
    └── pos

6 directories

因文件过多，这里仅展示了目录。其中每个 neg 或 pos 目录下均有 $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 TensorDataset

2.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$ 即 $1$。$0$ 代表是消极影评（负样本），$1$ 代表是积极影评（正样本）。

2.2 构建数据集

经过简单的统计可以发现，分词后，训练集中最短影评的长度为 $11$，最长影评的长度为 $2752$，平均长度为 $271$；测试集中最短影评的长度为 $8$，最长影评的长度为 $2623$，平均长度为 $265$。

因为批量加载句子需要统一长度，所以这里我们选择 $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

这里之所以加上 和 这两个特殊词元是因为本系列稍后的文章会将BERT应用在该数据集上。

三、封装

为了让项目更加结构化，我们将上述三个函数封装进 data_preprocess.py 文件中，之后需要加载数据集的时候，只需要简单导入：

from data_preprocess import load_imdb

---

References

[1] https://ai.stanford.edu/~amaas/data/sentiment/