HydroRoll 水系社区基础预训练模型

基于 MiniRBT (hfl/minirbt-h256) 的中文 TRPG（桌上角色扮演游戏）日志命名实体识别系统，支持训练、推理、ONNX 导出和 WebUI 可视化标注平台。

功能特性

• NER 实体识别: 自动识别 TRPG 日志中的发言者、时间戳、对话、动作、注释等实体
• 灵活训练: 支持初次训练、增量训练、标签维度扩展/缩减
• ONNX 导出: 支持 ONNX 格式导出，实现 CPU 高速推理
• WebUI 平台: Label Studio 风格的可视化标注界面 (WIP)
• 数据转换: 支持 Word-level/Char-level CoNLL 格式互转
• 自动修复: 智能修复截断的时间戳和发言者名称

实体类型

标签	说明	示例
speaker	发言者	"风雨"
timestamp	时间戳	"2024-06-08 21:44:59"
dialogue	对话内容	""呜哇...""
action	动作描述	"剧烈的疼痛从头颅深处一波波地涌出..."
comment	注释/旁白	"（红木家具上刻着一行小字）"

安装

环境要求

• Python >= 3.12
• CUDA (可选，用于 GPU 加速)

安装依赖

# 使用 uv (推荐)
uv sync

# 或使用 pip
pip install -r requirements.txt

1. 数据准备

将训练数据准备为 CoNLL 格式，支持两种格式：

Char-level 格式 (推荐)

-DOCSTART- -X- O
风 -X- _ B-speaker
雨 -X- _ I-speaker
  -X- _ O
2 -X- _ B-timestamp
0 -X- _ I-timestamp
...

Word-level 格式

风雨 O O B-speaker
2024-06-08 O O B-timestamp
21:44:59 O O I-timestamp

2. 初次训练

# 基础训练（使用默认参数）
uv run main.py --train --conll ./data

# 完整参数训练
uv run main.py --train \
    --conll ./data \
    --model hfl/minirbt-h256 \
    --output ./models/trpg-final \
    --epochs 20 \
    --batch 4

参数说明

参数	说明	默认值
--train	启用训练模式	-
--conll	CoNLL 文件或目录路径	./data
--model	基础模型名称	hfl/minirbt-h256
--output	模型输出目录	./models/trpg-ner-v1
--epochs	训练轮数	20
--batch	批处理大小	4
--resume	恢复检查点路径	None

3. 推理测试

# 单文本测试
uv run main.py --test "风雨 2024-06-08 21:44:59 剧烈的疼痛从头颅深处一波波地涌出..."

# 多文本测试
uv run main.py --test \
    "莎莎 2024-06-08 21:46:26 \"呜哇...\" 下意识去拿法杖" \
    "BOT 2024-06-08 21:50:03 莎莎 的出目是 D10+7=6+7=13"

输出格式

{
  "metadata": {
    "speaker": "风雨",
    "timestamp": "2024-06-08 21:44:59"
  },
  "content": [
    {
      "type": "comment",
      "content": "剧烈的疼痛从头颅深处一波波地涌出...",
      "confidence": 0.952
    }
  ]
}

增量训练

在已有模型基础上继续训练新数据：

# 继续训练（自动加载最新检查点）
uv run main.py --train \
    --conll ./new_data \
    --output ./models/trpg-final \
    --epochs 5

# 从指定检查点恢复
uv run main.py --train \
    --conll ./new_data \
    --output ./models/trpg-final \
    --resume ./models/trpg-final/checkpoint-200 \
    --epochs 5

标签维度管理

添加新标签训练

在 CoNLL 数据中添加新的标签类型（如 emotion），模型会自动适配：

# 添加新标签后重新训练（使用 ignore_mismatched_sizes）
uv run main.py --train \
    --conll ./data_with_emotion \
    --output ./models/trpg-final \
    --epochs 15

减少标签维度训练

如果需要减少标签类型，建议从头训练：

# 从基础模型重新训练（使用缩减后的标签集）
uv run main.py --train \
    --conll ./data_reduced_labels \
    --model hfl/minirbt-h256 \
    --output ./models/trpg-reduced \
    --epochs 20

ONNX 导出

将训练好的模型导出为 ONNX 格式，用于 CPU 推理加速。

注意: ONNX 导出功能会自动使用推理时找到的模型目录。如果模型目录不在默认位置，请用 --output 指定正确的模型目录。

# 方式一：从默认模型目录导出（自动查找 models/trpg-final）
uv run main.py --export_onnx \
    --onnx_path ./models/trpg-final/model.onnx

# 方式二：指定模型目录导出
uv run main.py --export_onnx \
    --output ./models/trpg-final \
    --onnx_path ./models/trpg-final/model.onnx

# 方式三：导出到其他路径
uv run main.py --export_onnx \
    --onnx_path ./models/trpg-optimized.onnx

ONNX 模型特点

• 支持 CPU 推理（无需 GPU）
• 模型大小约 50-100 MB
• 推理速度约 10-50 ms/句（取决于硬件）
• 兼容 Windows/Linux/macOS/Raspberry Pi

ONNX 模型测试

基础推理测试

# 使用 ONNX 模型进行推理
uv run tests/onnx_infer.py "风雨 2024-06-08 21:44:59 剧烈的疼痛..."

性能指标

Performance Results (n=100):
   Average latency: 25.32 ms
   P95 latency:     31.45 ms
   Max RAM usage:   120.5 MB
   Throughput:      39.5 sentences/sec

数据转换工具

CoNLL 转 Dataset

# 转换单个文件
uv run src/utils/conll_to_dataset.py data.conll \
    --output ./dataset/trpg \
    --format jsonl \
    --validate

# 转换整个目录
uv run src/utils/conll_to_dataset.py ./data \
    --output ./dataset/trpg \
    --format both \
    --validate

输出格式

• --format jsonl: 输出 JSONL 格式
• --format dataset: 输出 HuggingFace Dataset 格式
• --format both: 同时输出两种格式

Word-level 转 Char-level CoNLL

uv run src/utils/word_conll_to_char_conll.py \
    input_word.conll \
    output_char.conll

系统要求

训练环境

• CPU: Intel Core i5 或同等性能
• RAM: 8 GB+
• GPU: NVIDIA GPU（可选，用于加速）
• 存储: 5 GB+

ONNX 推理环境

• CPU: Intel Core i3 或同等性能
• RAM: 2 GB+
• 存储: 100 MB
• 支持 Raspberry Pi 4 (4GB RAM)

开源协议

本项目采用 AFL-3.0 协议开源。

相关链接

• MiniRBT 模型
• Transformers 文档
• ONNX Runtime