textgenrnn

只需几行代码，即可在任何文本数据集上轻松训练任意大小和复杂度的文本生成神经网络，或者使用预训练模型快速对一段文本进行训练。

textgenrnn 是一个基于 Keras/TensorFlow 的 Python 3 模块，用于创建 char-rnn 模型，并提供了许多酷炫的功能：

• 一种现代化的神经网络架构，利用注意力加权和跳跃嵌入等新技术来加速训练并提升模型质量。
• 可以在字符级别或单词级别上进行训练和生成文本。
• 可以配置 RNN 的规模、RNN 层数以及是否使用双向 RNN。
• 支持训练任何通用的输入文本文件，包括大型文件。
• 可以在 GPU 上训练模型，然后在 CPU 上使用该模型生成文本。
• 在 GPU 上训练时，可以利用强大的 CuDNN 实现的 RNN，相比传统的 LSTM 实现，能够大幅加快训练速度。
• 使用上下文标签进行训练，使模型在某些情况下能够更快地学习并产生更好的结果。

您可以在这个 Colaboratory Notebook 中免费试用 textgenrnn，并使用 GPU 训练任意文本文件！阅读 这篇博客文章 或 观看这个视频 以获取更多信息！

示例

from textgenrnn import textgenrnn

textgen = textgenrnn()
textgen.generate()

【剧透】还有谁觉得这篇帖子和那个人有点超出我真正喜欢《星球大战》的程度？他们在火堆旁讨论健康问题，还发布了2016年游戏信件中关于我家后院的报道。

附带的模型可以轻松地在新文本上进行训练，并且即使只遍历一次输入数据，也能生成合适的文本。

textgen.train_from_file('hacker_news_2000.txt', num_epochs=1)
textgen.generate()

项目状态 项目 火狐浏览器

模型权重相对较小（磁盘上约 2 MB），可以轻松保存并在新的 textgenrnn 实例中加载。因此，您可以尝试使用经过数百次数据遍历训练过的模型。（事实上，textgenrnn 学习得 太好了，以至于需要显著提高温度才能获得更具创造性的输出！）

textgen_2 = textgenrnn('/weights/hacker_news.hdf5')
textgen_2.generate(3, temperature=1.0)

为什么我们得到了钱“定期改变”

乌尔堡向火狐浏览器收购了奈尔夫·穆尔蒂·沙姆恩

谷歌的伯恩推出的 Kubernetes

您还可以通过在任何训练函数中添加 new_model=True 来训练一个新模型，该模型支持单词级别的嵌入和双向 RNN 层。

交互模式

您也可以逐步参与到输出内容的生成过程中。交互模式会为您建议下一个字符或单词的 前 N 个选项，并允许您从中选择一个。

在终端中运行 textgenrnn 时，只需在 generate 方法中传入 interactive=True 和 top_n=N 即可。N 的默认值为 3。

from textgenrnn import textgenrnn

textgen = textgenrnn()
textgen.generate(interactive=True, top_n=5)

这可以为输出增添一份 人性化的色彩；感觉就像您自己在写作一样！(参考)

使用方法

textgenrnn 可以通过 pip 从 pypi 安装：

pip3 install textgenrnn

要使用最新版本的 textgenrnn，您必须安装至少 TensorFlow 2.1.0 版本。

您可以在 这个 Jupyter Notebook 中查看常见功能和模型配置选项的演示。

/datasets 目录包含用于训练 textgenrnn 的 Hacker News/Reddit 数据示例。

/weights 目录包含上述数据集上的进一步预训练模型，可以加载到 textgenrnn 中。

/outputs 目录包含由上述预训练模型生成的文本示例。

神经网络架构与实现

textgenrnn 基于 Andrej Karpathy 的 char-rnn 项目，并加入了一些现代优化，例如能够处理非常短的文本序列。

