TOLE模型完整启动方法指南

TOLE (Token-level Optimization with Language Models) 是一种基于强化学习的可控文本生成方法，通过token级别的反馈实现对文本多个属性的精确控制。以下是完整的启动方法指南：

1. 环境准备

1.1 创建虚拟环境

conda create -n tole_rl python=3.9
conda activate tole_rl

1.2 安装依赖

# 基础依赖
pip install torch==2.0.0 transformers==4.30.2 datasets==2.14.4 rouge-score nltk# 强化学习依赖
pip install gymnasium==0.28.1 stable-baselines3# 其他工具
pip install numpy pandas tqdm tensorboard

2. 数据准备

2.1 数据集格式

确保数据集包含以下字段：

text: 原始文本
sentiment: 情感标签 (如positive/negative)
topic: 主题标签 (如politics/entertainment)

2.2 示例数据集结构

data/
├── train.jsonl
├── dev.jsonl
└── test.jsonl

3. 模型准备

3.1 预训练语言模型

下载并缓存预训练模型（如gpt2-medium）：

python -c "from transformers import GPT2LMHeadModel, GPT2Tokenizer; \
model = GPT2LMHeadModel.from_pretrained('gpt2-medium'); \
tokenizer = GPT2Tokenizer.from_pretrained('gpt2-medium')"

3.2 准备评分器(checkpoint)

确保已有训练好的情感分类器和主题分类器：

models/
├── sentiment_scorer/    # 情感评分器checkpoint
└── topic_scorer/        # 主题评分器checkpoint

4. 训练权重器(Weigher)

权重器用于平衡不同属性评分器的重要性：

python weigher.py \--sent_scorer_path models/sentiment_scorer \--topic_scorer_path models/topic_scorer \--train_data_path data/train.jsonl \--eval_data_path data/dev.jsonl \--output_dir models/weigher \--learning_rate 5e-5 \--batch_size 32 \--num_epochs 10

参数说明：

sent_scorer_path: 情感评分器路径
topic_scorer_path: 主题评分器路径
output_dir: 权重器保存路径

5. 运行Token-level RL训练

使用训练好的权重器和评分器进行策略模型训练：

python token_main.py \--sent_reward_model models/sentiment_scorer \--topic_reward_model models/topic_scorer \--weigher_ckpt models/weigher/final_checkpoint \--train_data_path data/train.jsonl \--eval_data_path data/dev.jsonl \--output_dir models/policy_model \--learning_rate 1e-5 \--batch_size 8 \--num_epochs 5 \--max_length 128 \--gamma 0.99 \--kl_coef 0.2

参数说明：

sent_reward_model: 情感奖励模型路径
topic_reward_model: 主题奖励模型路径
weigher_ckpt: 权重器检查点路径
gamma: 奖励折扣因子
kl_coef: KL散度惩罚系数

6. 模型推理与评估

6.1 生成文本

python generate.py \--model_path models/policy_model/final_checkpoint \--input_text "Once upon a time" \--sentiment positive \--topic entertainment \--output_file generated_texts.txt

6.2 评估模型

python evaluate.py \--model_path models/policy_model/final_checkpoint \--eval_data_path data/test.jsonl \--metrics_file metrics.json

7. 常见问题与解决方案

CUDA内存不足
- 降低batch_size
- 使用--gradient_accumulation_steps 4
训练不稳定
- 调整kl_coef（建议范围：0.1-0.5）
- 降低learning_rate
环境依赖冲突
- 使用pip freeze > requirements.txt保存当前环境
- 使用Docker容器化部署

8. 参考资料

论文链接：Reinforcement Learning with Token-level Feedback for Controllable Text Generation (NAACL 2024)
代码仓库：https://github.com/hust-nlp/TOLE
联系邮箱：wendili@hust.edu.cn

如果遇到任何问题，请通过邮箱联系作者获取支持。以下是基于强化学习的可控文本生成方法的概述，主要介绍TOLE模型外的代表性工作及其核心思想：