【NLP】TorchText介绍与（搜狗语料）文本语料处理案例

原创一笑清寒皮皮AI记

文 | 菊子皮 (转载请注明出处)

B站：科皮子菊

背景

自然语言处理领域包含很多任务，在日常的任务中序列标注、文本分类等任务也特别多。真正进入项目工程的时候你会发现，我们大多会花费很多时间在数据源处理（读取，切分词，清洗，标准化，特征提取等操作）上。在pytorch中众所周知的数据处理包是处理图片的torchvision，而处理文本的少有提及，快速处理文本数据的包也是有的，那就是torchtext^[1]。当然原文档内容较多，这里也只是介绍一些核心内容，后面有机会的话再详细（尝试翻译一下）介绍。除此之外还参考了TorchText用法示例及完整代码^[2]。废话不多说，来看看如何使用吧。

ps:需要自己准备好相关的环境哦(pip install torchtext)

1 torchtext概述

站在处理模型训练的角度上，我们需要的数据集是数字化的、可批次处理的。而torchtext主要包含Field，Dataset和迭代器三个部分，这三个部分与我们实际处理中所需要的内容相对应。具体如何使用，下面来看。

1.1 Field对象

Field对象指定要如何处理某个字段，也就是一条样本数据包含的一部分。例如一条文本分类样本包含文本内容及对应的文本标签，那么就可以创建两个Field对象。其中包含数据预处理的一些配置信息，例如：分词方式，是否转成大小写，起始字符，补全字符以及字典等。Field包含一些文本处理的通用参数的设置，并且还包含一个词典对象，可以把文本数据表示成数字类型（也就是将文字特征数字化的过程），进而就可以把文本转成深度学习所需的Tensor数据类型。它包含的主要参数如下：

sequential：是否把数据表示成序列，如果值为False，则不能进行分词，默认为True
use_vocab：是否使用词典对象，默认值为True，如果为False，则数据的类型必须已经是数值类型
init_token：每一条数据的其实字符，默认None
eos_token：每条数据的结尾字符，默认None
fix_length：修改每条数据的长度为该值，长度不够的使用pad_token补全，默认None
tensor_type：把数据转成的tensor类型，默认值为tensor.LongTensor
tokenize：分词函数，默认值为str.split

例如我们处理文本分类的语料,那么处理好的语料包含本文内容及其对应的标签.那么我们就可以定一个两个Field,即LALBEL，TEXT。

对于LABEL传入的是原始的未转换为数字的类别时，后面的处理中会使用vocab方法将其转化为类别对应的数字（数字化），如果传入的已经处理过的类别对应的数字的话，那就不需要进行数字化操作，这是use_vocab需要设置为False。由于标签是不需要分词的，所有sequential需要设置为False。
对应TEXT，通常需要对需要分类的文本进行分词，我们可以手写一个分词方法设置到属性中，前面已经提及默认的方法。根据不同的算法模型，还可进行pad处理等。

根据上面的描述我们可做下面的定义：

import jieba
from torchtext imort data, datsets

# 使用jieba自定义分词函数
def tokenizer(text):
    return [word for word in jieba.lcut(text) if word.strip()]

# 定义语料字段
TEXT = data.Field(sequential=True, tokenizer=tokenizer, fix_length=5)
LABEL = data.Field(sequential=False, use_vocab=True)
# 还需告诉fields处理哪些数据
tv_datafields = {'title':('title', LABEL), 'text':('text', TEXT)}

补充: tv_datafields 这个就是将定义的字段与实际的文本联系起来的，具体操作看看后面的就知道了。

1.2 Dataset

Dataset用于定义数据源信息，就是我们数据来自于哪里。其继承字PyTorch的Dataset，提供了TabularDataset等，可以指定路径，格式Field信息，这样就可以更方便加载数据了。与torchvision一样，TorchText也预先构建了常用的数据集的Dataset对象，可以直接使用。提供的有splits方法可以同时加载训练集，验证集和测试集；可以下载压缩数据并解压的方法（支持.zip,.gz,.tgz）。

TabularDataset可以很方便地读取CSV，TSV以及JSON格式的文件.

