gpt-2-simple

这是一个简单的 Python 包，用于封装 OpenAI 的 GPT-2 文本生成模型（特别是“small”1.24亿参数版本和“medium”3.55亿参数版本）的现有模型微调与生成脚本。此外，该包还支持更便捷地生成文本，可将生成结果保存到文件以便于后续整理，并允许指定前缀以强制文本以特定短语开头。

此包对以下内容进行了整合，并仅做了少量底层改动：

• 来自 OpenAI 官方 GPT-2 仓库（MIT 许可证）的模型管理
• 来自 Neil Shepperd 分支的 GPT-2 微调代码（MIT 许可证）
• 来自 textgenrnn 的文本生成输出管理（MIT 许可证／同样由我开发）

对于微调，强烈建议使用 GPU，尽管您也可以用 CPU 进行生成（但速度会慢很多）。如果您在云端训练，强烈推荐使用 Colaboratory Notebook 或配备 TensorFlow 深度学习镜像的 Google Compute Engine 虚拟机。（因为 GPT-2 模型托管在 GCP 上）

您可以使用 gpt-2-simple 在 此 Colaboratory Notebook 中免费利用 GPU 重新训练模型，该 Notebook 还演示了该包的其他功能。

注意：gpt-2-simple 的开发工作已基本被 aitextgen 取代，后者具备相似的 AI 文本生成能力，且训练时间和资源消耗更加高效。如果您不需要使用 TensorFlow，我建议改用 aitextgen。使用 gpt-2-simple 训练得到的检查点同样可以 通过 aitextgen 加载。

安装

gpt-2-simple 可通过 PyPI 安装：

pip3 install gpt-2-simple

此外，您还需要为您的系统安装对应的 TensorFlow 2.X 版本（最低 2.5.1），例如 tensorflow 或 tensorflow-gpu。

使用方法

以下是一个示例，展示如何将模型下载到本地、在数据集上进行微调并生成一些文本。

警告：预训练的 1.24 亿参数模型及其微调后的模型大小均为 500 MB！（预训练的 3.55 亿参数模型则为 1.5 GB）

import gpt_2_simple as gpt2
import os
import requests

model_name = "124M"
if not os.path.isdir(os.path.join("models", model_name)):
    print(f"正在下载 {model_name} 模型...")
    gpt2.download_gpt2(model_name=model_name)   # 模型会保存到当前目录下的 /models/124M/

file_name = "shakespeare.txt"
if not os.path.isfile(file_name):
    url = "https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"
    data = requests.get(url)

    with open(file_name, 'w') as f:
        f.write(data.text)


sess = gpt2.start_tf_sess()
gpt2.finetune(sess,
              file_name,
              model_name=model_name,
              steps=1000)   # steps 是最大训练步数

gpt2.generate(sess)

生成的模型检查点默认保存在 /checkpoint/run1。如果您想从该文件夹加载模型并生成文本：

import gpt_2_simple as gpt2

sess = gpt2.start_tf_sess()
gpt2.load_gpt2(sess)

gpt2.generate(sess)

与 textgenrnn 类似，您可以通过设置 return_as_list 参数来生成并保存文本，以便后续使用（例如作为 API 或机器人）。

single_text = gpt2.generate(sess, return_as_list=True)[0]
print(single_text)

如果您希望在 checkpoint 文件夹中存储或加载多个模型，可以在 finetune 和 load_gpt2 中传入 run_name 参数。

此外，该包还提供了命令行界面，支持微调与生成操作，默认配置适合直接在配备 GPU 的云虚拟机上运行。对于微调（如果模型尚未下载，则会自动下载）：

gpt_2_simple finetune shakespeare.txt

对于生成操作，会将生成的文本保存到 gen 文件夹中：

gpt_2_simple generate

大部分函数中可用的参数同样支持作为 CLI 参数，例如：

gpt_2_simple generate --temperature 1.0 --nsamples 20 --batch_size 20 --length 50 --prefix "<|startoftext|>" --truncate "<|endoftext|>" --include_prefix False --nfiles 5

请参阅下方了解部分 CLI 参数的具体作用。

注意：如果您想在另一个数据集上微调或加载另一个模型，请先重启 Python 会话。

gpt-2-simple 与其他文本生成工具的区别

GPT-2 生成文本的方式与其他工具（如 textgenrnn）略有不同——具体而言，textgenrnn 是在 GPU 上直接生成完整的文本序列，然后再进行解码，这种做法很难在不修改底层模型代码的情况下修复。因此：