内置的预训练模型遵循一种受 DeepMoji 启发的 神经网络架构。对于默认模型，textgenrnn 接收最多 40 个字符的输入，将每个字符转换为一个 100 维的字符嵌入向量，然后将其输入到一个包含 128 个单元的长短期记忆（LSTM）循环层中。该层的输出再被送入另一个 128 单元的 LSTM 层。随后，所有三层的输出都会被送入一个注意力层，以加权最重要的时间特征并对其进行平均（由于嵌入和第一层 LSTM 被跳过连接到注意力层，模型更新可以更容易地反向传播到这些层，从而防止梯度消失）。最终的输出会被映射为多达 394 种不同字符 的概率分布，表示它们是序列中的下一个字符，包括大写字母、小写字母、标点符号和表情符号。（如果在新数据集上训练新模型，上述所有数值参数都可以进行配置）

或者，如果每个文本文档都提供了上下文标签，模型可以以上下文模式进行训练，在这种模式下，模型会学习“给定上下文”的文本，从而使循环层学会“去情境化”的语言。仅基于文本的路径可以利用这些去情境化的层；总体而言，这使得训练速度更快，且模型的定量和定性表现都优于仅基于文本进行训练的情况。

软件包中包含的模型权重是在来自 Reddit 提交的数十万份文本文档上训练得到的（通过 BigQuery），这些文档来自非常 多样化 的子版块。该网络还采用了上述的去情境化方法进行训练，以提高训练效率并减轻作者偏见的影响。

当使用 textgenrnn 在新的文本数据集上对模型进行微调时，所有层都会被重新训练。然而，由于原始的预训练网络已经具备更为扎实的基础知识，因此新的 textgenrnn 模型在训练过程中会更快、更准确地收敛，并且有可能学习到原始数据集中不存在的新关系（例如，预训练的字符嵌入包含了现代互联网语法中各种可能的语境信息）。

此外，重新训练采用基于动量的优化器和线性衰减的学习率，这两种方法都可以防止梯度爆炸，从而大大降低模型在长时间训练后发生发散的可能性。

注意事项

• 即使使用经过大量训练的神经网络，你也无法保证每次生成的文本质量都高。这也是为什么那些利用神经网络文本生成技术的病毒式 博客文章/Twitter 推文通常会生成大量文本，然后再从中挑选出最好的内容的原因。

• 不同数据集的结果差异很大。由于预训练的神经网络规模相对较小，它无法像一些博客文章中展示的那样存储大量数据。为了获得最佳效果，建议使用至少包含 2,000–5,000 个文档的数据集。如果数据集较小，则需要通过在调用训练方法时增加 num_epochs 参数或从头开始训练新模型来延长训练时间。即便如此，目前仍然没有一个可靠的启发式方法来判断一个模型是否“良好”。

• 重新训练 textgenrnn 并不需要 GPU，但在 CPU 上训练会花费更长的时间。如果你使用 GPU，建议增加 batch_size 参数，以更好地利用硬件资源。

textgenrnn 的未来计划

• 更正式的文档
• 使用 tensorflow.js 的基于 Web 的实现（由于网络规模较小，效果尤为理想）
• 一种可视化注意力层输出的方法，以便观察网络是如何“学习”的。
• 一种允许将模型架构用于聊天机器人对话的模式（可能会作为独立项目发布）。
• 进一步深化对上下文的理解（位置上下文以及支持多个上下文标签）。
• 一个更大的预训练网络，能够处理更长的字符序列，并对语言有更深入的理解，从而生成更好的句子。
• 针对词级模型的层次 softmax 激活函数（一旦 Keras 对其提供良好支持）。
• 使用 FP16 进行超快速训练（适用于 Volta/TPU，一旦 Keras 对其提供良好支持）。

使用 textgenrnn 的文章/项目

文章

• Lifehacker：如何训练你自己的神经网络——Beth Skwarecki
• 纽约时报：让我们的算法为你选择万圣节服装——Janelle Shane
• CNN Business：这个古怪的实验凸显了人工智能面临的最大挑战——Rachel Metz

项目