在1.1节中我们很好地定义了语料所需的字段，下面我们就可以使用Dataset相关属性去读取实际的语料，并将语料与我们定义的字段联系在一起了。事例如下：

# 读取数据
train_set, dev_set = data.TabularDataset.splits(
                     path='/data/', 
                     train='train.json',
                     validation='dev.json',
                     format='json',
                     fields=tv_datafields)
# 查看数据
for i in range(len(train_set)):
    print(vars(train_set[i]))   # vars函数返回对象object的属性和属性值得对象

上面的代码，打印的是一个Example对象，Example对象其实就对应的是一条语料数据。可以通过：

train_set[0].__dict__.keys()

查看一个Example对象设置的键，不出意外的是，返回的结果是：

dict_keys(['title', 'text'])

并且，还可以通过下面的方式查看text对应的内容：

train_set[0].text

这时我们的数据已经加载完毕了，紧接着我们则需要将数据构建响应的字典以方便进一步的数字化处理。也可以加载一些预训练的word-embedding(词向量)。例如加载预训练的词向量“glove.6B.100d”：

TEXT.build_vocab(train, vectors="glove.6B.100d")

1.3 迭代器

迭代器返回模型所需要的处理后的批次数据。对于数据量比较大的样本，模型不可能一次性加上训练，毕竟内存就那么大，那么就需要构造迭代器去处理，也就是我们常说的batch。迭代器主要分为Iterator, BucketIerator, BPTTIterator三种。

在1.1，1.2处理完之后，我们将应该开始做batching操作了，也就是对我们语料进行打乱，排序，批次对齐等处理。主要的相关参数如下：

dataset:加载的数据集
batch_size：batch大小
batch_size_fn：产生动态的batch大小的函数
sort_key：排序的key
train：是否是一个训练集
repeat：是否在不同epoch中重复迭代
shuffle：是否对数据进行排序
sort：是否对数据进行排序
sort_within_batch:batch内部是否排序
device：数据存放的设备

如果数据运行在cpu上，则device=-1，如果运行在GPU上，则device=0。当然也可以在后需要的操作中再做处理。

注：TorchText使用了动态填充，也就是说，batch内的所有句子填充到batch中句子最长的长度。

事例代码：

train_iter, val_iter = data.Iterator.splits((train_set, dev_set),
                       sort_key=lambda x: len(x.Text),
                       batch_sizes = (256, 256),
                       device=-1)
for x in train_iter:
    print(x)   # 打印查看

2 具体使用

实战为王，下面我们针对一个语料，按照文本分类的任务去处理，具体看看如何操作。

2.1 语料选择

本案例选取的是搜狗实验室提供的来自若干新闻站点2012年6月—7月期间国内，国际，体育，社会，娱乐等18个频道的新闻数据，提供URL和正文信息。具体信息可参考：https://www.sogou.com/labs/resource/ca.php^[3]

每篇的数据格式如下：

<doc>
<url>页面URL</url>
<docno>页面ID</docno>
<contenttitle>页面标题</contenttitle>
<content>页面内容</content>
</doc>

注意：content字段去除了HTML标签，保存的是新闻正文文本案例选取的是“精简版(一个月数据, 437MB)：tar.gz格式”，原文的文本编码是GBK。

2.2 数据预处理

从数据格式可以看出，这些数据与我们输入模型的数据还存在很大的差别。先看一下语料样例，一个文件中包含多个如下的数据：

<doc>
<url>http://2008.163.com/comment/0074/2008_bbs/04/4CFRUM0400742437.html</url>
<docno>013e65ba3b398a67-b4f5d9a362314a50</docno>
<contenttitle></contenttitle>
<content></content>
</doc>
<doc>
<url>http://news.qq.com/a/20080620/000677_6.htm</url>
<docno>005323975d05d5a5-309908a255a73ab0</docno>
<contenttitle></contenttitle>
<content>组图：震前汶川风光震前汶川风光　ＱＱ群４９１４６６７．作者肚螂皮</content>
</doc>