• 总体而言，GPT-2 在整个生成过程中更能保持上下文连贯性，特别适合生成对话类文本。此外，生成的文本语法正确，大小写规范，错别字也较少。
• 原始 GPT-2 模型是在种类极其丰富的数据源上训练的，这使得模型能够融入输入文本中未出现过的习语。
• GPT-2 每次请求最多只能生成 1024 个标记（约相当于 3-4 段英文文本）。
• GPT-2 无法在遇到特定结束标记时提前停止。（解决方法：在 generate 函数中传入 truncate 参数，以仅收集到指定结束标记为止的文本。你可能需要适当调小 length 参数。）
• 较高的温度设置（例如 0.7 - 1.0）能生成更有趣的文本，而其他框架的最佳温度范围则在 0.2 - 0.5 之间。
• 在微调 GPT-2 时，它无法感知更大文本中文档的开头和结尾。你需要使用自定义字符序列来标记文档的起始和结束。然后，在生成时，你可以指定一个 prefix 来定位文档开头标记，以及一个 truncate 来定位文档结束标记。你还可以将 include_prefix=False 设置为 True，以便在生成时忽略前缀标记（例如，如果你不想生成 <|startoftext|> 这样的标记）。
• 如果你向 finetune() 传递一个单列的 .csv 文件，它会自动将 CSV 解析成适合 GPT-2 训练的格式（包括在每个文本文档前添加 <|startoftext|>，并在末尾添加 <|endoftext|>，因此上述 truncate 技巧在生成输出时非常有用）。这样做是为了正确处理每个文本文档中的引号和换行符。
• GPT-2 支持通过设置可被 nsamples 整除的 batch_size 来并行生成文本，从而大幅提升生成速度。配合 GPU 使用效果极佳（在 Colaboratory 的 K80 上，batch_size 最高可设为 20！）。
• 由于 GPT-2 的架构设计，它能很好地适应更强大的 GPU。对于 1.24 亿参数的模型，如果你想长时间训练，GCP 的 P100 GPU 比 K80/T4 快约 3 倍，但价格仅贵 3 倍，性价比相当（V100 比 P100 快约 1.5 倍，但价格贵约 2 倍）。P100 即使在 batch_size=1 时也能充分利用 GPU，而 V100 的利用率约为 88%。
• 如果你已有部分训练好的 GPT-2 模型，并希望继续微调，可以将 overwrite=True 设置为 True，这样就能在不创建副本的情况下继续训练并覆盖之前的模型迭代。这在迁移学习中尤其有用（例如，先用一个数据集对 GPT-2 进行深度微调，再用另一个数据集微调，实现两个数据集的“融合”）。
• 如果你的输入文本数据集非常庞大（超过 100 MB），你可以使用 gpt2.encode_dataset(file_path) 对数据集进行预编码和压缩。输出是一个压缩后的 .npz 文件，加载到 GPU 中进行微调的速度会快很多。
• 7.74 亿参数的“large”模型或许支持微调，因为现代 GPU 可能会因内存不足而崩溃（如果你在 Colaboratory 上使用 P100 GPU，可能会幸运一些）。不过，你仍然可以使用默认预训练模型进行生成，只需调用 gpt2.load_gpt2(sess, model_name='774M') 和 gpt2.generate(sess, model_name='774M')。
• 15.58 亿参数的“extra large”真模型可能无法直接在 Colaboratory Notebook 自带的 GPU 上运行。还需要更多测试来确定最适合它的配置。

使用 gpt-2-simple 的交互式应用

• gpt2-small — 使用默认 GPT-2 1.24 亿参数预训练模型的应用
• gpt2-reddit — 根据指定的 subreddit 和/或关键词生成 Reddit 标题的应用
• gpt2-mtg — 生成万智牌卡牌的应用

使用 gpt-2-simple 的文本生成示例

• ResetEra — 生成的电子游戏论坛讨论帖（GitHub 包含转储文件）
• /r/legaladvice — 标题生成（GitHub 包含转储文件）
• Hacker News — 数以万计的 Hacker News 提交标题生成

维护者/创作者

Max Woolf (@minimaxir)

Max 的开源项目由他的 Patreon 支持。如果你觉得这个项目对你有帮助，任何金钱上的捐助都将受到欢迎，并会被用于有意义的创作用途。

许可协议

MIT

免责声明

本仓库与 OpenAI 无任何关联或合作关系。

gpt-2-simple 快速上手指南

环境准备

• 系统要求：建议使用 NVIDIA GPU（如 T4/P100），CPU 也可运行但速度较慢
• 前置依赖：Python 3.6+，TensorFlow 2.5.1+
• 国内加速推荐：使用清华源加速安装依赖

安装步骤

pip install gpt_2_simple tensorflow -i https://pypi.tuna.tsinghua.edu.cn/simple

基本使用

import gpt_2_simple as gpt2
import requests

# 下载示例数据（莎士比亚文本）
url = "https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"
with open("shakespeare.txt", "w") as f:
    f.write(requests.get(url).text)

# 下载124M预训练模型（首次运行自动下载，约500MB）
gpt2.download_gpt2(model_name="124M")

# 启动TensorFlow会话
sess = gpt2.start_tf_sess()

# 微调模型（1000步）
gpt2.finetune(sess, "shakespeare.txt", steps=1000)

# 生成文本并打印
generated_text = gpt2.generate(sess, return_as_list=True)[0]
print(generated_text)

注意：若需使用自己的数据，替换 shakespeare.txt 为本地文件路径。首次运行会自动下载模型文件。