• 推文生成器 —— 训练一个针对 Twitter 用户优化的神经网络，用于生成推文
• Hacker News 模拟器 —— 一个基于 textgenrnn 训练的 Twitter 机器人，使用超过 30 万条 Hacker News 提交内容
• SubredditRNN —— 一个 Reddit 子版块，所有提交的内容均由 textgenrnn 生成的机器人发布
• 人机协作披萨 —— 使用 textgenrnn 生成食谱并在现实中制作的披萨
• 桌游标题
• 视频游戏讨论论坛标题
• 人工智能创作的蛋糕
• 人工智能创作的饼干
• 人工智能生成的歌曲

推文

• BuzzFeed YouTube 视频
• AWS 服务
• 食谱 + D&D 法术 + 重金属乐队名
• RPG 冒险名称
• 洋葱报 + 美丽佳人杂志
• 谷歌会议室名称
• 西斯尊主

维护者/创作者

Max Woolf (@minimaxir)

Max 的开源项目由他的 Patreon 支持。如果您觉得这个项目有帮助，欢迎向 Patreon 捐款，您的支持将用于富有创意的用途。

致谢

Andrej Karpathy 通过博客文章 循环神经网络的不合理有效性 提出了 char-rnn 的原始构想。

Daniel Grijalva 贡献了交互模式 (pull 请求 #52)。

许可证

MIT 许可证

注意层代码取自 DeepMoji（MIT 许可）。

textgenrnn 快速上手指南

textgenrnn 是一个基于 Keras/TensorFlow 的 Python 模块，旨在让用户通过几行代码轻松训练任意大小和复杂度的文本生成神经网络（Char-RNN）。它支持字符级或词级训练，并针对 GPU 训练进行了深度优化。

环境准备

• 操作系统：Linux, macOS 或 Windows
• Python 版本：Python 3.6+
• 核心依赖：
  • TensorFlow >= 2.1.0 (必需)
  • Keras (通常随 TensorFlow 安装)
• 硬件建议：
  • 训练：推荐使用 NVIDIA GPU 以利用 CuDNN 加速，大幅缩短训练时间。
  • 推理：CPU 即可流畅运行生成任务。

安装步骤

使用 pip 进行安装。国内开发者建议使用清华源或阿里源以提升下载速度。

# 使用默认源安装
pip3 install textgenrnn

# 推荐：使用清华大学镜像源加速安装
pip3 install textgenrnn -i https://pypi.tuna.tsinghua.edu.cn/simple

注意：安装前请确保您的环境中已安装满足版本要求的 TensorFlow (pip3 install "tensorflow>=2.1.0")。

基本使用

1. 快速生成文本（使用预训练模型）

加载默认预训练模型并直接生成文本，无需任何训练数据。

from textgenrnn import textgenrnn

textgen = textgenrnn()
textgen.generate()

2. 微调模型（基于自定义数据集）

只需一个文本文件即可开始训练。以下示例展示如何读取文件并进行单轮 epoch 训练，随后生成内容。

from textgenrnn import textgenrnn

textgen = textgenrnn()

# 训练模型，num_epochs 控制训练轮数
textgen.train_from_file('hacker_news_2000.txt', num_epochs=1)

# 生成文本
textgen.generate()

3. 加载已保存的模型权重

您可以保存训练好的模型权重（.hdf5 文件），并在后续会话中直接加载使用。

from textgenrnn import textgenrnn

# 加载指定权重的模型
textgen_2 = textgenrnn('/weights/hacker_news.hdf5')

# temperature 参数控制生成的随机性（值越大越具创造性）
textgen_2.generate(3, temperature=1.0)

4. 交互式生成模式

在终端运行时，可以开启交互模式，让 AI 提供多个候选词/字符供用户选择，实现人机协作写作。

from textgenrnn import textgenrnn

textgen = textgenrnn()
# top_n 设置每次提供的候选选项数量，默认为 3
textgen.generate(interactive=True, top_n=5)