可以从url中获取当前doc标签中的content内容对应的标签，如上，标签为:“news”，文本内容就是“组图：震前汶川风光震前汶川风光　ＱＱ群４９１４６６７．作者肚螂皮”，那么接下来使用正则表达式提取数据就简单多了。除此之外，为了更好地解码，我选取了gb18030（包含的文字多）对数据进行解码，本文主要介绍使用TorchText，这里我只提取.sina.com 链接中的数据。下面就可查看如何获取文章极其对应的标签了及相关类别的统计，程序如下：

#!/usr/bin/python
# -*- coding: UTF-8 -*-
"""
@author: juzipi
@file: sougou_data_deal.py
@time:2021/03/21
@description: 处理搜狗语料数据
"""
import re
import os
from collections import Counter


def create_data_set(file_dir: str):
    """
    提取语料中的文本及其标签
    :param file_dir: 处理的文件夹
    :return: contents_list, classes_list
    """
    # 定义提取数据的正则表达式
    # 获取所有文本对应url
    pattern_url = re.compile(r'<url>(.*?)</url>', re.S)
    # 提取文本内容
    pattern_content = re.compile(r'<content>(.*?)</content>', re.S)
    contents_list = []
    classes_list = []
    # 查看新闻的种类共有多少类以及各个类别有多少篇新闻
    for file in os.listdir(file_dir):
        file_path = os.path.join(file_dir, file)
        with open(file_path, 'r', encoding='gb18030') as reader:
            text = reader.read()
            # 正则表达式匹配出url和content
            urls = pattern_url.findall(text)
            contents = pattern_content.findall(text)
            for i in range(len(urls)):
                # 提取文本类别
                pattern_class = re.compile(r'http://(.*?).sina.com', re.S)
                class_type_find = pattern_class.findall(urls[i])
                class_type = class_type_find[0] if class_type_find else ''
                content = contents[i]
                if class_type and content:
                    classes_list.append(class_type)
                    contents_list.append(content)
    return contents_list, classes_list


if __name__ == '__main__':
    dir_path = 'Data/SogouCA.reduced'
    content_lst, class_lst = create_data_set(dir_path)
    print(Counter(class_lst))

程序运行结果如下：

Counter({'finance': 71942, 'sports': 44697, 'ent': 28032, 'eladies': 13977, 'edu': 9895, '2008': 9834, 'auto': 9588, 'mil.news': 3320, 'house': 1851, 'tech': 1732, 'cul.book': 440, 'news': 35, 'tour': 4})

我们看到有些类别不是我们常见的，如：2008，这个标签应该是2008年，除此之外，类别也极其不均衡，我们可以在代码中过滤一些少见的类别以及数据比较少的，下面就开始处理吧，这里处理的规则是选取类别数大于1000的，前1000条数据，并去除2008这个类别，然后将这些数据以6:2:2分成三批并存到指定目录，为了调试，我存储了两种格式：txt和pkl：

#!/usr/bin/python
# -*- coding: UTF-8 -*-
"""
@author: juzipi
@file: sougou_data_deal.py
@time:2021/03/21
@description: 处理搜狗语料数据
"""
import re
import os
from collections import Counter
from sklearn.model_selection import train_test_split


def create_data_set(file_dir: str):
    """
    提取语料中的文本及其标签
    :param file_dir: 处理的文件夹
    :return: contents_list, classes_list
    """
    # 定义提取数据的正则表达式
    # 获取所有文本对应url
    pattern_url = re.compile(r'<url>(.*?)</url>', re.S)
    # 提取文本内容
    pattern_content = re.compile(r'<content>(.*?)</content>', re.S)
    contents_list = []
    classes_list = []
    # 查看新闻的种类共有多少类以及各个类别有多少篇新闻
    for file in os.listdir(file_dir):
        file_path = os.path.join(file_dir, file)
        with open(file_path, 'r', encoding='gb18030') as reader:
            text = reader.read()
            # 正则表达式匹配出url和content
            urls = pattern_url.findall(text)
            contents = pattern_content.findall(text)
            for i in range(len(urls)):
                # 提取文本类别
                pattern_class = re.compile(r'http://(.*?).sina.com', re.S)
                class_type_find = pattern_class.findall(urls[i])
                class_type = class_type_find[0] if class_type_find else ''
                content = contents[i]
                if class_type and content:
                    classes_list.append(class_type)
                    contents_list.append(content)
    return contents_list, classes_list


