本章将介绍该系列文章中使用的数据集，并且编写预处理代码，处理成咱们需要的格式。
一、数据集介绍
咱们使用的数据集名称是chinese-poetry，是一个在github上开源的中文诗词数据集，根据仓库中readme.md中的介绍，该数据集是最全的中华古典文集数据库，包含 5.5 万首唐诗、26 万首宋诗、2.1 万首宋词和其他古典文集。诗人包括唐宋两朝近 1.4 万古诗人，和两宋时期 1.5 千古词人。
数据集的下载地址：https://github.com/chinese-poetry/chinese-poetry?tab=readme-ov-file，大家可以点击Code按钮，选择Download ZIP将该数据集下载到本地，如下图：
当然，作者收集数据也不易，大家可以顺手点一下star鼓励一下作者，如图：
如果你按照上面的步骤，把数据集下载到你本地了，解压后你可以看到如下图所示的目录结构
作者按照不同诗词类型进行了分类，并且在每个分类下提供了1个到多个的json文件，json文件里按照结构化数据组织了每一个诗词的信息，如下图
二、数据集预处理
上面咱们详细介绍了chinese-poetry数据集的下载方式和作者组织的结构，下面我们将提取每个诗词的标题和内容作为我们需要的部分，并聚合到一个文件中，以方便我们后续训练模型使用。
首先，我们需要把作者提供的诗词类目整理到一个数组中，方便我们后续进行目录的变量
classes = ['五代诗词', '元曲', '全唐诗', '四书五经', '宋词', '幽梦影', '御定全唐詩', '曹操诗集', '楚辞', '水墨唐诗',
'纳兰性德', '蒙学', '论语', '诗经']
然后，我们可以遍历该数组，拼接一个目录，遍历目录中中的文件，再进行文件处理
for cls in classes:
dir = base_dir + cls
files = os.listdir(dir)
for f in files:
f = f'{dir}/{f}'
if os.path.isdir(f):
if 'error' in f:
continue
for ff in os.listdir(f):
process_json(f'{f}/{ff}')
else:
process_json(f)
上面代码中，我们遍历每个类别的目录后，会列出该类别中所有的文件，文件如果是一个目录，则继续遍历这个目录，因为作者提供的目录结构会存在二级目录的情况。
最后，拿到每个json文件后，会调用process_json()函数处理对应的json文件。下面我们开始介绍process_json()函数。
process_json()函数会对上面代码中拿到的每个json文件进行处理，并且从json文件中提取我们需要的信息（诗词的标题和内容），重新组织结构，写入到一个新文件中；该函数还会根据一个简单的策略划分出训练集和测试集(训练集用来训练我们的模型，测试集用来在训练过程中测试模型的性能)。整体代码如下
def process_json(file):
if not file.endswith('.json'):
return
with open(file, 'r') as f:
json_content = f.read()
array = json.loads(json_content)
if type(array) != list:
return
if len(array) > 100:
train_array = array[:-1]
test_array = array[-1:]
else:
train_array = array
test_array = None
for item in train_array:
if 'title' not in item.keys() or 'paragraphs' not in item.keys():
continue
write_file(item, dst_train_file)
if test_array is not None:
for item in test_array:
if 'title' not in item.keys() or 'paragraphs' not in item.keys():
continue
write_file(item, dst_test_file)
在代码中，首先会打开该json文件，并读取json文件中的内容；读到内容后，通过json.loads()函数将它解码成在python中可以识别的数据结构。
接下来，我们根据该分类下诗词的数据决定是否要划分出测试集，策略很简单，如果个数大于100，我们就把最后一个作为测试集的一部分，当然这个策略可以根据你的需求进行调整。
最后，我们从json中拿到title和paragraphs属性通过一个write_file()函数写到我们的新文件中。
write_file()函数的实现也很简单，作用就是拿到title和paragraphs，组织好结构写入到一个新文件中；我们预处理后的文件不会像原数据集那样提供多个文件，而是全部写到同一个文件中，所以，此时就得考虑一个问题：所有的诗词在一个文件中，怎么标识出一首诗结束了呢？办法很简单，我们在没首诗结束的时候添加一个<|endoftext|>特殊标识，该标识很重要，因为在后面我们训练模型的时候，该标识也会根据此标识学习一首诗到哪结束了(不需要结束，咱们模型就无止境的输出了)。
def write_file(item, dst_file):
global error_count
title = item['title']
paragraphs = item['paragraphs']
content = f'\n{title}'
for p in paragraphs:
content = f'{content}\n{p}'
content = converter.convert(content)
if '𫗋' in content:
print(f'{content}----')
error_count += 1
return
content = content + '<|endoftext|>'
dst_file.write(content)
上面代码中，处理前面我们介绍的部分，存在两个特殊的地方
...
content = converter.convert(content)
...
if '𫗋' in content
第一个的作用是将繁体中文转换成简体字，因为原数据集中存在大量的繁体字，显然，我们不想让咱们的模型生成的诗词是繁体字形式，所以这里我选择将繁体字转换成简体字，这里借助了一个python的转换库opencc实现，大家可以通过pip3 install opencc-python-reimplemented进行安装，该库的使用方法如下
import opencc
# 繁转简
converter = opencc.OpenCC('t2s')
content = converter.convert(content)
第二个特殊的地方就是我们代码中有一个𫗋，这是因为，通过上述代码转换成简体字的时候会有一些字转换错误，所以我们这里直接将存在转换错误情况的诗过滤掉，当然，这种情况不会很多，大概几十首诗词，对于咱们几十万首诗词的数据集来说都是毛毛雨。
好了，上面就是咱们数据预处理的全部过程，最终你会得到一个如下结构的train.txt和test.txt分别代表咱们前面提到过的训练集和测试集。
最后，我把全部代码整理出来，方便大家可以复制到本地直接运行
import os, json
import opencc
base_dir = 'chinese-poetry-master/'
classes = ['五代诗词', '元曲', '全唐诗', '四书五经', '宋词', '幽梦影', '御定全唐詩', '曹操诗集', '楚辞', '水墨唐诗',
'纳兰性德', '蒙学', '论语', '诗经']
dst_train_file = open('./train.txt', 'w')
dst_test_file = open('./test.txt', 'w')
converter = opencc.OpenCC('t2s')
error_count = 0
def write_file(item, dst_file):
global error_count
title = item['title']
paragraphs = item['paragraphs']
content = f'\n{title}'
for p in paragraphs:
content = f'{content}\n{p}'
content = converter.convert(content)
if '𫗋' in content:
print(f'{content}----')
error_count += 1
return
content = content + '<|endoftext|>'
dst_file.write(content)
def process_json(file):
if not file.endswith('.json'):
return
with open(file, 'r') as f:
json_content = f.read()
array = json.loads(json_content)
if type(array) != list:
return
if len(array) > 100:
train_array = array[:-1]
test_array = array[-1:]
else:
train_array = array
test_array = None
for item in train_array:
if 'title' not in item.keys() or 'paragraphs' not in item.keys():
continue
write_file(item, dst_train_file)
if test_array is not None:
for item in test_array:
if 'title' not in item.keys() or 'paragraphs' not in item.keys():
continue
write_file(item, dst_test_file)
for cls in classes:
dir = base_dir + cls
files = os.listdir(dir)
for f in files:
f = f'{dir}/{f}'
if os.path.isdir(f):
if 'error' in f:
continue
for ff in os.listdir(f):
process_json(f'{f}/{ff}')
else:
process_json(f)
dst_train_file.close()
dst_test_file.close()
dst_train_file = open('./train.txt', 'r')
dst_test_file = open('./test.txt', 'r')
train_count = 0
test_count = 0
for line in dst_train_file:
if '<|endoftext|>' in line:
train_count += 1
for line in dst_test_file:
if '<|endoftext|>' in line:
test_count += 1
print(f'train_count: {train_count}, test_count: {test_count}, error_count: {error_count}')
