code/09-transformer/prepare_poems.py

"""准备训练数据：下载《全唐诗》，转成简体，只留下格律整齐的五言、七言绝句和律诗。

    python prepare_poems.py
准备：uv add opencc
数据来自 https://github.com/chinese-poetry/chinese-poetry （MIT 许可），约 25 MB，只在第一次运行时下载。
结果写到 .cache/poems.txt，一行一首诗。
"""
import json
import random
import re
import urllib.request
from pathlib import Path

import opencc

CACHE = Path(__file__).parent / ".cache"
RAW = CACHE / "tang"
URL = "https://raw.githubusercontent.com/chinese-poetry/chinese-poetry/master/%E5%85%A8%E5%94%90%E8%AF%97/poet.tang.{}.json"

RAW.mkdir(parents=True, exist_ok=True)
for i in range(0, 58000, 1000):
    path = RAW / f"poet.tang.{i}.json"
    if not path.exists():
        print(f"下载 {path.name}")
        urllib.request.urlretrieve(URL.format(i), path)

# 原始数据是繁体。OpenCC 按词组转换，比逐字替换准确（"餘"在"其餘"里转成"余"，"乾坤"的"乾"不会被转成"干"）
to_simple = opencc.OpenCC("t2s")
SENTENCE = re.compile(r"[一-鿿]+[，。]")  # 一句：若干汉字，加一个逗号或句号

total, kept = 0, []
for path in sorted(RAW.glob("poet.tang.*.json")):
    for poem in json.loads(path.read_text()):
        total += 1
        text = to_simple.convert("".join(poem["paragraphs"]))
        sentences = SENTENCE.findall(text)
        if "".join(sentences) != text:  # 有汉字和逗号句号以外的东西（注释、缺字符号、括号），不要
            continue
        lengths = {len(s) - 1 for s in sentences}
        # 五言或七言，四句（绝句）或八句（律诗），逗号句号交替
        if lengths in ({5}, {7}) and len(sentences) in (4, 8) and \
                all(s[-1] == "，。"[i % 2] for i, s in enumerate(sentences)):
            kept.append(text)

kept = sorted(set(kept))  # 去掉重复收录的诗
random.Random(0).shuffle(kept)
(CACHE / "poems.txt").write_text("\n".join(kept) + "\n")

chars = set("".join(kept))
print(f"《全唐诗》共 {total} 首，留下格律整齐的 {len(kept)} 首，共 {sum(map(len, kept))} 个字符，{len(chars)} 个不同的字符")
for kind, n_char, n_sent in [("五言绝句", 5, 4), ("七言绝句", 7, 4), ("五言律诗", 5, 8), ("七言律诗", 7, 8)]:
    print(f"  {kind}：{sum(1 for p in kept if len(p) == (n_char + 1) * n_sent)} 首")
print("前三首：")
for p in kept[:3]:
    print("  " + p)