def simple_extract(contents: list, kinds: list, res_root: str = 'Data/SougouCA.reduced_Selected'):
    """
    根据一定规则提取数据，并存放到指定文件中。其中包含三个数据集train.txt, dev.txt,test.txt
    :param res_root:
    :param contents: 文本内容
    :param kinds: 文本对应标签
    :return: 并拆分成三个数据集
    """
    # 构建格式化数据
    X = []
    y = []
    counts = Counter(kinds)
    statistics = {}
    for kind, content in zip(kinds, contents):
        if kind == '2008':
            continue
        if counts.get(kind) >= 1000:
            if kind in statistics:
                if statistics[kind] > 1000:
                    continue
                else:
                    statistics[kind] += 1
            else:
                statistics[kind] = 1
            X.append(content)
            y.append(kind)
    # 开始将数据拆分为多个集
    X_train, X_dev_test, y_train, y_dev_test = train_test_split(X, y, test_size=0.4)
    X_dev, X_test, y_dev, y_test = train_test_split(X_dev_test, y_dev_test, test_size=0.5)
    # 保存数据为：类别\t内容
    with open(os.path.join(res_root, 'train.txt'), 'w', encoding='utf8') as writer:
        for content, kind in zip(X_train, y_train):
            writer.write("{}\t{}\n".format(kind, content))
    with open(os.path.join(res_root, 'dev.txt'), 'w', encoding='utf8') as writer:
        for content, kind in zip(X_dev, y_dev):
            writer.write("{}\t{}\n".format(kind, content))
    with open(os.path.join(res_root, 'test.txt'), 'w', encoding='utf8') as writer:
        for content, kind in zip(X_test, y_test):
            writer.write("{}\t{}\n".format(kind, content))
    # 保存数据为pkl格式
    with open(os.path.join(res_root, 'train.pkl'), 'wb') as writer:
        pickle.dump(list(zip(X_dev, y_dev)), writer)

    with open(os.path.join(res_root, 'dev.pkl'), 'wb') as writer:
        pickle.dump(list(zip(X_train, y_train)), writer)

    with open(os.path.join(res_root, 'test.pkl'), 'wb') as writer:
        pickle.dump(list(zip(X_test, y_test)), writer)


if __name__ == '__main__':
    dir_path = 'Data/SogouCA.reduced'
    content_lst, class_lst = create_data_set(dir_path)
    simple_extract(content_lst, class_lst)

处理后的结果如下图所示：

2.3 使用TorchText进行数据预处理

好了，文章的主角到咯。这里处理主要分为4步，具体如下。

2.3.1 构建数据字段Field

这个则需要我们使用torchtext.data中的Field定义相关数据属性。首先我们定义一个分词方法如下，当然在实际场景中可以增加停用词，领域词等。

def tokenizer(text: str):
    """
    自定义分词方法
    :param text: 待分词的文本
    :return: 分词后的结果 list
    """
    return [word for word in jieba.lcut(text) if word.strip()]

构建这个文本分类相关数据字段，

    TEXT = Field(sequential=True, tokenize=tokenizer, batch_first=True, fix_length=200)
    LABEL = Field(sequential=False)

按照之前的定义，这个应该很简单了。上面两个Field对象，一个是存储文本对应的标签，一个是存储文本内容的，并且实际处理文本时只选取前200个词（不足时补齐），实际场景中则根据情况修改选取词语的个数。Sequential选择True是因为，我们的文本需要表示成序列的。而对于标签，我们则不需要进行分词，故此设置False.以及，我们自己为需要处理的序列书写了自定义的分词方法，不要忘了，默认使用的分词方法是str.split()哦。在实际的模型训练中，我们经常处理的数据是一批批的数据，所以，我们把batch_first的值设置为True。由于，我们处理的文本数据不是已经进行类别映射（数字化），所以使用use_vocab的默认设置。

2.3.2 封装数据成为datasets

torchtext.datasets实例能够对数据进行封装。TorchText预置了Dataset类至少需要传入examples和fields这两个参数。examples是TorchText中的example对象构造的列表，fields则是我们定义的字段，其应该与之前生成数据的列保持对应。下面，我们自定义一个Dataset。

