ver1: 实现深度学习训练框架，支持 Wiki GPT 与诗歌生成双任务

模型架构：
- 手写 Mini GPT（Transformer）：实现 PositionalEmbedding、TransformerDecoder 组件
- 手写 RNN：LSTM 堆叠结构
训练系统：
- Pipeline 框架：数据加载、Tokenizer、训练、生成全流程封装
- Checkpoint 机制：支持断点续训、分代模型保存和加载
- TensorBoard 训练监控
任务实现：
- Wiki GPT：基于中文维基语料的中文文本生成
- 诗歌生成器（GPT）：基于 Transformer 的诗歌生成
- 诗歌生成器（RNN）：基于 LSTM 的诗歌生成
工程支持：
- Gradio 交互界面
- Hugging Face Space 部署配置
- pytest 测试套件
其他特点：
- 多种采样方法：top-k、随机采样（temperature）、贪婪搜索
- Pipeline 重要支持组件：ModelBuidler、DataBundle

.gitattributes

@@ -0,0 +1 @@
+*.keras filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,3 @@
+__pycache__
+/.idea
+/local

AGENTS.md

@@ -0,0 +1,220 @@
+# Agent 编码规范
+
+## 防御性编程精简
+
+避免过度防御性编程，遵循以下原则：
+
+### 1. None 检查
+
+- **不要**进行显式的 None 检查
+- 信任输入数据，让程序在真正的错误点上失败
+- 避免 `if x is not None:` 这样的防御性代码
+
+```python
+# ❌ 避免
+def process(data):
+    if data is not None:
+        return data.value
+    return None
+
+# ✅ 推荐
+def process(data):
+    return data.value
+```
+
+### 2. 类型检查
+
+- **不要**使用 `isinstance`、`type()`、`typeof` 等进行运行时类型检查
+- 依靠类型提示和静态类型检查工具（如 mypy）
+- 让 Duck Typing 发挥作用
+
+```python
+# ❌ 避免
+def calculate(obj):
+    if isinstance(obj, int):
+        return obj * 2
+    elif isinstance(obj, str):
+        return obj * 2
+    else:
+        raise TypeError("不支持的类型")
+
+# ✅ 推荐
+def calculate(obj: int | str) -> int | str:
+    return obj * 2
+```
+
+### 3. 异常处理
+
+- **不要**滥用 try-catch 来压制异常
+- **不要**用 try-catch 让程序"容错"运行
+- 只在真正需要处理异常的地方捕获
+- 让未处理的异常自然抛出，暴露真正的问题
+
+```python
+# ❌ 避免 - 压制异常
+import logging
+
+logger = logging.getLogger(__name__)
+
+def parse_config(path):
+    try:
+        with open(path) as f:
+            return json.load(f)
+    except Exception as e:
+        logger.error(f"加载配置失败: {e}")
+        return {}  # 返回空配置让程序继续运行
+
+# ✅ 推荐 - 让异常传播
+def parse_config(path):
+    with open(path) as f:
+        return json.load(f)
+
+# ✅ 或仅在必要时转换异常类型
+def parse_config(path):
+    try:
+        with open(path) as f:
+            return json.load(f)
+    except json.JSONDecodeError as e:
+        raise ConfigError(f"配置文件格式错误: {e}") from e
+```
+
+### 4. 原则总结
+
+1. **早失败（Fail Fast）** - 让错误尽早暴露，不要试图掩盖
+2. **信任调用方** - 假设调用方会提供正确的输入
+3. **清晰错误信息** - 让异常信息直接指出问题所在
+4. **代码简洁** - 减少不必要的检查代码，专注于业务逻辑
+
+---
+
+**核心信条**：清晰的代码比健壮的代码更重要。让错误暴露，让问题可见。
+
+## 编辑文件时的精准修改原则
+
+在进行代码编辑时，**只修改必要的部分**，不要进行任何无关改动：
+
+### 禁止的无关改动
+
+- **不要**调整代码缩进或格式
+- **不要**重排 import 语句的顺序
+- **不要**添加或删除空行
+- **不要**修改注释（除非任务明确要求）
+- **不要**修改变量名、函数名等标识符（除非任务明确要求）
+- **不要**进行任何代码重构（除非任务明确要求）
+
+### ✅ 正确示例
+
+如果任务是将 `import config` 改为 `from mini_gpt import config`：
+
+```python
+# 修改前
+import config
+from typing import Callable
+
+# 修改后 - 只修改 import 语句，其他保持不变
+from mini_gpt import config
+from typing import Callable
+```
+
+### ❌ 错误示例
+
+#### 示例1：无关地调整 import 顺序
+
+```python
+# 修改前
+import config
+from typing import Callable
+
+# 错误 - 无关地调整了 import 顺序
+from typing import Callable
+from mini_gpt import config
+```
+
+#### 示例2：无关地修改函数参数格式
+
+```python
+# 修改前
+def my_function(
+        param1,
+        param2,
+        param3,
+):
+    pass
+
+# 错误 - 任务只要求修改函数体，却无关地修改了参数格式
+def my_function(
+    param1,  # 调整了缩进宽度
+    param2,
+    param3   # 去掉了尾部逗号
+):
+    pass
+```
+
+**原则**：最小化改动范围，只改必须改的地方。
+
+## 运行单元测试
+
+本项目使用 pytest 运行单元测试，必须在 `mini-gpt` conda 环境中执行。
+
+### 运行命令
+
+```bash
+/Users/run/anaconda3/envs/mini-gpt/bin/python -m pytest test/ -v
+```
+
+### 重要提示
+
+1. **必须使用 mini-gpt 环境** - 基础环境缺少 tensorflow 依赖，会导致测试收集失败
+2. **不要添加 `pytest.importorskip("tensorflow")`** - 这些测试依赖 tensorflow，跳过会掩盖真正的问题
+
+## Python 代码风格
+
+### 禁止尾逗号
+
+**任何情况下都不应出现尾逗号**（trailing comma）。
+
+```python
+# ❌ 避免 - 尾逗号
+my_list = [
+    1,
+    2,
+    3,
+]
+
+# ✅ 推荐
+my_list = [
+    1,
+    2,
+    3
+]
+
+# ❌ 避免 - 函数参数尾逗号
+def my_func(
+    arg1,
+    arg2,
+):
+    pass
+
+# ✅ 推荐
+def my_func(
+    arg1,
+    arg2
+):
+    pass
+
+# ❌ 避免 - 字典尾逗号
+my_dict = {
+    "key1": "value1",
+    "key2": "value2",
+}
+
+# ✅ 推荐
+my_dict = {
+    "key1": "value1",
+    "key2": "value2"
+}
+```
+
+## 禁止命令行参数
+
+永远不要在代码中使用命令行参数（如 `argparse`、`sys.argv` 等）。配置应通过代码中硬编码实现。

README.md