class MyDataset(data.Dataset):
    def __init__(self, data_tuple, text_field, label_field, test=False):
        fields = [('text', text_field), ('label', label_field)]
        examples = []
        if test:
            # 如果为测试集时，则不加载label
            for content, label in tqdm(data_tuple):
                examples.append(data.Example.fromlist([content, None], fields))
        else:
            for content, label in tqdm(data_tuple):
                examples.append(data.Example.fromlist([content, label], fields))
        # 使用父类方法初始化
        super().__init__(examples, fields)

然后就可以开始生成三个数据集了。

def gen_torch_text():
    TEXT = Field(sequential=True, tokenize=tokenizer, batch_first=True, fix_length=200)
    LABEL = Field(sequential=False)
    data_root_path = "Data/SougouCA.reduced_Selected"
    with open(os.path.join(data_root_path, 'train.pkl'), 'rb') as reader:
        train_set = pickle.load(reader)
    with open(os.path.join(data_root_path, 'dev.pkl'), 'rb') as reader:
        dev_set = pickle.load(reader)
    with open(os.path.join(data_root_path, 'test.pkl'), 'rb') as reader:
        test_set = pickle.load(reader)
    train_set = MyDataset(train_set, text_field=TEXT, label_field=LABEL, test=False)
    dev_set = MyDataset(dev_set, text_field=TEXT, label_field=LABEL, test=False)
    test_set = MyDataset(test_set, text_field=TEXT, label_field=LABEL, test=True)

2.3.3 构建词表

使用TorchText简单快捷，在上面的操作中基本上完成了对数据的加载了。然后使用Field的build_vocab方法传入相关参数即可构建词表，我们可以使用预训练的词向量构建，也可只是构建一个词语数字化的表，然后在什么模型中使用Embeding层处理。对于中文预训练词向量可以从：https://github.com/Embedding/Chinese-Word-Vectors^[4]中下载，使用预训练词向量方式如下：

 # 使用预训练的词向量，维度为300维
    vectors = Vectors(name="Data/sgns.sougou.word")
    TEXT.build_vocab(train_set, vectors=vectors)

其中，vectors参数可以直接传入一个字符串类型的值，指定所需的预训练词向量。构建数字映射类型字典方式如下：

TEXT.build_vocab(train_set)
LABEL.build_vocab(train_set)

当然也可以从多个数据集中构建字典，例如：

TEXT.build_vocab(train_set, dev_set)

构建词表完毕之后，我可热查看词频，以及各个词向量以及标签索引等。例如：

print(TEXT.vocab.freqs)  # 查看词频
print(LABEL.vocab.stoi)  # 查看标签索引
print(TEXT.vocab.vectors) # 查询每个词的向量

具体结果，读者可以自己测试。

2.3.4 构建迭代器

模型的训练通常是批次训练的，而实际数据量是比较大的，比较好的解决方法是提供一个迭代器，模型需要时再获取对应的批次数据。Iterator是TorchText到模型的输出，提供了很多有用的功能，例如：打乱，排序等。通常：

验证集和训练集使用BucketIterator.splits()方法构建迭代器，目的是进行shuffle，padding等，这样可以通过将长度相似的句子shuflle一起，提高训练效率。
测试集用Iterator，如：test_iter = Iterator(test_set, batch_size=128,devic-1,train=False)
sort是对全体数据按照升序进行排序，而sort_within_batch仅对一个batch内部的数据进行排序
sort_within_batch参数设置为True时，sort_key对每个小批次内的数据进行降序排序。当我们想使用pack_padded_sequence将padded序列转换成PackedSequence对象时，那么这项操作就是必须的（RNN中会使用到）。
repeat是否在不同epoch中重复迭代
device可以是torch.device

好了，现在我们就开始构建我们的迭代器吧。

    train_iter, dev_iter = BucketIterator.splits((train_set, dev_set),
                                                 batch_size=128,
                                                 device=torch.device('cuda' if torch.cuda.is_available() else 'cpu'),
                                                 sort_key=lambda x: len(x.text),
                                                 sort_within_batch=False,
                                                 shuffle=True,
                                                 repeat=False)
    # 检验
    for train_data in train_iter:
        print(train_data.text)
        print(train_data.label)
        exit()

处理后的结果如下，即打印一个batch_size的数据：

tensor([[  181,   250,    36,  ...,     1,     1,     1],
        [  278,  9612,   284,  ...,   124,     2,    79],
        [ 7945,    95, 13780,  ...,     1,     1,     1],
        ...,
        [ 1521,  6603,   412,  ..., 39968,     3, 16072],
        [  278,  6125,   234,  ...,     8,  2062,     8],
        [38512, 12811, 45726,  ...,  6453, 13975, 14064]], device='cuda:0')
tensor([8, 7, 4, 2, 3, 4, 6, 3, 4, 4, 3, 3, 7, 1, 1, 9, 3, 1, 6, 6, 7, 6, 1, 4,
        7, 6, 8, 5, 6, 2, 2, 9, 8, 4, 7, 1, 4, 7, 5, 5, 6, 4, 3, 5, 3, 9, 6, 5,
        7, 4, 7, 8, 3, 8, 1, 9, 2, 6, 1, 2, 4, 7, 8, 8, 9, 4, 9, 4, 7, 1, 7, 8,
        9, 5, 2, 7, 6, 2, 9, 3, 1, 4, 7, 9, 1, 2, 3, 5, 8, 3, 7, 8, 8, 2, 1, 2,
        1, 5, 8, 9, 7, 1, 1, 5, 6, 8, 8, 5, 5, 1, 3, 8, 2, 7, 6, 7, 6, 9, 3, 9,
        4, 1, 1, 8, 6, 6, 6, 3], device='cuda:0')

3 总结

总的来说，这个工具是孰能生巧，用多了，你会爱死他的。用了的都说好。

附录

全部代码如下：

#!/usr/bin/python
# -*- coding: UTF-8 -*-
"""
@author: juzipi
@file: sougou_data_deal.py
@time:2021/03/21
@description: 处理搜狗语料数据
"""
import re
import os
from collections import Counter

import torch
from sklearn.model_selection import train_test_split
import jieba
from torchtext.data import Field
from torchtext import data
from tqdm import tqdm
import pickle
from torchtext.vocab import Vectors
from torchtext.data import BucketIterator


class MyDataset(data.Dataset):
    def __init__(self, data_tuple, text_field, label_field, test=False):
        fields = [('text', text_field), ('label', label_field)]
        examples = []
        if test:
            # 如果为测试集时，则不加载label
            for content, label in tqdm(data_tuple):
                examples.append(data.Example.fromlist([content, None], fields))
        else:
            for content, label in tqdm(data_tuple):
                examples.append(data.Example.fromlist([content, label], fields))
        # 使用父类方法初始化
        super().__init__(examples, fields)


def create_data_set(file_dir: str):
    """
    提取语料中的文本及其标签
    :param file_dir: 处理的文件夹
    :return: contents_list, classes_list
    """
    # 定义提取数据的正则表达式
    # 获取所有文本对应url
    pattern_url = re.compile(r'<url>(.*?)</url>', re.S)
    # 提取文本内容
    pattern_content = re.compile(r'<content>(.*?)</content>', re.S)
    contents_list = []
    classes_list = []
    # 查看新闻的种类共有多少类以及各个类别有多少篇新闻
    for file in os.listdir(file_dir):
        file_path = os.path.join(file_dir, file)
        with open(file_path, 'r', encoding='gb18030') as reader:
            text = reader.read()
            # 正则表达式匹配出url和content
            urls = pattern_url.findall(text)
            contents = pattern_content.findall(text)
            for i in range(len(urls)):
                # 提取文本类别
                pattern_class = re.compile(r'http://(.*?).sina.com', re.S)
                class_type_find = pattern_class.findall(urls[i])
                class_type = class_type_find[0] if class_type_find else ''
                content = contents[i]
                if class_type and content:
                    classes_list.append(class_type)
                    contents_list.append(content)
    return contents_list, classes_list