@@ -0,0 +1,146 @@
+---
+title: General Deep Learning
+emoji: 🏃
+colorFrom: yellow
+colorTo: gray
+sdk: gradio
+sdk_version: 6.12.0
+app_file: app.py
+pinned: false
+license: mit
+short_description: General Deep Learning is a practical deep learning experimen
+---
+
+# 通用深度学习（General Deep Learning）
+
+## 项目简介
+
+**通用深度学习（General Deep Learning）** 是一个面向实践的深度学习实验平台，致力于打造"训练-部署-体验"一体化的完整工作流。
+
+### ✨ 为什么适合你？
+
+**🎯 我的愿景**
+- 构建一个**从零开始、透明可学、工程模块化**的深度学习平台。
+
+**🎓学习友好**
+- ✅ **纯手工从零构建** - Transformer、RNN 都是一行行代码手撸
+- ✅ **代码即教程** - 没有黑盒封装，每个组件清晰可见
+- ✅ **完整的训练闭环** - 从数据处理到部署，全流程透明
+-
+**🔧 技术特性**
+- ✅ **覆盖主流模型** - Transformer、RNN，未来将扩展至 CNN、Diffusion 等
+- ✅ **模块化架构** - 可插拔设计，新模型/新数据集快速接入
+- ✅ **生产级部署** - 一键部署到 Hugging Face，支持断点续训、TensorBoard 监控
+
+### 📅 关于这个项目
+
+> *历时俩月，忙里偷闲。*
+
+这不是一个追求最新模型的项目，而是一个**"代码即教程"**的个人实验场。
+
+**已完成功能**：
+- Wiki GPT - 基于中文维基的手写 Transformer
+- 诗歌生成器 - GPT 和 RNN 双版本对比
+
+**未来规划**：
+4 月有事不再投入，5 月开始计划每月新增一个模型，探索更多架构（CNN、Diffusion...）
+
+- 🔮 逐步扩展至 CV、多模态等领域
+- 🔮 保持"从零手撸"的风格，让每个新模型都成为学习素材
+
+**欢迎一起折腾** —— 反馈问题、贡献代码，或单纯聊聊技术！
+
+### 🤗 在线体验
+
+[![Hugging Face Space](https://img.shields.io/badge/🤗-Hugging%20Face%20Space-blue)](https://huggingface.co/spaces/yetrun/general-deep-learning)
+
+🚀 **在线体验**：[点击访问 Hugging Face Space](https://huggingface.co/spaces/yetrun/general-deep-learning)
+
+本项目已部署到 Hugging Face Space，你可以在线体验以下功能：
+
+- **Wiki GPT 文本生成**：基于 Transformer 架构的中文文本生成，训练数据来自中文维基语料库
+- **诗歌生成器（GPT）**：基于 Transformer 的中文诗歌生成，支持五言、七言诗等
+- **诗歌生成器（RNN）**：基于 RNN 架构的中文诗歌生成，支持五言、七言诗等
+
+## 部署说明
+
+本项目已配置为 Hugging Face Space 兼容格式，如需更新部署：
+
+```bash
+# 1. 在 Hugging Face 创建新的 Space（选择 Gradio SDK）
+# 2. 绑定 Space 远程仓库
+git remote add huggingface https://huggingface.co/spaces/YOUR_USERNAME/YOUR_SPACE_NAME
+# 3. 确保依赖同步（生成 requirements.txt）
+python3 generate_requirements.py
+# 4. 提交并推送
+git push huggingface master
+```
+
+## 本地开发
+
+### Conda 环境使用
+
+使用方法：
+```bash
+# 创建环境
+conda env create -f <environment.yml>
+# 激活环境
+conda activate general-dl
+# 更新 environment.yml
+conda env update -f <environment.yml> --prune
+```
+
+上述 `<environment.yml>` 是环境配置文件的路径，需要替换成实际的文件名：
+
+- 如果你是本地开发，使用 `environment.yml`（Mac Intel 64 环境，`ENV=test`）
+- 如果你是在远程服务器上运行，使用 `environments-linux.yml`（Linux 服务器环境，`ENV=production`）
+
+> **插曲：**
+> 
+> 环境配置出现了问题，强制重新安装 tensorflow-text 才修复。
+> 
+> ```bash
+> pip uninstall tensorflow-text -y
+> pip install --no-cache-dir --force-reinstall tensorflow-text==2.20.0
+> ```
+
+### 开发工具配置
+
+#### TensorBoard 说明
+
+训练时，调用 `tensorboard --logdir=<logdir>` 来启动 TensorBoard，默认访问地址是 http://localhost:6006/.
+
+`<logdir>` 通常是 `local/tasks/<project_name>/tensorboard`.
+
+> 冷知识：tensorboard 中的代数与我们常规认为的代数不一致，第一代的计数是 0.
+
+#### JetBrains 远程开发配置
+
+配置本地代码映射：
+
+1. 菜单栏：Tools → Deployment → Configuration
+2. 配置目录映射：切换到Mappings标签页，Deployment path 设置远程目录路径
+3. 配置排除目录，一般可排除的本地目录包括：`data/dev`, `local`, `test`.
+
+手工同步：
+
+- 右键文件/目录 → Deployment → Upload to...
+
+## 数据集说明
+
+### WIKI 数据集 
+
+*（本项目中 `wiki_gpt` 任务使用了中文维基语料库进行训练）*
+
+下载维基百科的数据。
+
+```bash
+wget https://dumps.wikimedia.org/other/mediawiki_content_current/zhwiki/2026-01-01/xml/bzip2/zhwiki-2026-01-01-p1p5254490.xml.bz2
+wget https://dumps.wikimedia.org/other/mediawiki_content_current/zhwiki/2026-01-01/xml/bzip2/zhwiki-2026-01-01-p5254491p9382552.xml.bz2
+```
+
+维基百科的数据分成两个文件，可使用 cat 命令合并成一个文件：
+
+```bash
+cat zhwiki-2026-01-01-p1p5254490.xml.bz2 zhwiki-2026-01-01-p5254491p9382552.xml.bz2 > zhwiki-2026-01-01.xml.bz2
+```

app.py

@@ -0,0 +1,58 @@
+"""
+AI 文本生成工具集 - 多页面 Gradio 应用
+
+入口点，提供导航到各个子应用：
+- /：首页导航
+- /wiki_gpt：Wiki GPT 文本生成器
+- /poetry_gpt：诗歌生成器（GPT）
+- /poetry_rnn：诗歌生成器（RNN）
+
+特点：
+- 每个子页面可以独立运行测试
+"""
+import gradio as gr
+
+from tasks.wiki_gpt.gradio import demo as wiki_gpt_demo
+from tasks.poetry_gpt.gradio import demo as poetry_gpt_demo
+from tasks.poetry_rnn.gradio import demo as poetry_rnn_demo
+
+
+with gr.Blocks(title="AI 文本生成工具集") as demo:
+    gr.Markdown("# AI 文本生成工具集")
+    gr.Markdown("请选择要使用的应用：")
+
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("## 诗歌生成器（GPT）")
+            gr.Markdown("基于 Transformer 的中文诗歌生成，支持五言、七言诗等。")
+            gr.Button("进入诗歌生成器", link="/poetry_gpt")
+
+        with gr.Column():
+            gr.Markdown("## 诗歌生成器（RNN）")
+            gr.Markdown("基于 RNN 的中文诗歌生成，支持五言、七言诗等。")
+            gr.Button("进入诗歌生成器", link="/poetry_rnn")
+
+        with gr.Column():
+            gr.Markdown("## Wiki GPT 文本生成")
+            gr.Markdown("基于 Transformer 的中文文本生成，训练来自于中文维基语料库。")
+            gr.Button("进入 Wiki GPT", link="/wiki_gpt")
+
+    gr.Markdown("---")
+    gr.Markdown("### 说明")
+    gr.Markdown("每个应用都是独立加载的，进入页面后需要等待模型加载完成。")
+
+
+with demo.route("诗歌生成器（GPT）", "/poetry_gpt"):
+    poetry_gpt_demo.render()
+
+
+with demo.route("诗歌生成器（RNN）", "/poetry_rnn"):
+    poetry_rnn_demo.render()
+
+
+with demo.route("Wiki GPT", "/wiki_gpt"):
+    wiki_gpt_demo.render()
+
+
+if __name__ == "__main__":
+    demo.launch()

data/__init__.py

@@ -0,0 +1,25 @@
+"""数据集模块
+
+提供统一的数据集接口，包括 Wiki 和诗歌数据集。
+
+Usage:
+    from data import WikiDataset, PoetryDataset
+
+    # Wiki 数据集
+    wiki = WikiDataset(data_dir="~/data/wiki/mini_c4")
+    doc_ds = wiki.doc_ds()
+    tokens_ds = wiki.tokens_ds(seq_length=256, batch_size=32)
+    wiki.stat(seq_length=256)
+
+    # 诗歌数据集
+    poetry = PoetryDataset(data_dir="~/data/Poetry")
+    doc_ds = poetry.doc_ds()
+    tokens_ds = poetry.tokens_ds(seq_length=100, batch_size=128)
+    poetry.stat(seq_length=100)
+"""
+
+from data.base import DataBundle, TokenizerBundle
+from data.wiki import WikiDataset
+from data.poetry import PoetryDataset
+
+__all__ = ["DataBundle", "TokenizerBundle", "WikiDataset", "PoetryDataset"]

data/base.py

@@ -0,0 +1,85 @@
+"""数据集抽象基类模块
+
+定义 DataBundle 抽象基类，统一数据集的接口规范。
+每个具体的数据集（如 Wiki、诗歌）都应该继承此类并实现相应方法。
+"""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Callable, Optional
+
+import tensorflow as tf
+
+
+@dataclass
+class TokenizerBundle:
+    """分词器信息包装类
+
+    将分词器相关的属性打包在一起，简化 DataBundle 接口。
+    """
+
+    tokenizer: Callable
+    decode: Callable
+    end_of_text: int
+    vocab_size: int
+    vocab_path: str = ""
+
+
+@dataclass
+class DataBundle(ABC):
+    """数据集抽象基类
+
+    将数据加载、分词、统计等功能绑定在一起，提供统一的数据集接口。
+
+    Usage:
+        dataset = WikiDataset(data_dir="~/data/wiki")
+        doc_ds = dataset.doc_ds()
+        tokens_ds = dataset.tokens_ds(seq_length=256, batch_size=32)
+        dataset.stat()
+    """
+
+    data_dir: str
+    sequence_length: int = 256
+
+    @abstractmethod
+    def doc_ds(self) -> tf.data.Dataset:
+        """返回原始文档数据集
+
+        Returns:
+            TensorFlow Dataset，每个元素是一个文档字符串
+        """
+        pass
+
+    @abstractmethod
+    def tokens_ds(self, seq_length: int, batch_size: int) -> tf.data.Dataset:
+        """返回 tokenized 数据集
+
+        将原始文档转换为 token ID 序列，并分割为训练样本。
+
+        Args:
+            seq_length: 序列长度
+            batch_size: 批次大小
+
+        Returns:
+            TensorFlow Dataset，每个元素是 (input_ids, target_ids) 对
+        """
+        pass
+
+    @abstractmethod
+    def tokenizer_bundle(self) -> TokenizerBundle:
+        """返回分词器信息"""
+        pass
+
+    def stat(self, seq_length: int | None = None) -> None:
+        """打印数据集统计信息
+
+        Args:
+            seq_length: 序列长度，用于估算训练样本数
+        """
+        from data.common import collect_stats
+
+        info = self.tokenizer_bundle()
+        stats = collect_stats(
+            name=self.__class__.__name__, loader=self.doc_ds, tokenizer=info.tokenizer
+        )
+        stats.print_report(seq_length=seq_length)

data/common.py

@@ -0,0 +1,147 @@
+"""数据集共享工具模块
+
+提供数据集统计、报告生成等共享功能。
+"""
+
+import pathlib
+from dataclasses import dataclass
+from typing import Callable
+
+import numpy as np
+import tensorflow as tf
+from keras import layers
+
+
+@dataclass
+class DatasetStats:
+    """数据集统计结果"""
+
+    name: str
+    doc_count: int
+    total_chars: int
+    total_tokens: int
+    max_length: int
+    median_length: int
+
+    def print_report(self, seq_length: int | None = 256):
+        """打印统一格式的统计报表
+
+        Args:
+            seq_length: 序列长度，用于估算训练样本数。
+                       为 None 时表示不切割，一个文档一个样本。
+        """
+        avg_chars = self.total_chars / self.doc_count if self.doc_count > 0 else 0
+        avg_tokens = self.total_tokens / self.doc_count if self.doc_count > 0 else 0
+
+        print()
+        print("=" * 60)
+        print(f"{self.name} 数据集统计")
+        print("=" * 60)
+        print(f"{'文档数:':<20} {self.doc_count:>15,}")
+        print(f"{'总字符数:':<20} {self.total_chars:>15,}")
+        print(f"{'总 Token 数:':<20} {self.total_tokens:>15,}")
+        print("-" * 60)
+        print(f"{'平均每文档字符数:':<20} {avg_chars:>15.1f}")
+        print(f"{'平均每文档 Token 数:':<20} {avg_tokens:>15.1f}")
+        print(f"{'最长文档字符数:':<20} {self.max_length:>15,}")
+        print(f"{'文档长度中位数:':<20} {self.median_length:>15,}")
+        print("=" * 60)
+
+        if self.total_tokens > 0:
+            print()
+            if seq_length is None:
+                print(f"训练样本数: {self.doc_count:,} 个 (一个文档一个样本)")
+            else:
+                print(f"训练样本预估 (seq={seq_length}):")
+                print(f"  可生成约 {self.total_tokens // seq_length:,} 个训练样本")
+
+
+def collect_stats(
+    name: str, loader: Callable[[], tf.data.Dataset], tokenizer: Callable
+) -> DatasetStats:
+    """从 DatasetLoader 收集统计数据
+
+    Args:
+        name: 数据集名称（用于报表显示）
+        loader: 返回 tf.data.Dataset 的加载器函数
+        tokenizer: 分词器函数，接收文本返回 token ID 列表
+
+    Returns:
+        DatasetStats 统计结果对象
+    """
+    ds = loader()
+
+    doc_count = 0
+    total_chars = 0
+    total_tokens = 0
+    lengths = []
+
+    for item in ds:
+        text = item.numpy().decode("utf-8")
+        if not text.strip():
+            continue
+
+        doc_count += 1
+        total_chars += len(text)
+        lengths.append(len(text))
+
+        # Token 统计，过滤掉末尾的 padding (值为 0 的 token)
+        try:
+            import keras
+
+            token_ids = keras.ops.convert_to_numpy(tokenizer(text))
+        except ImportError:
+            # Fallback: assume tokenizer returns numpy array directly
+            token_ids = np.array(tokenizer(text))
+
+        # 只去掉末尾的 0，保留中间内容（包括中间的 OOV/padding）
+        valid_tokens = np.trim_zeros(token_ids, "b")
+        total_tokens += len(valid_tokens)
+
+    return DatasetStats(
+        name=name,
+        doc_count=doc_count,
+        total_chars=total_chars,
+        total_tokens=total_tokens,
+        max_length=max(lengths) if lengths else 0,
+        median_length=int(np.median(lengths)) if lengths else 0,
+    )
+
+
+def save_vocabulary(vocab: list[str], vocab_path: pathlib.Path) -> None:
+    """保存词汇表到文件
+
+    Args:
+        vocab: 词汇表列表
+        vocab_path: 保存路径
+    """
+    vocab_path.parent.mkdir(parents=True, exist_ok=True)
+    with open(vocab_path, "w", encoding="utf-8") as f:
+        for char in vocab:
+            written = char if char != "\n" else r"\n"
+            f.write(written + "\n")
+
+
+def build_vocab_from_dataset(
+    doc_ds: tf.data.Dataset, vocab_path: pathlib.Path
+) -> list[str]:
+    """从文档数据集构建词汇表
+
+    Args:
+        doc_ds: 文档数据集
+        vocab_path: 词汇表保存路径
+
+    Returns:
+        词汇表列表
+    """
+    vectorizer = layers.TextVectorization(
+        output_mode="int", split="character", standardize=None
+    )
+    vectorizer.adapt(doc_ds, batch_size=128)
+
+    vocab = vectorizer.get_vocabulary()
+    if "$" not in vocab:
+        vocab = [*vocab, "$"]
+
+    save_vocabulary(vocab, vocab_path)
+    return vocab

data/dev/mini_c4/file1.txt

@@ -0,0 +1,3 @@
+first document of first file
+second document of first file
+third document of first file

data/dev/mini_c4/file2.txt

@@ -0,0 +1,4 @@
+first document of second file
+second document of second file
+third document of second file
+fourth document of second file

data/dev/mini_c4/file3.txt

@@ -0,0 +1,3 @@
+first document of third file
+second document of third file
+third document of third file

data/dev/poetry/元.csv

@@ -0,0 +1,11 @@
+标题,朝代,作者,体裁,内容
+西洱河,元,述律杰,五言排律,洱水何雄壮，源流自邓川。两关龙首尾，九曲势蜿蜒。大理城池固，金汤铁石坚。四洲从古号，三岛至今传。罗阁凭巘崄，蒙人恃极边。要当兵十万，不数客三千。世祖亲征日，初还一统天。雨师清瘴疠，风伯扫氛烟。民物因蕃富，封疆近百年。点苍山色好，铭刻尚依然。
+陟玩春山纪兴,元,忽必烈,七言律诗,时膺韶景陟兰峰，不惮跻攀谒粹容。花色映霞祥彩混，垆烟拂雾瑞光重。雨沾琼干岩边竹，风袭琴声岭际松。净刹玉毫瞻礼罢，回程仙驾驭苍龙。
+结联,元,奥鲁赤,句,久立危栏须北望，无边秋色杳冥冥。
+八月初四日雪坡太守周门拓入云居山中复度岭饮于水月尼寺赋诗书似太守及苏州刺史周义卿,元,杨维桢,七言律诗,文章太守早休牙，五马传呼处士家。好客新分朱露酒，题诗近在白云窝。山中子落千年桂，海上人归八月槎。水月楼头横玉笛，误猜萼绿是韶华。
+用顾松江韵复理贰守并柬雪坡刺史,元,杨维桢,七言律诗,仙客归来隘九州，身骑黄鹤记南游。乌衣故国江山在，铜柱荒台草木秋。起舞刘琨空有志，登高王粲不胜愁。问君蔗境今何在，祇忆当年顾虎头。
+寄小蓬莱主者闻梅涧并简沈元方宇文仲美贤主宾,元,杨维桢,七言律诗,罗浮主者是仙才，东老诸孙亦俊哉。风雨春城花落尽，江山故国燕归来。酒盟自有乌巾在，笑口应随皓齿开。十八仙人重会处，劫灰不到小蓬莱。
+次韵奉答倪元镇,元,杨维桢,七言律诗,坐断深林事不闻，西窗风日爱余曛。旧经高赤寻三传，新咏山王削五君。翠筱侵床落苍雪，石池洗砚动玄云。东邻书屋最相忆，莫遣草堂移浪文。
+送谢太守,元,杨维桢,七言律诗,朝廷遣使航东海，万里南来送玺书。著屐登山良不恶，分符典郡复何如。白苏事业千年后，吴楚封疆百战馀。今日养民方急务，肯将徵算及舟车。
+送谢太守,元,杨维桢,七言律诗,
+回上张太尉（一云"谢赐玳瑁笔见征楚国公碑文"）,元,杨维桢,七言律诗,昨夜文星照南极，今朝客省过东维。锦囊颖脱千年兔，斑管光摇九尾龟。墨卷风云随王气，恩分雨露出天池。老夫来草平蛮策，先写新封楚国碑。

data/dev/poetry/先秦.csv

@@ -0,0 +1,5 @@
+标题,朝代,作者,体裁,内容
+禹玉牒辞,先秦,无名氏,古风,祝融司方发其英，沐日浴月百宝生。
+衣铭,先秦,无名氏,古风,桑蚕苦，女工难，得新捐故后必寒。
+书车,先秦,无名氏,古风,出畏之，入惧之。
+击壤歌,先秦,无名氏,古风,日出而作。日入而息。凿井而饮。耕田而食。帝力于我何有哉。

data/dev/poetry/南北朝.csv

@@ -0,0 +1,6 @@
+标题,朝代,作者,体裁,内容
+悬瓠方丈竹堂飨侍臣联句诗,南北朝,元宏,古风,白日光天兮无不曜。江左一隅独未照。愿从圣明兮登衡会。万国驰诚混内外。云雷大振兮天门辟。率土来宾一正历。舜舞干戚兮天下归。文德远被莫不思。皇风一鼓兮九地匝。戴日依天清六合。遵彼汝坟兮昔化贞。未若今日道风明。文王政教兮晖江沼。宁如大化光四表。
+歌,南北朝,元宏,句,两菖蒲，新野乐。
+应制赋铜鞮山松诗,南北朝,元协,古风,问松林。松林经几冬。山川何如昔。风云与古同。
+绝命诗二首 其一 ,南北朝,元熙,古风,义实动君子，主辱死忠臣。何以明是节，将解七尺身。
+绝命诗二首 其二 ,南北朝,元熙,古风,平生方寸心，殷勤属知己。从今一销化，悲伤无极已。

data/poetry/__init__.py

@@ -0,0 +1,21 @@
+"""诗歌数据集模块
+
+从以下 github 地址下载数据集到目录 ./data/Poetry：
+
+> https://github.com/xiu-ze/Poetry.git
+
+数据集的格式是多文件 CSV 格式，统计结果：
+
+> 找到 22 个 CSV 文件
+>
+> 诗歌总数: 1014507
+> 最长字符数: 4872
+> 平均字符数: 66.04
+> 中位数: 48
+
+因此可设置序列长度为 100.
+"""
+
+from data.poetry.dataset import PoetryDataset
+
+__all__ = ["PoetryDataset"]

data/poetry/dataset.py

@@ -0,0 +1,104 @@
+"""诗歌数据集主模块
+
+实现 PoetryDataset 类，继承自 DataBundle。
+"""
+
+import pathlib
+from dataclasses import dataclass, field
+from typing import Optional
+
+import tensorflow as tf
+
+from data.base import DataBundle, TokenizerBundle
+from data.poetry.loader import doc_load_with_eot
+from data.poetry.transformer import transform
+from data.poetry.tokenizer import load_vectorizer
+
+
+@dataclass
+class PoetryDataset(DataBundle):
+    """诗歌数据集
+
+    将文档加载、分词、统计等功能绑定在一起的数据集类。
+
+    Usage:
+        dataset = PoetryDataset(
+            data_dir="~/data/Poetry/诗歌数据集",
+            vocab_path="~/data/Poetry/vocabulary.txt",
+            sequence_length=100
+        )
+
+        # 获取文档数据集
+        doc_ds = dataset.doc_ds()
+
+        # 获取 token 数据集
+        tokens_ds = dataset.tokens_ds(seq_length=100, batch_size=128)
+
+        # 打印统计信息
+        dataset.stat(seq_length=100)
+    """
+
+    vocab_path: str = ""
+
+    _data_path: pathlib.Path = field(init=False, repr=False)
+    _vocab_path: pathlib.Path = field(init=False, repr=False)
+    _tokenizer_info: Optional[TokenizerBundle] = field(
+        init=False, repr=False, default=None
+    )
+
+    def __post_init__(self):
+        self._data_path = pathlib.Path(self.data_dir).expanduser()
+        self._vocab_path = pathlib.Path(self.vocab_path).expanduser()
+
+    def _load_tokenizer(self):
+        """懒加载分词器"""
+        if self._tokenizer_info is None:
+            tokenizer = load_vectorizer(self._vocab_path, self.sequence_length + 1)
+            vocab = tokenizer.get_vocabulary()
+            end_of_text = vocab.index("$")
+            vocab_size = len(vocab)
+
+            def decode(token_ids: list[int]) -> str:
+                chars = [
+                    vocab[token_id] for token_id in token_ids if token_id < len(vocab)
+                ]
+                return "".join(chars)
+
+            self._tokenizer_info = TokenizerBundle(
+                tokenizer=tokenizer,
+                decode=decode,
+                end_of_text=end_of_text,
+                vocab_size=vocab_size,
+                vocab_path=str(self._vocab_path)
+            )
+
+    def doc_ds(self) -> tf.data.Dataset:
+        """返回原始文档数据集
+
+        Returns:
+            TensorFlow Dataset，每个元素是带结束标记的诗歌内容
+        """
+        return doc_load_with_eot(self._data_path)
+
+    def tokens_ds(self, seq_length: int, batch_size: int) -> tf.data.Dataset:
+        """返回 tokenized 数据集
+
+        Args:
+            seq_length: 序列长度（诗歌中此参数主要用于兼容性）
+            batch_size: 批次大小
+
+        Returns:
+            TensorFlow Dataset，每个元素是 (input_ids, target_ids) 对
+        """
+        self._load_tokenizer()
+        ds = self.doc_ds()
+        return transform(
+            ds=ds,
+            tokenizer=self._tokenizer_info.tokenizer,
+            batch_size=batch_size,
+        )
+
+    def tokenizer_bundle(self) -> TokenizerBundle:
+        """返回分词器信息"""
+        self._load_tokenizer()
+        return self._tokenizer_info

data/poetry/loader.py

@@ -0,0 +1,59 @@
+"""诗歌数据集文档加载模块
+
+从 CSV 文件加载诗歌文本数据。
+"""
+
+import glob
+import os
+import pathlib
+
+import tensorflow as tf
+
+
+def _parse_csv_line(line: tf.Tensor) -> tf.Tensor:
+    """解析 CSV 行，返回内容列"""
+    fields = tf.io.decode_csv(
+        line,
+        use_quote_delim=False,  # 行内的引号是普通字符
+        record_defaults=["", "", "", "", ""],
+    )
+    return fields[4]  # 返回 '内容' 列的值
+
+
+def doc_load(data_dir: pathlib.Path) -> tf.data.Dataset:
+    """加载诗歌数据集
+
+    从指定目录下的 CSV 文件中加载诗歌文本数据。
+    每个 CSV 文件应该包含以下列：标题、作者、朝代、类型、内容。
+
+    Args:
+        data_dir: 数据目录路径
+
+    Returns:
+        TensorFlow Dataset，每个元素是诗歌内容字符串
+    """
+    csv_files = glob.glob(os.path.join(data_dir, "*.csv"))
+    if not csv_files:
+        raise ValueError(f"在目录 {data_dir} 中未找到任何 CSV 文件！")
+
+    files_ds = tf.data.Dataset.from_tensor_slices(csv_files)
+    csv_line_ds = files_ds.interleave(
+        lambda csv_file: tf.data.TextLineDataset(csv_file).skip(1),
+        cycle_length=1,
+    )
+    return csv_line_ds.map(_parse_csv_line, num_parallel_calls=tf.data.AUTOTUNE).filter(
+        lambda x: tf.strings.length(x) > 0
+    )
+
+
+def doc_load_with_eot(data_dir: pathlib.Path) -> tf.data.Dataset:
+    """加载诗歌数据集，每行末尾添加结束标记
+
+    Args:
+        data_dir: 数据目录路径
+
+    Returns:
+        TensorFlow Dataset，每个元素是带结束标记的诗歌内容
+    """
+    ds = doc_load(data_dir)
+    return ds.map(lambda x: tf.strings.join([x, "$"]))

data/poetry/runner.py

@@ -0,0 +1,38 @@
+"""诗歌数据集 Runner
+
+Usage:
+    python data/poetry/runner.py build_vocab
+    python data/poetry/runner.py test_dataset
+    ENV=production python data/poetry/runner.py build_vocab
+"""
+
+import pathlib
+import sys
+
+sys.path.insert(0, str(pathlib.Path(__file__).parent.parent.parent))
+
+from data.poetry.dataset import PoetryDataset
+from data.runner import DatasetRunner
+from env.resolve import resolve_path, resolve_env, resolve_saved
+
+
+dataset = PoetryDataset(
+    data_dir=str(
+        resolve_env(resolve_path("data/dev/poetry"), resolve_path("~/data/Poetry/诗歌数据集"))
+    ),
+    vocab_path=str(
+        resolve_env(
+            resolve_saved("vocab/poetry/vocab.txt"),
+            resolve_path("~/data/Poetry/vocabulary.txt"),
+        )
+    ),
+    sequence_length=100,
+)
+
+runner = DatasetRunner(
+    dataset=dataset,
+    name="poetry",
+)
+
+if __name__ == "__main__":
+    runner()

data/poetry/tokenizer.py

@@ -0,0 +1,67 @@
+"""诗歌数据集分词器模块
+
+提供诗歌数据集专用的分词器实现。
+"""
+
+import pathlib
+
+from keras import layers
+
+
+def load_vocabulary(vocab_path: pathlib.Path):
+    """从文本文件加载词汇表，每行一个字符。
+
+    Args:
+        vocab_path: 词汇表文件路径
+
+    Returns:
+        词汇表列表
+    """
+
+    def extract_word(line: str) -> str:
+        word = line[:-1]  # 去掉行末的换行符
+        return word if word != r"\n" else "\n"
+
+    with open(vocab_path, "r", encoding="utf-8") as f:
+        vocab = [extract_word(line) for line in f]
+    return vocab
+
+
+def load_vectorizer(
+    vocab_path: pathlib.Path, sequence_length: int = 101
+) -> layers.TextVectorization:
+    """从词汇表文件加载分词器
+
+    Args:
+        vocab_path: 词汇表文件路径
+        sequence_length: 输出序列长度，默认为 101
+                       （多一位是为了在训练时构建输入和目标偏移一位）
+
+    Returns:
+        TextVectorization 层
+    """
+    vectorizer = layers.TextVectorization(
+        output_mode="int",
+        split="character",
+        output_sequence_length=sequence_length,
+        standardize=None,
+    )
+
+    vocab = load_vocabulary(vocab_path)
+    vectorizer.set_vocabulary(vocab)
+
+    return vectorizer
+
+
+def create_vectorizer(sequence_length: int = 101) -> layers.TextVectorization:
+    """创建新的分词器（用于训练词汇表）
+
+    Args:
+        sequence_length: 输出序列长度，默认为 101
+
+    Returns:
+        TextVectorization 层
+    """
+    return layers.TextVectorization(
+        output_mode="int", split="character", standardize=None
+    )

data/poetry/transformer.py

@@ -0,0 +1,38 @@
+"""诗歌数据集 token 转换模块
+
+将诗歌文档数据集转换为训练用的 token 序列。
+"""
+
+from typing import Callable
+
+import tensorflow as tf
+
+
+def transform(
+    ds: tf.data.Dataset,
+    tokenizer: Callable,
+    batch_size: int,
+) -> tf.data.Dataset:
+    """转换诗歌数据集为训练数据集
+
+    诗歌数据集已经生成了固定数量的 token 序列，不足的部分会 padding。
+
+    Args:
+        ds: 文档数据集
+        tokenizer: 分词器函数
+        batch_size: 批次大小
+
+    Returns:
+        训练数据集，每个元素是 (input_ids, target_ids) 对
+    """
+    # 文本向量化；对于诗歌数据集来说，已经生成了固定数量的 token 序列了，不足的部分会 padding
+    ds = ds.map(tokenizer, num_parallel_calls=8)
+
+    # 构建输入和目标（偏移一位）
+    # 无需在这里添加结束标记，因为在 doc_load 中已经添加了结束标记
+    ds = ds.map(lambda x: (x[:-1], x[1:]))
+
+    # 重新设置批次大小并预取数据以提高性能
+    ds = ds.batch(batch_size).prefetch(8)
+
+    return ds

data/runner.py

@@ -0,0 +1,142 @@
+"""数据集 Runner 公共模块
+
+提供通用的数据集测试和词汇表生成功能。
+
+Usage:
+    # 在各自 runner.py 中实例化
+    from data.runner import DatasetRunner
+    from data.poetry.dataset import PoetryDataset
+    from env.resolve import resolve, resolve_saved, resolve_env
+
+    dataset = PoetryDataset(
+        data_dir=str(resolve_env(resolve("data/dev/poetry"), resolve("~/data/Poetry/诗歌数据集"))),
+        vocab_path=str(resolve_env(resolve_saved("poetry/vocab.txt"), resolve("~/data/Poetry/vocabulary.txt"))),
+        sequence_length=100,
+    )
+    runner = DatasetRunner(dataset=dataset, name="poetry")
+    runner()
+"""
+
+from data.base import DataBundle
+from data.common import build_vocab_from_dataset
+from env.resolve import resolve_saved
+from env.runner import ActionRunner
+
+
+class DatasetRunner(ActionRunner):
+    """数据集 Runner
+
+    提供通用的数据集测试和词汇表生成功能。
+
+    Args:
+        dataset: 数据集实例（PoetryDataset 或 WikiDataset）
+        name: 数据集英文名称（如 "poetry", "wiki"）
+        max_docs: 测试时显示的文档数量，默认 5
+        max_samples: 测试时显示的 token 样本数量，默认 3
+        max_doc_chars: 文档显示的最大字符数，默认 200
+        max_text_display: token 文本显示的最大字符数，默认 80
+
+    Usage:
+        runner = DatasetRunner(dataset=poetry_dataset, name="poetry")
+        runner.test_dataset()  # 或 runner.build_vocab()
+    """
+
+    # 中英文名称映射
+    NAME_MAP = {
+        "poetry": "诗歌",
+        "wiki": "Wiki",
+    }
+
+    def __init__(
+        self,
+        dataset: DataBundle,
+        name: str,
+        max_docs: int = 5,
+        max_samples: int = 3,
+        max_doc_chars: int = 200,
+        max_text_display: int = 80,
+    ):
+        self.dataset = dataset
+        self.name = name
+        self.display_name = self.NAME_MAP.get(name, name)
+        self.vocab_path = resolve_saved(f"vocab/{name}/vocab.txt")
+        self.max_docs = max_docs
+        self.max_samples = max_samples
+        self.max_doc_chars = max_doc_chars
+        self.max_text_display = max_text_display
+
+    def build_vocab(self) -> None:
+        """生成字符词汇表"""
+        print(f"正在加载数据集...")
+        ds = self.dataset.doc_ds()
+
+        print(f"正在保存词汇表到: {self.vocab_path}")
+        vocab = build_vocab_from_dataset(ds, self.vocab_path)
+
+        print(f"词汇表大小: {len(vocab)}")
+        print("完成！")
+
+    def test_dataset(self) -> None:
+        """测试数据集"""
+        print("\n" + "=" * 60)
+        print(f"{self.display_name} 数据集测试")
+        print("=" * 60)
+
+        self._view_documents(self.dataset.doc_ds())
+        self._view_tokens(self.dataset)
+        self._show_vocab_info(self.dataset.tokenizer_bundle())
+
+        print("\n" + "=" * 60)
+        print("测试完成")
+        print("=" * 60)
+
+    def _view_documents(self, doc_ds) -> None:
+        """查看原始文档"""
+        print("\n【原始文档查看】")
+        print("-" * 60)
+        count = 0
+        for doc in doc_ds.take(self.max_docs):
+            count += 1
+            text = doc.numpy().decode("utf-8")
+            if len(text) > self.max_doc_chars:
+                text = text[: self.max_doc_chars] + "..."
+            print(f"\n第 {count} 个文档:")
+            print(f"  {text}")
+        print(f"\n共显示 {count} 个文档")
+
+    def _view_tokens(self, dataset) -> None:
+        """查看 tokenized 数据"""
+        print("\n【Tokenized 数据查看】")
+        print("-" * 60)
+
+        tokenizer_info = dataset.tokenizer_bundle()
+        tokens_ds = dataset.tokens_ds(seq_length=dataset.sequence_length, batch_size=1)
+
+        count = 0
+        for batch_input, batch_target in tokens_ds.take(self.max_samples):
+            count += 1
+            input_ids = batch_input[0].numpy()
+            target_ids = batch_target[0].numpy()
+
+            input_text = tokenizer_info.decode(input_ids.tolist())
+            target_text = tokenizer_info.decode(target_ids.tolist())
+
+            if len(input_text) > self.max_text_display:
+                input_text = input_text[: self.max_text_display] + "..."
+            if len(target_text) > self.max_text_display:
+                target_text = target_text[: self.max_text_display] + "..."
+
+            print(f"\n第 {count} 个样本:")
+            print(f"  输入 tokens: {input_ids[:20]}... (长度: {len(input_ids)})")
+            print(f"  目标 tokens: {target_ids[:20]}... (长度: {len(target_ids)})")
+            print(f"  输入文本: {input_text}")
+            print(f"  目标文本: {target_text}")
+        print(f"\n共显示 {count} 个样本")
+
+    @staticmethod
+    def _show_vocab_info(tokenizer_info) -> None:
+        """显示词汇表信息"""
+        print("\n【词汇表信息】")
+        print("-" * 60)
+        print(f"  词汇表大小: {tokenizer_info.vocab_size}")
+        print(f"  结束标记 ID: {tokenizer_info.end_of_text}")

data/tokenizers.py

@@ -0,0 +1,89 @@
+"""
+GPT模型的共享组件模块：
+
+- 分词器
+"""
+
+import keras
+import keras_hub
+from keras import layers
+
+
+def sentence_piece():
+    # 用预训练好的分词器，也就是说我们不去自己训练分词器了
+    vocabulary_file = keras.utils.get_file(
+        origin="https://hf-mirror.com/mattdangerw/spiece/resolve/main/vocabulary.proto"
+    )
+    # [Note] 依然需要 tensorflow_text 包
+    tokenizer = keras_hub.tokenizers.SentencePieceTokenizer(vocabulary_file)
+
+    end_of_text = tokenizer.token_to_id("<|endoftext|>")
+
+    def decode(tokens: list[int]) -> str:
+        return tokenizer.detokenize(tokens)
+
+    return tokenizer, end_of_text, decode
+
+
+def character_vectorization():
+    """简单的字符级分词器，适用于测试"""
+    vectorizer = layers.TextVectorization(output_mode="int", split="character")
+    vectorizer.set_vocabulary(
+        list("abcdefghijklmnopqrstuvwxyz0123456789 .,!?;:()[]{}<>-_\n")
+        + ["<|endoftext|>"]  # 兼容 sentence_piece 分词器的特殊标记
+    )
+
+    vocab = vectorizer.get_vocabulary()
+    for idx, word in enumerate(vocab):
+        if word == "<|endoftext|>":
+            end_of_text = idx
+            break
+    else:
+        raise ValueError("Vocabulary does not contain <|endoftext|> token.")
+
+    def decode(tokens: list[int]) -> str:
+        words = [vocab[token] for token in tokens]
+        return "".join(words)
+
+    return vectorizer, end_of_text, decode
+
+
+def poetry_character_vectorization(
+    vocab_path: str = "local/saved/vocab/poetry/vocab.txt",
+):
+    """从文本文件加载诗歌字符级分词器。
+
+    词汇表文件格式：每行一个字符，第一行必须是 <|endoftext|>。
+
+    Args:
+        vocab_path: 词汇表文件路径，默认为 "local/saved/poetry/vocab.txt"
+
+    Returns:
+        (vectorizer, end_of_text, decode): 分词器、结束标记ID、解码函数
+    """
+    from env.resolve import resolve_path
+
+    # 读取词汇表
+    vocab_file = resolve_path(vocab_path)
+    with open(vocab_file, "r", encoding="utf-8") as f:
+        vocab = [line.rstrip("\n") for line in f]
+
+    # 创建 TextVectorization 层
+    vectorizer = layers.TextVectorization(
+        output_mode="int", split="character", standardize=None
+    )
+    vectorizer.set_vocabulary(vocab)
+
+    # 找到 end_of_text 的索引
+    for idx, word in enumerate(vocab):
+        if word == "<|endoftext|>":
+            end_of_text = idx
+            break
+    else:
+        raise ValueError("Vocabulary does not contain <|endoftext|> token.")
+
+    def decode(tokens: list[int]) -> str:
+        words = [vocab[token] for token in tokens]
+        return "".join(words)
+
+    return vectorizer, end_of_text, decode

data/wiki/__init__.py

@@ -0,0 +1,8 @@
+"""Wiki 数据集模块
+
+导出 WikiDataset 类。
+"""
+
+from data.wiki.dataset import WikiDataset
+
+__all__ = ["WikiDataset"]

data/wiki/dataset.py

@@ -0,0 +1,100 @@
+"""Wiki 数据集主模块
+
+实现 WikiDataset 类，继承自 DataBundle。
+"""
+
+import pathlib
+from dataclasses import dataclass, field
+from typing import Optional
+
+import tensorflow as tf
+
+from data.base import DataBundle, TokenizerBundle
+from data.wiki.loader import doc_load
+from data.wiki.transformer import transform
+from data.wiki.tokenizer import sentence_piece, character_vectorization
+
+
+@dataclass
+class WikiDataset(DataBundle):
+    """Wiki 数据集
+
+    将文档加载、分词、统计等功能绑定在一起的数据集类。
+
+    Usage:
+        dataset = WikiDataset(
+            data_dir="~/data/wiki/mini_c4",
+            tokenizer_type="sentence_piece"  # 或 "character"
+        )
+
+        # 获取文档数据集
+        doc_ds = dataset.doc_ds()
+
+        # 获取 token 数据集
+        tokens_ds = dataset.tokens_ds(seq_length=256, batch_size=32)
+
+        # 打印统计信息
+        dataset.stat(seq_length=256)
+    """
+
+    glob_pattern: str = "*"
+    tokenizer_type: str = "sentence_piece"
+
+    _data_path: pathlib.Path = field(init=False, repr=False)
+    _tokenizer_bundle: Optional[TokenizerBundle] = field(
+        init=False, repr=False, default=None
+    )
+
+    def __post_init__(self):
+        self._data_path = pathlib.Path(self.data_dir).expanduser()
+
+    def _load_tokenizer(self):
+        """懒加载分词器"""
+        if self._tokenizer_bundle is None:
+            if self.tokenizer_type == "sentence_piece":
+                tokenizer, end_of_text, decode = sentence_piece()
+            elif self.tokenizer_type == "character":
+                tokenizer, end_of_text, decode = character_vectorization()
+            else:
+                raise ValueError(f"Unknown tokenizer type: {self.tokenizer_type}")
+
+            vocab_size = tokenizer.vocabulary_size()
+            self._tokenizer_bundle = TokenizerBundle(
+                tokenizer=tokenizer,
+                decode=decode,
+                end_of_text=end_of_text,
+                vocab_size=vocab_size
+            )
+
+    def doc_ds(self) -> tf.data.Dataset:
+        """返回原始文档数据集
+
+        Returns:
+            TensorFlow Dataset，每个元素是一个文档字符串
+        """
+        return doc_load(self._data_path, glob_pattern=self.glob_pattern)
+
+    def tokens_ds(self, seq_length: int, batch_size: int) -> tf.data.Dataset:
+        """返回 tokenized 数据集
+
+        Args:
+            seq_length: 序列长度
+            batch_size: 批次大小
+
+        Returns:
+            TensorFlow Dataset，每个元素是 (input_ids, target_ids) 对
+        """
+        self._load_tokenizer()
+        ds = self.doc_ds()
+        return transform(
+            ds=ds,
+            tokenizer=self._tokenizer_bundle.tokenizer,
+            end_of_text=self._tokenizer_bundle.end_of_text,
+            sequence_length=seq_length,
+            batch_size=batch_size,
+        )
+
+    def tokenizer_bundle(self) -> TokenizerBundle:
+        """返回分词器信息"""
+        self._load_tokenizer()
+        return self._tokenizer_bundle

data/wiki/loader.py

@@ -0,0 +1,52 @@
+"""Wiki 数据集文档加载模块
+
+从 mini_c4 格式加载文档数据集。
+"""
+
+import pathlib
+
+import tensorflow as tf
+
+
+def doc_load(
+    data_dir: pathlib.Path, glob_pattern: str = "*", cycle_length: int = 32
+) -> tf.data.Dataset:
+    """加载并处理文档数据集为 TensorFlow Dataset。
+
+    递归查找指定目录下匹配 glob_pattern 的所有文件，使用 doc_extract 函数
+    将每个文件转换为 TensorFlow Dataset，然后使用 interleave 进行并行处理。
+
+    目录下的文件格式要求每行一个文档，其中的换行符使用 "\\n" 转义。
+
+    Args:
+        data_dir: 数据目录路径
+        glob_pattern: 文件匹配模式，如 "*.txt"，默认为 "*" 匹配所有文件
+        cycle_length: interleave 的 cycle_length 参数，控制并行处理的文件数量，默认为 32
+
+    Returns:
+        合并后的 TensorFlow Dataset，包含所有文件处理后的数据
+    """
+    # 获取所有文件（过滤掉目录），递归查找子目录
+    files = [str(file) for file in data_dir.rglob(glob_pattern) if file.is_file()]
+    if not files:
+        raise FileNotFoundError(f"在目录 {data_dir} 中未找到匹配 {glob_pattern} 的文件")
+
+    # 排序文件列表以确保一致的处理顺序
+    files = sorted(files)
+
+    # 创建数据集管道
+    ds = tf.data.Dataset.from_tensor_slices(files)
+    ds = ds.interleave(
+        _line_doc_extract,
+        cycle_length=cycle_length,
+        num_parallel_calls=tf.data.AUTOTUNE,
+    )
+
+    return ds
+
+
+def _line_doc_extract(path: str) -> tf.data.Dataset:
+    """Mini-c4 format: one document per line."""
+    return tf.data.TextLineDataset(path).map(
+        lambda x: tf.strings.regex_replace(x, r"\\n", "\n")
+    )

data/wiki/runner.py

@@ -0,0 +1,32 @@
+"""Wiki 数据集 Runner
+
+Usage:
+    python data/wiki/runner.py test_dataset
+    ENV=production python data/wiki/runner.py test_dataset
+"""
+
+import pathlib
+import sys
+
+sys.path.insert(0, str(pathlib.Path(__file__).parent.parent.parent))
+
+from data.runner import DatasetRunner
+from data.wiki.dataset import WikiDataset
+from env.resolve import resolve_path, resolve_env
+
+
+dataset = WikiDataset(
+    data_dir=str(
+        resolve_env(resolve_path("data/dev/mini_c4"), resolve_path("~/data/wiki/mini_c4"))
+    ),
+    tokenizer_type=resolve_env("character", "sentence_piece"),
+    sequence_length=256,
+)
+
+runner = DatasetRunner(
+    dataset=dataset,
+    name="wiki",
+)
+
+if __name__ == "__main__":
+    runner()

data/wiki/tokenizer.py

@@ -0,0 +1,60 @@
+"""Wiki 数据集分词器模块
+
+提供 Wiki 数据集专用的分词器实现。
+"""
+
+import keras
+import keras_hub
+from keras import layers
+
+
+def sentence_piece():
+    """SentencePiece 分词器
+
+    使用预训练好的分词器，无需自己训练。
+
+    Returns:
+        (tokenizer, end_of_text, decode): 分词器、结束标记ID、解码函数
+    """
+    # 用预训练好的分词器，也就是说我们不去自己训练分词器了
+    vocabulary_file = keras.utils.get_file(
+        origin="https://hf-mirror.com/mattdangerw/spiece/resolve/main/vocabulary.proto"
+    )
+    # [Note] 依然需要 tensorflow_text 包
+    tokenizer = keras_hub.tokenizers.SentencePieceTokenizer(vocabulary_file)
+
+    end_of_text = tokenizer.token_to_id("<|endoftext|>")
+
+    def decode(tokens: list[int]) -> str:
+        return tokenizer.detokenize(tokens)
+
+    return tokenizer, end_of_text, decode
+
+
+def character_vectorization():
+    """字符级分词器
+
+    简单的字符级分词器，适用于测试。
+
+    Returns:
+        (tokenizer, end_of_text, decode): 分词器、结束标记ID、解码函数
+    """
+    vectorizer = layers.TextVectorization(output_mode="int", split="character")
+    vectorizer.set_vocabulary(
+        list("abcdefghijklmnopqrstuvwxyz0123456789 .,!?;:()[]{}\u003c\u003e-_\n")
+        + ["<|endoftext|>"]  # 兼容 sentence_piece 分词器的特殊标记
+    )
+
+    vocab = vectorizer.get_vocabulary()
+    for idx, word in enumerate(vocab):
+        if word == "<|endoftext|>":
+            end_of_text = idx
+            break
+    else:
+        raise ValueError("Vocabulary does not contain <|endoftext|> token.")
+
+    def decode(tokens: list[int]) -> str:
+        words = [vocab[token] for token in tokens]
+        return "".join(words)
+
+    return vectorizer, end_of_text, decode

data/wiki/transformer.py

@@ -0,0 +1,47 @@
+"""Wiki 数据集 token 转换模块
+
+将文档数据集转换为训练用的 token 序列。
+"""
+
+from typing import Callable
+
+import numpy as np
+import tensorflow as tf
+
+
+def transform(
+    ds: tf.data.Dataset,
+    tokenizer: Callable,
+    end_of_text: int,
+    sequence_length: int,
+    batch_size: int,
+) -> tf.data.Dataset:
+    """转换文档数据集为训练数据集
+
+    将文档转换为 token ID，添加结束标记，分割为固定长度的序列。
+
+    Args:
+        ds: 文档数据集
+        tokenizer: 分词器函数
+        end_of_text: 结束标记的 token ID
+        sequence_length: 序列长度
+        batch_size: 批次大小
+
+    Returns:
+        训练数据集，每个元素是 (input_ids, target_ids) 对
+    """
+    ds = ds.map(tokenizer, num_parallel_calls=8)
+
+    # 将文档之间添加 end_of_text 标记分隔
+    ds = ds.map(lambda x: tf.concat([x, np.array([end_of_text])], -1))
+
+    # 重新设置样本大小为固定长度序列
+    ds = ds.rebatch(sequence_length + 1, drop_remainder=True)
+
+    # 构建输入和目标（偏移一位）
+    ds = ds.map(lambda x: (x[:-1], x[1:]))
+
+    # 重新设置批次大小并预取数据以提高性能
+    ds = ds.batch(batch_size).prefetch(8)
+
+    return ds

data/wiki/wiki_cleaner.py

@@ -0,0 +1,122 @@
+"""
+Wiki 文本清洗模块。
+
+提供多种过滤器用于清洗 wiki 格式的文本数据。
+"""
+
+import re
+
+
+def filter_single_line(text: str) -> str | None:
+    """
+    过滤只有一行的数据（通常是重定向页面）。
+
+    Args:
+        text: 输入文本
+
+    Returns:
+        如果只有一行返回 None，否则返回原文本
+    """
+    lines = [line for line in text.split("\n") if line.strip()]
+    if len(lines) <= 1:
+        return None
+    return text
+
+
+def filter_empty_brackets(text: str) -> str:
+    """
+    移除文本中的空括号对。
+
+    例如：（）、()、（ ）、( )、[ ]、【 】、{ } 等
+
+    Args:
+        text: 输入文本
+
+    Returns:
+        移除空括号后的文本
+    """
+    # 匹配空括号对：() （） [] 【】 {} 等，中间可有空白
+    pattern = re.compile(r"[\(\)（）\[\]【】{}]\s*[\(\)（）\[\]【】{}]")
+    return pattern.sub("", text)
+
+
+def filter_html_tags(text: str) -> str:
+    """
+    移除 HTML/XML 标签（HTML 实体编码格式）。
+
+    例如：&lt;templatestyles src="ShareCSS/infobox.css" /&gt;
+
+    Args:
+        text: 输入文本
+
+    Returns:
+        移除 HTML 标签后的文本
+    """
+    # 匹配 &lt;...&gt; 格式的实体编码标签
+    pattern = re.compile(r"&lt;[^&]+&gt;")
+    return pattern.sub("", text)
+
+
+def filter_lang_tags(text: str) -> str:
+    """
+    移除特殊的语言标记（支持嵌套）。
+
+    例如：-{H|zh-hans:重定向;zh-hant:重新导向;}-
+    嵌套例如：-{T|zh:-{zh|}-;zh-hans:-{zh-hans|}-;}-
+
+    Args:
+        text: 输入文本
+
+    Returns:
+        移除语言转换标记后的文本
+    """
+    # 使用非贪婪匹配，循环处理嵌套
+    pattern = re.compile(r"-\{[^{}]+?}-")
+    while True:
+        new_text = pattern.sub("", text)
+        if new_text == text:  # 没有更多匹配了
+            break
+        text = new_text
+    return text
+
+
+def clean(text: str) -> str | None:
+    """
+    应用所有过滤器清洗文本。
+
+    过滤顺序：
+    1. 单行检查（重定向页面）
+    2. HTML 标签
+    3. 空白括号行
+    4. 语言转换标记
+    5. 最终空检查
+
+    Args:
+        text: 输入文本
+
+    Returns:
+        清洗后的文本，如果应该丢弃则返回 None
+    """
+    # 1. 检查单行
+    result = filter_single_line(text)
+    if result is None:
+        return None
+
+    # 2. 移除 HTML 标签
+    result = filter_html_tags(result)
+
+    # 3. 移除空白括号行
+    result = filter_empty_brackets(result)
+
+    # 4. 移除语言转换标记
+    result = filter_lang_tags(result)
+
+    # 5. 多个连续空行替换为一个空行
+    result = re.sub(r"\n\s*\n", "\n\n", result)
+    result = result.strip()
+
+    # 6. 最终检查：如果结果为空或只剩空白，返回 None
+    if not result.strip():
+        return None
+
+    return result

docs/TODOs.md

@@ -0,0 +1,3 @@
+- [ ] `<doc>` 格式由于计算图的限制还无法实现，未来打算实现。
+- [ ] 希望能通过 Callback 或 train_step 截取到训练过程中的数据。
+- [ ] wiki 训练后不能回答事实性问题，感觉是过拟合了，将 dropout 调成 0.5 试一试（当前 0.1）。

docs/pycharm.md

@@ -0,0 +1,10 @@
+# PyCharm 开发指南
+
+最近，我在项目里尝试将 PyCharm 的代码连接到远程服务器运行，遇到了一些莫名的问题。现在将一些解决方案记录下来，供以后参考。
+
+总的来说，我直接应用远程环境就会出错，有各种各样的问题。我需要重新构建一个全新的环境才使得它正常运作。记录如下：
+
+1. 在远程服务器上创建一个新的 conda 环境。
+2. 创建一个新的 Python 项目（我直接移动了我的项目目录，并删除目录下的 .idea, .ruff_cache, .pytest_cache 等文件夹）。
+3. 在本地 PyCharm 中配置远程 Python 解释器，指向远程服务器上的新环境。这一步骤中，注意配置好目录映射，和不自动上传文件。
+4. 等待一段时间，就能正常运作了。

env/keras.py

@@ -0,0 +1,11 @@
+"""Keras 相关工具模块
+
+提供 Keras 配置相关的功能。
+"""
+
+import keras
+
+
+def enable_mixed_precision():
+    """开启混合精度训练/推理"""
+    keras.config.set_dtype_policy("mixed_float16")

env/logger.py

@@ -0,0 +1,52 @@
+import logging
+from functools import wraps
+
+
+def get_logger(name: str, filepath: str = None):
+    logger = logging.getLogger(name)
+    logger.setLevel(logging.INFO)
+
+    # 控制台
+    console_handler = logging.StreamHandler()
+    logger.addHandler(console_handler)
+
+    # 文件
+    if filepath:
+        file_handler = logging.FileHandler(filepath)
+        logger.addHandler(file_handler)
+
+    return logger
+
+
+def log(enter_message: str = "", exit_message: str = ""):
+    return _Log(enter_message=enter_message, exit_message=exit_message)
+
+
+class _Log:
+    def __init__(
+        self,
+        enter_message: str = "",
+        exit_message: str = ""
+    ):
+        self.enter_message = enter_message
+        self.exit_message = exit_message
+
+    def __enter__(self):
+        if self.enter_message:
+            print(self.enter_message)
+        return self
+
+    def __exit__(self, exc_type, exc, tb):
+        if self.exit_message:
+            print(self.exit_message)
+        print("")
+        return False
+
+    def __call__(self, func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            with _Log(self.enter_message, self.exit_message):
+                return func(*args, **kwargs)
+            return None
+
+        return wrapper

env/resolve.py

@@ -0,0 +1,85 @@
+import os
+from enum import Enum, StrEnum
+from pathlib import Path
+
+
+"""定义项目的根路径"""
+PROJECT_ROOT = Path(__file__).parent.parent.absolute()
+
+
+"""定义根据环境变量选择配置的函数"""
+class Env(StrEnum):
+    TEST = "test"
+    PRODUCTION = "production"
+
+def resolve_env[T](test_conf: T = Env.TEST, prod_conf: T = Env.PRODUCTION) -> T:
+    env = os.environ.get("ENV", str(Env.TEST))
+    return prod_conf if env == str(Env.PRODUCTION) else test_conf
+
+
+"""定义一些预设的目录"""
+SAVED_DIR = resolve_env(
+    PROJECT_ROOT / "local" / "saved",
+    PROJECT_ROOT / "saved",
+)
+TASKS_DIR = PROJECT_ROOT / "local" / "tasks"
+
+
+"""定义一些路径解析函数，方便在项目中使用"""
+def resolve_saved(path: str | Path = None) -> Path:
+    """解析相对于 saved 目录的路径
+
+    1. 如果本身就是 Path 对象，直接返回。
+    2. 如果 path 是 None，返回 saved 目录本身。
+    3. 否则，将 path 解析为相对于 saved 目录的路径。
+    """
+    if isinstance(path, Path):
+        return path
+    return SAVED_DIR / path if path else SAVED_DIR
+
+
+def resolve_task_dir(task_name: str) -> Path:
+    """解析任务所在的目录
+
+    Args：
+        task_name: 任务名称，即定义在 Pipeline 中的 name 字段，例如 "poetry_gpt" 或 "poetry_rnn"。
+    """
+    return TASKS_DIR / task_name
+
+
+def resolve_path(path: str | Path) -> Path:
+    """从项目根目录解析路径
+
+    1. 如果路径是 Path 对象，直接返回。
+    2. 如果路径是以 ~ 或 / 开头的绝对路径，则直接返回该路径。
+    3. 如果路径是相对路径，则将其解析为相对于项目根目录的路径。
+
+    Args:
+        path: 相对于项目根目录的路径
+
+    Returns:
+        解析后的绝对路径
+
+    Example:
+        >>> resolve_path("data/dev/mini_c4/file.txt")
+        PosixPath('/Users/.../universal_deeplearning/data/dev/mini_c4/file.txt')
+    """
+    if isinstance(path, Path):
+        return path
+    elif path.startswith("~") or path.startswith("/"):
+        return Path(path).expanduser().resolve()
+    else:
+        return PROJECT_ROOT / path
+
+
+def display_path(path: str | Path) -> str:
+    """将路径转换为适合展示的字符串
+
+    如果路径位于项目根目录内，则显示为相对项目根目录的路径；
+    否则显示绝对路径。
+    """
+    resolved = resolve_path(path)
+    try:
+        return str(resolved.relative_to(PROJECT_ROOT))
+    except ValueError:
+        return str(resolved)

env/runner.py

@@ -0,0 +1,23 @@
+import sys
+from typing import Callable
+
+
+class ActionRunner:
+    def __call__(self, default_method: str | Callable = None):
+        if len(sys.argv) > 1:
+            method = self._resolve_method(sys.argv[1])
+        else:
+            method = default_method
+            if type(method) == str:
+                method = self._resolve_method(method)
+
+        if method:
+            method()
+        else:
+            raise ValueError("没有指定要执行的方法")
+
+    def _resolve_method(self, method_name: str) -> Callable:
+        method = getattr(self, method_name, None)
+        if method is None:
+            raise ValueError(f"没有找到对应的方法：{method_name}")
+        return method

env/vocab.py

@@ -0,0 +1,2 @@
+# 定义所有词典中 PAD 的 id token
+PAD = 0

environment-linux.yml

@@ -0,0 +1,15 @@
+name: general-dl
+channels:
+  - defaults
+dependencies:
+  - python=3.12
+  - pip
+  - numpy
+  - tensorflow
+  - tensorflow-text
+  - keras
+  - pip:
+    - keras-hub
+    - gradio
+variables:
+  ENV: production

environment.yml

@@ -0,0 +1,17 @@
+name: general-dl
+channels:
+  - defaults
+dependencies:
+  - python=3.12
+  - pip
+  - setuptools>=68,<70
+  - numpy
+  - ruff
+  - pytest
+  - pytest-mock
+  - tensorflow
+  - keras
+  - pip:
+      - gradio
+variables:
+  ENV: test

generate_requirements.py

@@ -0,0 +1,110 @@
+#!/usr/bin/env python3
+"""
+从 environment-linux.yml 生成 requirements.txt
+YAML 中的版本号优先级最高
+未指定版本号时查询当前环境的实际版本
+排除 python 和 pip
+"""
+
+import os
+from datetime import datetime
+import yaml
+from importlib.metadata import version, PackageNotFoundError
+
+
+# 排除的包（不加入 requirements.txt）
+EXCLUDE_PACKAGES = {"python", "pip"}
+
+
+def get_installed_version(package_name):
+    """获取包的安装版本，未安装返回 None"""
+    try:
+        return version(package_name)
+    except PackageNotFoundError:
+        return None
+
+
+def parse_package_string(dep):
+    """
+    解析包字符串，返回 (包名, yaml版本号或None)
+    例如: "tensorflow=2.15.0" -> ("tensorflow", "2.15.0")
+          "numpy" -> ("numpy", None)
+    """
+    if "=" in dep:
+        parts = dep.split("=")
+        pkg_name = parts[0]
+        pkg_version = parts[1]
+        return pkg_name, pkg_version
+    else:
+        return dep, None
+
+
+def parse_environment_yml(filepath):
+    """解析 environment-linux.yml，提取包列表和版本信息"""
+    with open(filepath, "r") as f:
+        env = yaml.safe_load(f)
+
+    packages = []
+
+    for dep in env.get("dependencies", []):
+        if isinstance(dep, str):
+            # 简单字符串格式：package 或 package=version
+            pkg_name, yaml_version = parse_package_string(dep)
+            if pkg_name not in EXCLUDE_PACKAGES:
+                packages.append((pkg_name, yaml_version))
+        elif isinstance(dep, dict) and "pip" in dep:
+            # pip 子列表
+            for pip_dep in dep["pip"]:
+                pkg_name, yaml_version = parse_package_string(pip_dep)
+                if pkg_name not in EXCLUDE_PACKAGES:
+                    packages.append((pkg_name, yaml_version))
+
+    return packages
+
+
+def main():
+    yml_file = "environment-linux.yml"
+    output_file = "requirements.txt"
+
+    print(f"读取 {yml_file}...")
+    packages = parse_environment_yml(yml_file)
+    print(f"发现 {len(packages)} 个包（排除 {EXCLUDE_PACKAGES}）")
+
+    lines = []
+    for pkg_name, yaml_version in packages:
+        if yaml_version:
+            # YAML 中有版本号，优先使用
+            lines.append(f"{pkg_name}=={yaml_version}")
+            print(f"  ✓ {pkg_name}=={yaml_version} (来自 YAML)")
+        else:
+            # YAML 中没有版本号，查询当前环境
+            env_version = get_installed_version(pkg_name)
+            if env_version:
+                lines.append(f"{pkg_name}=={env_version}")
+                print(f"  ✓ {pkg_name}=={env_version} (来自当前环境)")
+            else:
+                lines.append(pkg_name)
+                print(f"  ⚠ {pkg_name} (未安装，无版本号)")
+
+    # 添加头部注释
+    header_lines = [
+        f"# Generated from {yml_file}",
+        f"# Timestamp: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}",
+        f"# Environment: {os.environ.get('ENV', 'unknown')}",
+        "#",
+    ]
+
+    # 合并所有行
+    all_lines = header_lines + lines
+
+    with open(output_file, "w") as f:
+        f.write("\n".join(all_lines) + "\n")
+
+    print(f"\n已生成 {output_file}：")
+    print("-" * 40)
+    print("\n".join(all_lines))
+    print("-" * 40)
+
+
+if __name__ == "__main__":
+    main()

models/mini_gpt/__init__.py

@@ -0,0 +1 @@
+from models.mini_gpt.model_builder import GptModelBuilder

models/mini_gpt/gpt_components.py

@@ -0,0 +1,61 @@
+"""
+GPT模型的共享组件模块：
+
+- Positional Encoding
+- Transformer Decoder
+"""
+
+import keras
+from keras import layers, ops
+
+class PositionalEmbedding(keras.Layer):
+    def __init__(self, sequence_length, input_dim, output_dim, **kwargs):
+        super().__init__(**kwargs)
+        self.token_embeddings = layers.Embedding(input_dim, output_dim)
+        self.position_embeddings = layers.Embedding(sequence_length, output_dim)
+
+    def call(self, inputs, reverse=False):
+        if reverse:
+            token_embeddings = self.token_embeddings.embeddings
+            return ops.matmul(inputs, ops.transpose(token_embeddings))
+        positions = ops.cumsum(ops.ones_like(inputs), axis=-1) - 1
+        embedded_tokens = self.token_embeddings(inputs)
+        embedded_positions = self.position_embeddings(positions)
+        return embedded_tokens + embedded_positions
+
+
+class TransformerDecoder(keras.Layer):
+    def __init__(self, hidden_dim, intermediate_dim, num_heads, **kwargs):
+        super().__init__(**kwargs)
+
+        self.hidden_dim = hidden_dim
+        self.intermediate_dim = intermediate_dim
+
+        key_dim = hidden_dim // num_heads
+
+        # self-attention 层
+        self.self_attention = layers.MultiHeadAttention(num_heads, key_dim, dropout=0.1)
+        self.self_attention_layernorm = layers.LayerNormalization()
+
+        # feed-forward 层
+        self.feed_forward_1 = layers.Dense(intermediate_dim, activation="relu")
+        self.feed_forward_2 = layers.Dense(hidden_dim)
+        self.feed_forward_layernorm = layers.LayerNormalization()
+        self.dropout = layers.Dropout(0.1)
+
+    def call(self, inputs):
+        # self-attention 计算
+        residual = x = inputs
+        x = self.self_attention(query=x, key=x, value=x, use_causal_mask=True)
+        x = self.dropout(x)
+        x = x + residual
+        x = self.self_attention_layernorm(x)
+
+        # feed-forward 计算
+        residual = x
+        x = self.feed_forward_1(x)
+        x = self.feed_forward_2(x)
+        x = self.dropout(x)
+        x = x + residual
+        x = self.feed_forward_layernorm(x)
+        return x

models/mini_gpt/model_builder.py

@@ -0,0 +1,54 @@
+from dataclasses import dataclass
+from functools import partial
+
+import keras
+from keras import layers
+
+from models.mini_gpt.gpt_components import PositionalEmbedding, TransformerDecoder
+from pipeline.base.generation import generate_with_training_model
+from pipeline.base.model_builder import ModelArtifact
+
+
+@dataclass
+class GptModelBuilder:
+    hidden_dim: int
+    intermediate_dim: int
+    num_heads: int
+    num_layers: int
+
+    def build_training_artifact(
+        self,
+        vocab_size: int,
+        sequence_length: int
+    ) -> ModelArtifact:
+        inputs = keras.Input(shape=(None,), dtype="int32", name="inputs")
+        embedding = PositionalEmbedding(
+            sequence_length,
+            vocab_size,
+            self.hidden_dim,
+            name="embedding"
+        )
+        x = embedding(inputs)
+        x = layers.LayerNormalization(name="input_layer_norm")(x)
+
+        for i in range(self.num_layers):
+            decoder = TransformerDecoder(
+                self.hidden_dim,
+                self.intermediate_dim,
+                self.num_heads,
+                name=f"decoder_{i}"
+            )
+            x = decoder(x)
+
+        outputs = embedding(x, reverse=True)
+        model = keras.Model(inputs, outputs, name="mini_gpt")
+        return ModelArtifact(
+            model=model,
+            generate=partial(generate_with_training_model, model)
+        )
+
+    def build_inference_artifact(
+        self,
+        training_artifact: ModelArtifact
+    ) -> ModelArtifact:
+        return training_artifact

models/rnn/__init__.py

@@ -0,0 +1 @@
+from models.rnn.model_builder import RNNModelBuilder

models/rnn/model_builder.py

@@ -0,0 +1,114 @@
+from dataclasses import dataclass
+from functools import partial
+
+import keras
+import tensorflow as tf
+from keras import layers
+
+from pipeline.base.generation import generate_with_stateful_model, generate_with_training_model
+from pipeline.base.model_builder import ModelArtifact
+
+
+@dataclass
+class RNNModelBuilder:
+    num_layers: int = 2
+    embedding_dim: int = 100
+    hidden_dim: int = 1024
+
+    def build_training_artifact(
+        self,
+        vocab_size: int,
+        sequence_length: int
+    ) -> ModelArtifact:
+        inputs = keras.Input(shape=(None,), dtype="int32", name="inputs")
+        x = layers.Embedding(
+            input_dim=vocab_size,
+            output_dim=self.embedding_dim,
+            mask_zero=True,
+            name="embedding"
+        )(inputs)
+
+        for i in range(self.num_layers):
+            x = layers.LSTM(
+                self.hidden_dim,
+                return_sequences=True,
+                recurrent_dropout=0.1,
+                name=f"lstm_{i}"
+            )(x)
+            x = layers.Dropout(0.1, name=f"dropout_{i}")(x)
+
+        outputs = layers.Dense(vocab_size, name="logits")(x)
+        model = keras.Model(inputs=inputs, outputs=outputs, name="rnn_training")
+        return ModelArtifact(
+            model=model,
+            generate=partial(generate_with_training_model, model)
+        )
+
+    def build_inference_artifact(
+        self,
+        training_artifact: ModelArtifact
+    ) -> ModelArtifact:
+        inference_model = self._build_inference_model_from_training_model(
+            training_artifact.model
+        )
+        return ModelArtifact(
+            model=inference_model,
+            generate=partial(
+                generate_with_stateful_model,
+                inference_model,
+                initial_states=self._initial_states(batch_size=1)
+            )
+        )
+
+    def _build_inference_model_from_training_model(
+        self,
+        training_model: keras.Model
+    ) -> keras.Model:
+        token_input = keras.Input(shape=(None,), dtype="int32", name="token_input")
+        state_inputs = []
+        for i in range(self.num_layers):
+            h_input = keras.Input(shape=(self.hidden_dim,), name=f"h_{i}_input")
+            c_input = keras.Input(shape=(self.hidden_dim,), name=f"c_{i}_input")
+            state_inputs.extend([h_input, c_input])
+
+        embedding = training_model.get_layer("embedding")
+        logits_layer = training_model.get_layer("logits")
+        x = embedding(token_input)
+
+        new_states = []
+        inference_lstm_layers = []
+        for i in range(self.num_layers):
+            inference_lstm = layers.LSTM(
+                self.hidden_dim,
+                return_sequences=i < self.num_layers - 1,
+                return_state=True,
+                recurrent_dropout=0.1,
+                name=f"lstm_{i}"
+            )
+            h_input = state_inputs[i * 2]
+            c_input = state_inputs[i * 2 + 1]
+            x, new_h, new_c = inference_lstm(x, initial_state=[h_input, c_input])
+            new_states.extend([new_h, new_c])
+            dropout = training_model.get_layer(f"dropout_{i}")
+            x = dropout(x)
+            inference_lstm_layers.append(inference_lstm)
+
+        logits = logits_layer(x)
+        inference_model = keras.Model(
+            [token_input] + state_inputs,
+            [logits] + new_states,
+            name="rnn_inference"
+        )
+
+        for i, inference_lstm in enumerate(inference_lstm_layers):
+            training_lstm = training_model.get_layer(f"lstm_{i}")
+            inference_lstm.set_weights(training_lstm.get_weights())
+
+        return inference_model
+
+    def _initial_states(self, batch_size: int) -> list:
+        states = []
+        for _ in range(self.num_layers):
+            states.append(tf.zeros((batch_size, self.hidden_dim)))
+            states.append(tf.zeros((batch_size, self.hidden_dim)))
+        return states

pipeline/__init__.py

@@ -0,0 +1,3 @@
+from .runner import PipelineRunner
+from .pipeline import Pipeline
+from .base.configs import CheckpointConfig

pipeline/base/checkpoint.py

@@ -0,0 +1,147 @@
+"""
+模型工具模块
+
+包含模型构建、检查点管理等通用功能。
+"""
+
+import pathlib
+import re
+import warnings
+
+from env.resolve import resolve_path
+
+
+def extract_number_of_filename(filename: str) -> int:
+    """
+    从文件名中提取数字，无论数字出现在文件名的哪个位置。
+
+    例如：
+    - "model_epoch_001.weights.h5" -> 1
+    - "checkpoint_2024_06_30_epoch_002.weights.h5" -> 2
+    - "model_epoch_final.weights.h5" -> 抛出异常
+
+    :param filename: 包含数字的文件名字符串
+    :return: 提取的数字，如果没有数字则返回0
+    """
+    numbers = re.findall(r"\d+", filename)
+    if numbers:
+        return int(numbers[-1])  # 返回最后一个数字，假设它是代数
+    else:
+        raise ValueError(f"No number found in filename: {filename}")
+
+
+def resolve_checkpoint(
+    dirs: list[pathlib.Path | str] | None = None,
+    path: pathlib.Path | str | None = None,
+    epoch: int | None = None,
+    suffix: str | None = None
+):
+    """统一解析模型检查点路径
+
+    支持直接指定检查点文件路径或在目录中查找检查点文件。
+
+    参数:
+        dirs: 检查点目录列表
+        path: 直接指定的检查点文件路径（支持绝对路径和相对路径）
+        epoch: 指定的 epoch，用于查找对应的 .weights.h5 文件
+        suffix: 指定检查点文件后缀
+
+    返回:
+        (resolved_path, epoch): 绝对路径和 epoch 数
+
+    抛出:
+        FileNotFoundError: 当指定的路径不存在或未找到检查点文件时
+        ValueError: 当参数无效时
+    """
+    resolved_dirs = _resolve_checkpoint_dirs(dirs)
+
+    if path is not None:
+        path = pathlib.Path(path)
+
+        if not path.is_absolute():
+            if not resolved_dirs:
+                raise ValueError("path 是相对路径时，必须提供 dirs")
+            path = _resolve_relative_checkpoint_path(path, resolved_dirs)
+        else:
+            if dirs is not None:
+                warnings.warn(
+                    "警告：path 是绝对路径，dirs 参数将被忽略",
+                    UserWarning
+                )
+
+        if not path.exists():
+            raise FileNotFoundError(f"检查点文件不存在: {path}")
+        if suffix is not None and not path.name.endswith(suffix):
+            raise FileNotFoundError(f"检查点文件后缀不匹配: {path}")
+
+        try:
+            epoch_num = extract_number_of_filename(path.stem)
+        except ValueError:
+            epoch_num = 0
+
+        return path, epoch_num
+
+    if not resolved_dirs:
+        raise ValueError("必须提供 dirs 或 path")
+
+    files_with_number = _collect_checkpoint_files(
+        checkpoint_dirs=resolved_dirs,
+        suffix=suffix
+    )
+
+    if epoch is not None:
+        matches = [(f, num) for f, num in files_with_number if num == epoch]
+        if not matches:
+            raise FileNotFoundError(f"未找到 epoch {epoch} 对应的检查点文件")
+        if len(matches) > 1:
+            raise RuntimeError(
+                f"找到多个 epoch {epoch} 对应的检查点文件: {[match[0].name for match in matches]}"
+            )
+        return matches[0]
+
+    if not files_with_number:
+        return None, 0
+
+    return max(files_with_number, key=lambda item: item[1])
+
+
+def _resolve_checkpoint_dirs(
+    dirs: list[pathlib.Path | str] | None
+) -> list[pathlib.Path]:
+    if dirs is None:
+        return []
+    return [resolve_path(path) for path in dirs]
+
+
+def _resolve_relative_checkpoint_path(
+    checkpoint_path: pathlib.Path,
+    checkpoint_dirs: list[pathlib.Path]
+) -> pathlib.Path:
+    for checkpoint_dir in checkpoint_dirs:
+        candidate = checkpoint_dir / checkpoint_path
+        if candidate.exists():
+            return candidate
+    return checkpoint_dirs[0] / checkpoint_path
+
+
+def _collect_checkpoint_files(
+    checkpoint_dirs: list[pathlib.Path],
+    suffix: str | None
+) -> list[tuple[pathlib.Path, int]]:
+    files_with_number = []
+    for checkpoint_dir in checkpoint_dirs:
+        if not checkpoint_dir.exists():
+            continue
+        for file_path in sorted(checkpoint_dir.iterdir()):
+            if not file_path.is_file():
+                continue
+            if suffix is not None and not file_path.name.endswith(suffix):
+                continue
+            if suffix is None and not _is_checkpoint_file(file_path):
+                continue
+            files_with_number.append((file_path, extract_number_of_filename(file_path.stem)))
+    return files_with_number
+
+
+def _is_checkpoint_file(file_path: pathlib.Path) -> bool:
+    return file_path.name.endswith(".keras") or file_path.name.endswith(".weights.h5")

pipeline/base/configs.py

@@ -0,0 +1,69 @@
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Callable
+
+
+@dataclass
+class CheckpointConfig:
+    dirs: list[Path] | None = None
+    path: Path = None
+    epoch: int = None
+    suffix: str = None
+
+
+@dataclass
+class ModelConfig:
+    sequence_length: int = 256
+    hidden_dim: int = 512
+    intermediate_dim: int = 2056
+    num_heads: int = 8
+    num_layers: int = 8
+
+
+@dataclass
+class TrainingRule:
+    batch_size: int = 128
+    epochs: int = 1
+    steps_per_epoch: int = 30
+    validation_batches: int = 1
+
+
+@dataclass
+class GenerationRule:
+    prompts_generator: Callable
+    sample_strategy: Callable
+
+
+@dataclass
+class CheckpointRules:
+    training: CheckpointConfig = field(default_factory=CheckpointConfig)
+    testing: CheckpointConfig = field(default_factory=CheckpointConfig)
+    deployment: CheckpointConfig = field(default_factory=CheckpointConfig)
+
+    def resolve_training_rule(
+        self,
+        default_dirs: list[Path | str] | None = None
+    ) -> dict:
+        return self._resolve_rule(self.training, default_dirs)
+
+    def resolve_testing_rule(
+        self,
+        default_dirs: list[Path | str] | None = None
+    ) -> dict:
+        return self._resolve_rule(self.testing, default_dirs)
+
+    def resolve_deployment_rule(
+        self,
+        default_dirs: list[Path | str] | None = None
+    ) -> dict:
+        return self._resolve_rule(self.deployment, default_dirs)
+
+    @staticmethod
+    def _resolve_rule(checkpoint: CheckpointConfig, default_dirs: list[Path | str] | None) -> dict:
+        dirs = checkpoint.dirs if checkpoint.dirs is not None else default_dirs
+        return {
+            "dirs": dirs,
+            "path": checkpoint.path,
+            "epoch": checkpoint.epoch,
+            "suffix": checkpoint.suffix
+        }

pipeline/base/generation.py

@@ -0,0 +1,174 @@
+"""
+与生成有关的组件
+"""
+
+import pathlib
+from dataclasses import dataclass
+from typing import Any, Callable
+
+import keras
+import numpy as np
+from keras import callbacks, ops
+
+from env.vocab import PAD
+from env.logger import get_logger
+from pipeline.base.model_builder import GenerationContext, GenerationResult, ModelArtifact
+
+
+def generate_with_training_model(
+    model: keras.Model,
+    context: GenerationContext,
+    prompt_tokens: list[int]
+) -> GenerationResult:
+    prompt_length = len(prompt_tokens)
+
+    if prompt_length == 0:
+        return GenerationResult([], "<|empty|>")
+
+    tokens = prompt_tokens + [PAD] * (context.max_length - prompt_length)
+
+    for i in range(prompt_length, context.max_length):
+        prediction = model.predict(np.array([tokens]), verbose=0)
+        prediction = prediction[0, i - 1]
+        next_token = ops.convert_to_numpy(context.sample_fn(prediction))
+        next_token_id = np.array(next_token).item()
+        tokens[i] = next_token_id
+
+        if next_token_id == context.end_of_text:
+            return GenerationResult(tokens[:i], "<|endoftext|>")
+        if next_token_id == PAD:
+            return GenerationResult(tokens[:i], "<|pad|>")
+
+    return GenerationResult(tokens, "<|maxlength|>")
+
+
+def generate_with_stateful_model(
+    model: keras.Model,
+    context: GenerationContext,
+    prompt_tokens: list[int],
+    initial_states: list
+) -> GenerationResult:
+    if not prompt_tokens:
+        return GenerationResult([], "<|empty|>")
+
+    tokens = list(prompt_tokens)
+    batch_tokens = np.array([tokens])
+    logits, *states = model.predict([batch_tokens] + initial_states, verbose=0)
+
+    for _ in range(len(tokens), context.max_length):
+        next_token = ops.convert_to_numpy(context.sample_fn(logits[0]))
+        next_token_id = np.array(next_token).item()
+        tokens.append(next_token_id)
+
+        if next_token_id == context.end_of_text:
+            return GenerationResult(tokens[:-1], "<|endoftext|>")
+        if next_token_id <= PAD:
+            return GenerationResult(tokens, "<|pad|>")
+
+        logits, *states = model.predict([np.array([[next_token_id]])] + states, verbose=0)
+
+    return GenerationResult(tokens, "<|maxlength|>")
+
+
+@dataclass
+class TextGenerationResult:
+    text: str
+    stop_reason: str
+
+
+class TextGenerator:
+    def __init__(
+        self,
+        artifact: ModelArtifact,
+        tokenizer: Any,
+        decode: Callable,
+        end_of_text: int,
+        sample_fn: Callable,
+        max_length: int
+    ):
+        self.artifact = artifact
+        self.tokenizer = tokenizer
+        self.decode = decode
+        self.context = GenerationContext(
+            end_of_text=end_of_text,
+            max_length=max_length,
+            sample_fn=sample_fn
+        )
+
+    def generate_tokens(
+        self,
+        prompt: str,
+        max_length: int | None = None,
+        sample_fn: Callable | None = None
+    ) -> GenerationResult:
+        context = GenerationContext(
+            end_of_text=self.context.end_of_text,
+            max_length=max_length if max_length is not None else self.context.max_length,
+            sample_fn=sample_fn if sample_fn is not None else self.context.sample_fn
+        )
+        prompt_tokens = self._tokenize_prompt(prompt)
+        return self.artifact.generate(context, prompt_tokens)
+
+    def generate_text(
+        self,
+        prompt: str,
+        max_length: int | None = None,
+        sample_fn: Callable | None = None
+    ) -> TextGenerationResult:
+        result = self.generate_tokens(prompt, max_length, sample_fn)
+        return TextGenerationResult(
+            text=self.decode(result.token_ids),
+            stop_reason=result.stop_reason
+        )
+
+    def _tokenize_prompt(self, prompt: str) -> list[int]:
+        prompt_tokens = list(ops.convert_to_numpy(self.tokenizer(prompt)))
+        return [token for token in prompt_tokens if token > PAD]
+
+
+class GenerationCallback(callbacks.Callback):
+    def __init__(
+        self,
+        prompts: list[str],
+        log_file: pathlib.Path,
+        tokenizer: Any,
+        decode: Callable,
+        end_of_text: int,
+        max_length: int,
+        sample_fn: Callable,
+        training_artifact: ModelArtifact
+    ):
+        super().__init__()
+        self.prompts = prompts
+        self.tokenizer = tokenizer
+        self.decode = decode
+        self.end_of_text = end_of_text
+        self.max_length = max_length
+        self.sample_fn = sample_fn
+        self.training_artifact = training_artifact
+        self.logger = self.init_logger(log_file)
+
+    def on_epoch_end(self, epoch, logs=None):
+        generator = TextGenerator(
+            artifact=self.training_artifact,
+            tokenizer=self.tokenizer,
+            decode=self.decode,
+            end_of_text=self.end_of_text,
+            max_length=self.max_length,
+            sample_fn=self.sample_fn
+        )
+        self.logger.info(f"\nGenerated text after epoch {epoch + 1}:")
+        for i, prompt in enumerate(self.prompts):
+            result = generator.generate_text(prompt)
+            self.logger.info(f"Prompt {i + 1:2}: {prompt}")
+            self.logger.info(f"Generated: {result.text}{result.stop_reason}\n")
+
+    @staticmethod
+    def init_logger(log_file: pathlib.Path):
+        if not log_file.parent.exists():
+            log_file.parent.mkdir(parents=True)
+
+        logger = get_logger("GenerationCallback", filepath=str(log_file))
+        logger.info("Initialized GenerationCallback logger")
+
+        return logger