def simple_extract(contents: list, kinds: list, res_root: str = 'Data/SougouCA.reduced_Selected'):
    """
    根据一定规则提取数据，并存放到指定文件中。其中包含三个数据集train.txt, dev.txt,test.txt
    :param res_root:
    :param contents: 文本内容
    :param kinds: 文本对应标签
    :return: 并拆分成三个数据集
    """
    # 构建格式化数据
    X = []
    y = []
    counts = Counter(kinds)
    statistics = {}
    for kind, content in zip(kinds, contents):
        if kind == '2008':
            continue
        if counts.get(kind) >= 1000:
            if kind in statistics:
                if statistics[kind] > 1000:
                    continue
                else:
                    statistics[kind] += 1
            else:
                statistics[kind] = 1
            X.append(content)
            y.append(kind)
    # 开始将数据拆分为多个集
    X_train, X_dev_test, y_train, y_dev_test = train_test_split(X, y, test_size=0.4)
    X_dev, X_test, y_dev, y_test = train_test_split(X_dev_test, y_dev_test, test_size=0.5)
    # 保存数据为：类别\t内容
    with open(os.path.join(res_root, 'train.txt'), 'w', encoding='utf8') as writer:
        for content, kind in zip(X_train, y_train):
            writer.write("{}\t{}\n".format(kind, content))
    with open(os.path.join(res_root, 'dev.txt'), 'w', encoding='utf8') as writer:
        for content, kind in zip(X_dev, y_dev):
            writer.write("{}\t{}\n".format(kind, content))
    with open(os.path.join(res_root, 'test.txt'), 'w', encoding='utf8') as writer:
        for content, kind in zip(X_test, y_test):
            writer.write("{}\t{}\n".format(kind, content))
    # 保存数据为pkl格式
    with open(os.path.join(res_root, 'train.pkl'), 'wb') as writer:
        pickle.dump(list(zip(X_dev, y_dev)), writer)

    with open(os.path.join(res_root, 'dev.pkl'), 'wb') as writer:
        pickle.dump(list(zip(X_train, y_train)), writer)

    with open(os.path.join(res_root, 'test.pkl'), 'wb') as writer:
        pickle.dump(list(zip(X_test, y_test)), writer)


def tokenizer(text: str):
    """
    自定义分词方法
    :param text: 待分词的文本
    :return: 分词后的结果 list
    """
    return [word for word in jieba.lcut(text) if word.strip()]


def gen_torch_text():
    TEXT = Field(sequential=True, tokenize=tokenizer, batch_first=True, fix_length=200)
    LABEL = Field(sequential=False)
    data_root_path = "Data/SougouCA.reduced_Selected"
    with open(os.path.join(data_root_path, 'train.pkl'), 'rb') as reader:
        train_set = pickle.load(reader)
    with open(os.path.join(data_root_path, 'dev.pkl'), 'rb') as reader:
        dev_set = pickle.load(reader)
    with open(os.path.join(data_root_path, 'test.pkl'), 'rb') as reader:
        test_set = pickle.load(reader)
    train_set = MyDataset(train_set, text_field=TEXT, label_field=LABEL, test=False)
    dev_set = MyDataset(dev_set, text_field=TEXT, label_field=LABEL, test=False)
    test_set = MyDataset(test_set, text_field=TEXT, label_field=LABEL, test=True)
    # 使用预训练词向量
    # 使用预训练的词向量，维度为300维
    # vectors = Vectors(name="Data/sgns.sougou.word")
    # TEXT.build_vocab(train_set, vectors=vectors)
    TEXT.build_vocab(train_set)
    LABEL.build_vocab(train_set)
    # print(TEXT.vocab.freqs)  # 查看词频
    # print(LABEL.vocab.stoi)  # 查看标签索引
    train_iter, dev_iter = BucketIterator.splits((train_set, dev_set),
                                                 batch_size=128,
                                                 device=torch.device('cuda' if torch.cuda.is_available() else 'cpu'),
                                                 sort_key=lambda x: len(x.text),
                                                 sort_within_batch=False,
                                                 shuffle=True,
                                                 repeat=False)
    # 检验
    for train_data in train_iter:
        print(train_data.text)
        print(train_data.label)
        exit()



if __name__ == '__main__':
    # dir_path = 'Data/SogouCA.reduced'
    # content_lst, class_lst = create_data_set(dir_path)
    # simple_extract(content_lst, class_lst)
    gen_torch_text()