update

examples/playground/chat_modelscope.py

@@ -0,0 +1,144 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+"""
+https://github.com/jingyaogong/minimind/blob/master/eval_llm.py
+"""
+import argparse
+import os
+from pathlib import Path
+import platform
+import time
+
+if platform.system() in ("Windows", "Darwin"):
+    from project_settings import project_path, temp_directory
+else:
+    project_path = os.path.abspath("../../")
+    project_path = Path(project_path)
+    temp_directory = Path("/root/autodl-tmp/OpenMiniMind/temp")
+
+import torch
+from modelscope import AutoTokenizer, AutoModelForCausalLM
+from transformers import TextStreamer
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--pretrained_model_name_or_path",
+        default="qgyd2021/Qwen2.5-0.5B-ultrachat-sft-deepspeed",
+        type=str
+    )
+    parser.add_argument(
+        "--model_cache_dir",
+        default=(temp_directory / "hub_models").as_posix(),
+        type=str
+    )
+    parser.add_argument(
+        "--max_new_tokens",
+        default=8192, # 8192, 128
+        type=int, help="最大生成长度（注意：并非模型实际长文本能力）"
+    )
+    parser.add_argument("--top_p", default=0.85, type=float, help="nucleus采样阈值（0-1）")
+    parser.add_argument("--temperature", default=0.85, type=float, help="生成温度，控制随机性（0-1，越大越随机）")
+
+    parser.add_argument(
+        "--show_speed",
+        default=1,  # 1, 0
+        type=int, help="显示decode速度（tokens/s）"
+    )
+
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = get_args()
+
+    os.environ["MODELSCOPE_CACHE"] = args.model_cache_dir
+
+    if torch.cuda.is_available():
+        device = "cuda"
+    elif torch.backends.mps.is_available():
+        # device = "mps"
+        device = "cpu"
+    else:
+        device = "cpu"
+    print(f"device: {device}")
+
+    model = AutoModelForCausalLM.from_pretrained(
+        args.pretrained_model_name_or_path,
+        cache_dir=args.model_cache_dir,
+        trust_remote_code=True,
+    )
+    tokenizer = AutoTokenizer.from_pretrained(
+        args.pretrained_model_name_or_path,
+        cache_dir=args.model_cache_dir,
+        trust_remote_code=True,
+    )
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+        tokenizer.pad_token_id = tokenizer.eos_token_id
+
+    model = model.eval().to(device)
+    # print(tokenizer)
+    # print(model)
+
+    prompts = [
+        "你有什么特长？",
+        "为什么天空是蓝色的",
+        "请用Python写一个计算斐波那契数列的函数",
+        '解释一下"光合作用"的基本过程',
+        "如果明天下雨，我应该如何出门",
+        "比较一下猫和狗作为宠物的优缺点",
+        "解释什么是机器学习",
+        "推荐一些中国的美食"
+    ]
+    input_mode = int(input("[0] 自动测试\n[1] 手动输入\n"))
+
+    streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+
+    # conversation = list()
+    conversation = [
+        {"role": "system", "content": "You are a helpful assistant"}
+    ]
+    while True:
+        if input_mode == 0:
+            if len(prompts) == 0:
+                break
+            user_input = prompts.pop(0)
+            print(f"💬: {user_input}")
+        else:
+            user_input = input("💬: ")
+            user_input = str(user_input).strip()
+        conversation.append({"role": "user", "content": user_input})
+        inputs = tokenizer.apply_chat_template(
+            conversation=conversation,
+            tokenize=False,
+            add_generation_prompt=True
+        )
+        inputs = tokenizer.__call__(
+            inputs,
+            return_tensors="pt",
+            truncation=True
+        )
+        inputs = inputs.to(device)
+        # print(inputs)
+
+        print("🤖: ", end="")
+        st = time.time()
+        generated_ids = model.generate(
+            inputs=inputs["input_ids"], attention_mask=inputs["attention_mask"],
+            max_new_tokens=args.max_new_tokens, do_sample=True, streamer=streamer,
+            pad_token_id=tokenizer.pad_token_id, eos_token_id=tokenizer.eos_token_id,
+            top_p=args.top_p, temperature=args.temperature, repetition_penalty=3.0,
+        )
+        response = tokenizer.decode(generated_ids[0][len(inputs["input_ids"][0]):], skip_special_tokens=True)
+        conversation.append({"role": "assistant", "content": response})
+        gen_tokens = len(generated_ids[0]) - len(inputs["input_ids"][0])
+        print(f"\n[Speed]: {gen_tokens / (time.time() - st):.2f} tokens/s\n\n") if args.show_speed else print("\n\n")
+
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/playground/generation.py

@@ -16,8 +16,8 @@ def get_args():
     parser = argparse.ArgumentParser()
     parser.add_argument(
         "--pretrained_model_name_or_path",
-        # default=(project_path / "trained_models/gpt2-sst2-generation"),
-        default=(project_path / "trained_models/gpt2-sst2-generation-20260213-2048"),
+        default=(project_path / "trained_models/gpt2-sst2-generation"),
+        # default=(project_path / "trained_models/gpt2-sst2-ppo/checkpoint-150"),
         type=str
     )
     parser.add_argument(
@@ -50,9 +50,9 @@ def main():
 
     tokenized = tokenizer(
         # "this",
-        # "this is ",
+        "this is ",
         # "who needs mind-bending",
-        "eldom has a movie",
+        # "eldom has a movie",
         # "thanks to scott 's charismatic",
         return_tensors="pt"
     )

examples/tutorials/dpo/ultrafeedback-dpo/requirements.txt

@@ -0,0 +1,6 @@
+transformers
+torch
+modelscope
+datasets
+trl
+deepspeed

examples/tutorials/dpo/ultrafeedback-dpo/step_2_train_dpo_model_single_gpu.py

@@ -0,0 +1,208 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+"""
+https://huggingface.co/docs/trl/v0.16.1/en/sft_trainer
+
+单卡 V00 32G 全参微调
+python3 step_2_train_sft_model_single_gpu.py
+
+"""
+import argparse
+import os
+from pathlib import Path
+import platform
+
+# os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
+
+debug_mode = True if platform.system() in ("Windows", "Darwin") else False
+print(f"debug_mode: {debug_mode}")
+
+if platform.system() in ("Windows", "Darwin"):
+    from project_settings import project_path, temp_directory
+else:
+    project_path = os.path.abspath("../../../")
+    project_path = Path(project_path)
+    temp_directory = Path("/root/autodl-tmp/OpenMiniMind/temp")
+
+from datasets import load_dataset
+import torch
+
+from modelscope import AutoModelForCausalLM
+from transformers import AutoTokenizer
+from trl import DPOConfig, DPOTrainer
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model_name",
+        default=(project_path / "pretrained_models/jingyaogong/MiniMind2").as_posix() if debug_mode else "qgyd2021/Qwen2.5-0.5B-ultrachat-sft-deepspeed",
+        type=str
+    ),
+    parser.add_argument(
+        "--dataset_path",
+        default="HuggingFaceH4/ultrafeedback_binarized",
+        # default="miyuki2026/tutorials" if debug_mode else "HuggingFaceH4/ultrachat_200k",
+        type=str
+    ),
+    parser.add_argument(
+        "--dataset_cache_dir",
+        default=(temp_directory / "hub_datasets").as_posix(),
+        type=str
+    ),
+    parser.add_argument(
+        "--model_cache_dir",
+        default=(temp_directory / "hub_models").as_posix(),
+        type=str
+    ),
+    parser.add_argument(
+        "--output_model_dir",
+        default=(temp_directory / "trained_models/qwen2_5-0_5B-ultrafeedback-dpo-single-gpu").as_posix(),
+        type=str
+    ),
+    parser.add_argument(
+        "--num_workers",
+        default=None if debug_mode else os.cpu_count() // 2,
+        type=int
+    ),
+    args = parser.parse_args()
+    return args
+
+
+def format_func(examples, tokenizer):
+    chosen = examples["chosen"]
+    rejected = examples["rejected"]
+
+    chosen_prompt = chosen[:-1]
+    chosen_response = chosen[-1]
+
+    rejected_prompt = rejected[:-1]
+    rejected_response = rejected[-1]
+
+    chosen_prompt_text = tokenizer.apply_chat_template(
+        conversation=chosen_prompt,
+        tokenize=False,
+        add_generation_prompt=True,  # DPO 需要添加生成提示，让模型知道要从这里开始生成
+    )
+    rejected_prompt_text = tokenizer.apply_chat_template(
+        conversation=rejected_prompt,
+        tokenize=False,
+        add_generation_prompt=True,  # DPO 需要添加生成提示，让模型知道要从这里开始生成
+    )
+    if chosen_prompt_text != rejected_prompt_text:
+        raise AssertionError()
+
+    chosen_response_role = chosen_response["role"]
+    chosen_response_text = chosen_response["content"]
+    if chosen_response_role != "assistant":
+        raise AssertionError()
+
+    rejected_response_role = rejected_response["role"]
+    rejected_response_text = rejected_response["content"]
+    if rejected_response_role != "assistant":
+        raise AssertionError()
+
+    result = {
+        "prompt": chosen_prompt_text,
+        "chosen": chosen_response_text,
+        "rejected": rejected_response_text,
+    }
+    return result
+
+
+def main():
+    args = get_args()
+
+    os.environ["MODELSCOPE_CACHE"] = args.model_cache_dir
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    model = AutoModelForCausalLM.from_pretrained(
+        args.model_name,
+        cache_dir=args.model_cache_dir,
+        trust_remote_code=True,
+        dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    )
+    ref_model = AutoModelForCausalLM.from_pretrained(
+        args.model_name,
+        cache_dir=args.model_cache_dir,
+        trust_remote_code=True,
+        dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    )
+    tokenizer = AutoTokenizer.from_pretrained(
+        args.model_name,
+        cache_dir=args.model_cache_dir,
+        trust_remote_code=True,
+    )
+    model = model.to(device)
+    ref_model = ref_model.to(device)
+
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+        tokenizer.pad_token_id = tokenizer.eos_token_id
+
+    print(model)
+    print(ref_model)
+    print(tokenizer)
+
+    dataset_dict = load_dataset(
+        path=args.dataset_path,
+        cache_dir=args.dataset_cache_dir,
+    )
+    train_dataset = dataset_dict["train_prefs"]
+    # test_dataset = dataset_dict["test_prefs"]
+
+    train_dataset = train_dataset.map(
+        lambda x: format_func(x, tokenizer),
+        batched=False,
+        num_proc=args.num_workers,
+        remove_columns=train_dataset.column_names,
+    )
+
+    dpo_config = DPOConfig(
+        output_dir=args.output_model_dir,
+        num_train_epochs=1,
+        per_device_train_batch_size=1 if debug_mode else 2,
+        gradient_accumulation_steps=1 if debug_mode else 8,
+        save_strategy="steps",
+        save_steps=100,
+        save_total_limit=2,
+        logging_steps=10,
+        learning_rate=2e-5,
+        warmup_steps=100,
+        lr_scheduler_type="cosine",
+        fp16=True if torch.cuda.is_available() else False,
+        gradient_checkpointing=False,  # 如果内存紧张，可以设为 True
+        optim="adamw_torch",
+        report_to="none",
+        max_length=1024 if debug_mode else 2048,  # prompt + chosen 的最大长度
+        max_prompt_length=512 if debug_mode else 1024,  # prompt 的最大长度
+        # DPO 特定参数
+        beta=0.1,  # DPO 的温度参数，控制对 preference 的置信度
+        remove_unused_columns=False,
+        dataloader_pin_memory=False,
+    )
+
+    trainer = DPOTrainer(
+        model=model,
+        ref_model=ref_model,  # 提供参考模型
+        args=dpo_config,
+        train_dataset=train_dataset,
+        # DPOTrainer 会自动处理数据，不需要 data_collator
+    )
+
+    # 开始训练
+    print("开始 DPO 训练...")
+    trainer.train()
+
+    # 保存模型
+    print(f"保存模型到: {args.output_model_dir}")
+    trainer.save_model()
+    tokenizer.save_pretrained(args.output_model_dir)
+
+    print("DPO 训练完成!")
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/tutorials/rlhf/gpt2_sst2/step_5_ppo_rlhf.py

@@ -1,12 +1,11 @@
 #!/usr/bin/python3
 # -*- coding: utf-8 -*-
 """
-https://huggingface.co/docs/trl/v0.16.1/en/ppo_trainer
+PPO Training with TRL on SST-2 dataset
+基于 https://huggingface.co/docs/trl/v0.16.1/en/ppo_trainer 的实现
 """
 import argparse
-import copy
 import os
-import random
 from pathlib import Path
 import platform
 from typing import Optional, Tuple, List, Dict, Union
@@ -14,12 +13,12 @@ from typing import Optional, Tuple, List, Dict, Union
 import numpy as np
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 from torch.utils.data import DataLoader
 from datasets import load_dataset
 from transformers import (
-    AutoTokenizer, AutoModelForCausalLM, GPT2PreTrainedModel,
-    GPT2Config, GPT2Model, GPT2LMHeadModel, DataCollatorWithPadding
+    AutoTokenizer,
+    GPT2LMHeadModel,
+    DataCollatorWithPadding
 )
 
 # 路径配置
@@ -29,6 +28,8 @@ else:
     project_path = Path(os.path.abspath("../../../"))
     temp_directory = Path("/root/autodl-tmp/OpenMiniMind/temp")
 
+from trl import PPOTrainer, PPOConfig, AutoModelForCausalLMWithValueHead
+
 
 def get_args():
     parser = argparse.ArgumentParser()
@@ -44,7 +45,7 @@ def get_args():
     parser.add_argument("--valid_dataset_size", default=1000, type=int)
 
     # 训练参数
-    parser.add_argument("--batch_size", default=16, type=int)  # CPU上用小一点的batch
+    parser.add_argument("--batch_size", default=16, type=int)
     parser.add_argument("--ppo_epochs", default=4, type=int)
     parser.add_argument("--mini_batch_size", default=4, type=int)
     parser.add_argument("--kl_beta", default=0.2, type=float)
@@ -63,371 +64,184 @@ def get_args():
 
     # 其他
     parser.add_argument("--num_workers", default=0 if platform.system() == "Windows" else 2, type=int)
-    parser.add_argument("--device", default="cpu", type=str)  # 强制用CPU
+    parser.add_argument("--device", default="cpu", type=str)
 
     return parser.parse_args()
 
 
-class ValueHead(nn.Module):
-    """价值头，为每个token预测一个价值"""
-
-    def __init__(self, hidden_size: int):
-        super().__init__()
-        self.linear = nn.Linear(hidden_size, 1)
-        self._init_weights()
-
-    def _init_weights(self):
-        nn.init.normal_(self.linear.weight, std=1.0 / np.sqrt(self.linear.in_features + 1))
-        nn.init.zeros_(self.linear.bias)
-
-    def forward(self, hidden_states):
-        return self.linear(hidden_states).squeeze(-1)
-
-
-class GPT2ActorCritic(GPT2PreTrainedModel):
-    """Actor-Critic模型，同时输出logits和values"""
-
-    def __init__(self, config: GPT2Config):
-        super().__init__(config)
-        self.lm = GPT2LMHeadModel(config)
-        self.value_head = ValueHead(config.hidden_size)
-        self.post_init()
-
-    def forward(self, input_ids, attention_mask=None):
-        outputs = self.lm(
-            input_ids,
-            attention_mask=attention_mask,
-            output_hidden_states=True
-        )
-        # values来自最后一层hidden states
-        values = self.value_head(outputs.hidden_states[-1])
-        return outputs.logits, values
-
-    def generate(self, *args, **kwargs):
-        return self.lm.generate(*args, **kwargs)
-
-    @classmethod
-    def from_pretrained(cls, pretrained_model_name):
-        """从预训练GPT2LMHeadModel加载"""
-        config = GPT2Config.from_pretrained(pretrained_model_name)
-        model = cls(config)
-        pretrained = GPT2LMHeadModel.from_pretrained(pretrained_model_name)
-        model.lm.load_state_dict(pretrained.state_dict(), strict=False)
-        return model
-
-
-class GPT2RewardModel(GPT2PreTrainedModel):
-    """奖励模型，为每个token预测奖励"""
-
-    def __init__(self, config: GPT2Config):
-        super().__init__(config)
-        self.transformer = GPT2Model(config)
-        self.reward_head = nn.Linear(config.hidden_size, 1)
-        self.post_init()
-
-    def forward(self, input_ids, attention_mask=None):
-        outputs = self.transformer(
-            input_ids,
-            attention_mask=attention_mask,
-            output_hidden_states=True
-        )
-        rewards = self.reward_head(outputs.hidden_states[-1]).squeeze(-1)
-        return torch.sigmoid(rewards)  # [batch, seq_len]
-
-
-class PPOAgent:
-    """PPO训练Agent，封装所有训练逻辑"""
-
-    def __init__(self, args):
-        self.args = args
-        self.device = torch.device(args.device)
-
-        # 加载tokenizer
-        self.tokenizer = AutoTokenizer.from_pretrained(args.sft_model_name)
-        self.tokenizer.pad_token = self.tokenizer.eos_token
-        self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
-
-        # 加载模型
-        print("Loading models...")
-        self.actor_critic = GPT2ActorCritic.from_pretrained(args.sft_model_name).to(self.device)
-        self.reward_model = GPT2RewardModel.from_pretrained(args.reward_model_name).to(self.device)
-        self.reward_model.eval()
-
-        # 参考模型（冻结）
-        self.ref_model = copy.deepcopy(self.actor_critic).to(self.device)
-        self.ref_model.eval()
-
-        # 优化器
-        self.optimizer = torch.optim.Adam(self.actor_critic.parameters(), lr=args.lr)
-
-        # 训练状态
-        self.training_step = 0
-
-    def prepare_dataset(self):
-        """准备训练数据集"""
-        print("Loading dataset...")
-        dataset = load_dataset(
-            path=self.args.dataset_path,
-            cache_dir=self.args.dataset_cache_dir,
-            split="train"
-        )
-
-        def filter_and_truncate(example):
-            # 只保留足够长的句子
-            tokens = self.tokenizer(example["sentence"])["input_ids"]
-            if len(tokens) <= 8:
-                return False
-
-            # 随机截取前2-6个token作为query
-            example["query_ids"] = tokens[:random.randint(2, 6)]
-            return True
-
-        dataset = dataset.filter(filter_and_truncate)
-        dataset = dataset.select(range(min(len(dataset), 5000)))  # CPU上用小数据集
-
-        return dataset
-
-    def collect_rollouts(self, batch):
-        """收集一轮交互数据"""
-        query_ids_list = []
-        response_ids_list = []
-        rewards_list = []
-
-        for i in range(len(batch["query_ids"])):
-            query_ids = torch.tensor(batch["query_ids"][i]).to(self.device)
-            query_ids_list.append(query_ids)
-
-            # 生成response
-            with torch.no_grad():
-                response_len = random.randint(
-                    self.args.min_response_len,
-                    self.args.max_response_len
-                )
-                full_ids = self.actor_critic.generate(
-                    input_ids=query_ids.unsqueeze(0),
-                    max_new_tokens=response_len,
-                    do_sample=True,
-                    top_p=self.args.top_p,
-                    temperature=self.args.temperature,
-                    pad_token_id=self.tokenizer.pad_token_id,
-                    eos_token_id=self.tokenizer.eos_token_id,
-                )[0]
-
-                response_ids = full_ids[len(query_ids):]
-                response_ids_list.append(response_ids)
-
-                # 计算奖励（只取最后一个token的奖励）
-                reward = self.reward_model(
-                    full_ids.unsqueeze(0),
-                    attention_mask=torch.ones_like(full_ids).unsqueeze(0)
-                )[0, -1]
-                # 缩放到[-1, 1]
-                rewards_list.append(2 * (reward - 0.5))
-
-        return query_ids_list, response_ids_list, rewards_list
-
-    def compute_advantages_and_returns(self, log_probs, values, rewards, masks):
-        """计算GAE advantages和returns"""
-        seq_len = rewards.shape[1]
-        advantages = torch.zeros_like(rewards)
-        returns = torch.zeros_like(rewards)
-
-        gae = 0
-        for t in reversed(range(seq_len)):
-            if t == seq_len - 1:
-                next_value = 0
-            else:
-                next_value = values[:, t + 1]
-
-            delta = rewards[:, t] + self.args.gamma * next_value - values[:, t]
-            gae = delta + self.args.gamma * self.args.lam * gae
-            advantages[:, t] = gae
-            returns[:, t] = advantages[:, t] + values[:, t]
-
-        # 只对有效位置进行whiten
-        advantages = self.masked_whiten(advantages, masks)
-        return advantages, returns
-
-    def masked_whiten(self, values, mask):
-        """带mask的whitening"""
-        mask = mask.float()
-        mean = (values * mask).sum() / mask.sum()
-        var = (((values - mean) * mask) ** 2).sum() / mask.sum()
-        whitened = (values - mean) * torch.rsqrt(var + 1e-8)
-        return whitened * mask
-
-    def ppo_step(self, batch_data):
-        """单步PPO更新"""
-        (query_ids_list, response_ids_list, old_log_probs,
-         advantages, returns, masks) = batch_data
-
-        # 拼接完整的query+response
-        full_ids_list = []
-        for q, r in zip(query_ids_list, response_ids_list):
-            full_ids_list.append(torch.cat([q, r]))
-
-        # padding
-        padded = self.tokenizer.pad(
-            {"input_ids": full_ids_list},
-            padding=True,
-            return_tensors="pt"
-        )
-        input_ids = padded["input_ids"].to(self.device)
-        attention_mask = padded["attention_mask"].to(self.device)
-
-        # 前向传播
-        logits, values = self.actor_critic(input_ids, attention_mask)
-
-        # 计算新的log_probs
-        log_probs = F.log_softmax(logits[:, :-1, :], dim=-1)
-        log_probs = torch.gather(
-            log_probs, 2,
-            input_ids[:, 1:].unsqueeze(-1)
-        ).squeeze(-1)
-
-        # 只保留response部分的log_probs
-        response_start = [len(q) for q in query_ids_list]
-        new_log_probs = []
-        for i, start in enumerate(response_start):
-            new_log_probs.append(log_probs[i, start - 1:start - 1 + len(response_ids_list[i])])
-        new_log_probs = torch.cat(new_log_probs)
-
-        # 计算ratio和PPO损失
-        old_log_probs = old_log_probs.detach()
-        ratio = torch.exp(new_log_probs - old_log_probs)
-
-        # 裁剪的policy loss
-        surr1 = ratio * advantages
-        surr2 = torch.clamp(ratio, 1 - self.args.clip_epsilon,
-                            1 + self.args.clip_epsilon) * advantages
-        policy_loss = -torch.min(surr1, surr2).mean()
-
-        # value loss
-        value_pred = []
-        for i, start in enumerate(response_start):
-            value_pred.append(values[i, start - 1:start - 1 + len(response_ids_list[i])])
-        value_pred = torch.cat(value_pred)
-        value_loss = F.mse_loss(value_pred, returns)
-
-        # 总loss
-        loss = policy_loss + 0.5 * value_loss
-
-        return loss, policy_loss, value_loss
-
-    def train_epoch(self, dataset):
-        """训练一个epoch"""
-        total_policy_loss = 0
-        total_value_loss = 0
-        num_batches = 0
-
-        for batch_idx in range(0, len(dataset), self.args.batch_size):
-            # 1. 收集数据
-            batch = dataset[batch_idx:batch_idx + self.args.batch_size]
-            query_ids_list, response_ids_list, rewards_list = self.collect_rollouts(batch)
-
-            # 2. 计算旧的log_probs和values
-            old_log_probs_list = []
-            values_list = []
-            masks_list = []
-
-            with torch.no_grad():
-                for q_ids, r_ids in zip(query_ids_list, response_ids_list):
-                    full_ids = torch.cat([q_ids, r_ids]).unsqueeze(0).to(self.device)
-                    attn_mask = torch.ones_like(full_ids)
-
-                    logits, values = self.actor_critic(full_ids, attn_mask)
-
-                    # 计算response部分的log_probs
-                    log_probs = F.log_softmax(logits[:, :-1, :], dim=-1)
-                    log_probs = torch.gather(
-                        log_probs, 2,
-                        full_ids[:, 1:].unsqueeze(-1)
-                    ).squeeze(-1)
-
-                    start = len(q_ids) - 1
-                    end = start + len(r_ids)
-                    old_log_probs_list.append(log_probs[0, start:end])
-                    values_list.append(values[0, start:end])
-
-                    # 创建mask
-                    mask = torch.zeros(len(r_ids))
-                    mask[-1] = 1  # 最后一个token有真实奖励
-                    masks_list.append(mask)
-
-            # 转换为tensor
-            old_log_probs = torch.cat(old_log_probs_list).to(self.device)
-            values = torch.cat(values_list).to(self.device)
-            masks = torch.cat(masks_list).to(self.device)
-            rewards = torch.zeros_like(values).to(self.device)
-
-            # 设置奖励（只在最后一个token加上环境奖励）
-            for i, (r, mask) in enumerate(zip(rewards_list, masks_list)):
-                if mask[-1] > 0:
-                    # KL惩罚
-                    kl = old_log_probs[i] - old_log_probs[i]  # 这里简化了，实际要用ref_model
-                    kl_penalty = -self.args.kl_beta * kl
-                    rewards[i] = kl_penalty + r
-
-            # 3. 计算advantages和returns
-            advantages, returns = self.compute_advantages_and_returns(
-                old_log_probs.unsqueeze(0),
-                values.unsqueeze(0),
-                rewards.unsqueeze(0),
-                masks.unsqueeze(0)
-            )
-
-            # 4. PPO多次更新
-            batch_data = (query_ids_list, response_ids_list, old_log_probs,
-                          advantages.squeeze(0), returns.squeeze(0), masks)
-
-            for _ in range(self.args.ppo_epochs):
-                loss, policy_loss, value_loss = self.ppo_step(batch_data)
-
-                self.optimizer.zero_grad()
-                loss.backward()
-                torch.nn.utils.clip_grad_norm_(self.actor_critic.parameters(), 1.0)
-                self.optimizer.step()
-
-                total_policy_loss += policy_loss.item()
-                total_value_loss += value_loss.item()
-                num_batches += 1
-                self.training_step += 1
-
-            if batch_idx % 100 == 0:
-                print(f"Batch {batch_idx}/{len(dataset)}: "
-                      f"policy_loss={total_policy_loss / num_batches:.4f}, "
-                      f"value_loss={total_value_loss / num_batches:.4f}")
-
-        return total_policy_loss / num_batches, total_value_loss / num_batches
-
-    def train(self):
-        """主训练循环"""
-        dataset = self.prepare_dataset()
-        print(f"Dataset size: {len(dataset)}")
-
-        for epoch in range(self.args.max_epochs):
-            print(f"\n=== Epoch {epoch + 1}/{self.args.max_epochs} ===")
-            policy_loss, value_loss = self.train_epoch(dataset)
-            print(f"Epoch {epoch + 1} finished: "
-                  f"policy_loss={policy_loss:.4f}, value_loss={value_loss:.4f}")
+def build_dataset(tokenizer, dataset_path, dataset_cache_dir, valid_dataset_size):
+    """
+    构建SST-2数据集，返回query列表
+    """
+    dataset = load_dataset(
+        dataset_path,
+        cache_dir=dataset_cache_dir,
+        split="train"
+    )
+
+    # 只取前valid_dataset_size条数据用于演示
+    dataset = dataset.select(range(min(valid_dataset_size, len(dataset))))
+
+    def tokenize_function(examples):
+        return tokenizer(examples["sentence"], truncation=True, max_length=128)
+
+    dataset = dataset.map(tokenize_function, batched=True)
+    dataset.set_format(type='torch', columns=['input_ids', 'attention_mask'])
+
+    return dataset
+
+
+class RewardModelWrapper:
+    """
+    奖励模型包装器，用于计算生成文本的奖励
+    """
+
+    def __init__(self, reward_model_name, tokenizer, device):
+        self.device = device
+        self.tokenizer = tokenizer
+        # 加载你的GPT2RewardModel或标准模型
+        from transformers import GPT2ForSequenceClassification
+        self.model = GPT2ForSequenceClassification.from_pretrained(reward_model_name).to(device)
+        self.model.eval()
+
+    def get_reward(self, texts: List[str]) -> List[float]:
+        """
+        计算文本的奖励分数（SST-2情感分类）
+        """
+        inputs = self.tokenizer(texts, padding=True, truncation=True, return_tensors="pt").to(self.device)
+        with torch.no_grad():
+            outputs = self.model(**inputs)
+            # SST-2是二分类，取正面情感的logits作为奖励
+            logits = outputs.logits
+            # 如果模型输出是logits，使用softmax获取正面概率
+            probs = torch.softmax(logits, dim=-1)
+            # 假设标签1是正面
+            rewards = probs[:, 1].cpu().tolist()
+        return rewards
 
 
 def main():
     args = get_args()
     print("PPO Training with CPU")
-    print(f"Arguments: {args}")
 
-    # 创建agent并开始训练
-    agent = PPOAgent(args)
-    agent.train()
+    # 设备设置
+    device = torch.device(args.device)
+
+    # 1. 加载tokenizer
+    tokenizer = AutoTokenizer.from_pretrained(args.sft_model_name, cache_dir=args.model_cache_dir)
+    tokenizer.pad_token = tokenizer.eos_token
+
+    # 2. 构建数据集
+    dataset = build_dataset(
+        tokenizer,
+        args.dataset_path,
+        args.dataset_cache_dir,
+        args.valid_dataset_size
+    )
+
+    # 3. 加载模型（使用TRL的AutoModelForCausalLMWithValueHead）
+    # 这会在原有LM基础上自动添加value head
+    model = AutoModelForCausalLMWithValueHead.from_pretrained(args.sft_model_name)
+    model.to(device)
+
+    # 4. 加载参考模型（用于KL散度计算）
+    ref_model = AutoModelForCausalLMWithValueHead.from_pretrained(args.sft_model_name)
+    ref_model.to(device)
+
+    # 5. 加载奖励模型
+    reward_model = RewardModelWrapper(args.reward_model_name, tokenizer, device)
+
+    # 6. 配置PPO参数
+    ppo_config = PPOConfig(
+        model_name=args.sft_model_name,
+        learning_rate=args.lr,
+        batch_size=args.batch_size,
+        mini_batch_size=args.mini_batch_size,
+        ppo_epochs=args.ppo_epochs,
+        clip_epsilon=args.clip_epsilon,
+        gamma=args.gamma,
+        lam=args.lam,
+        kl_penalty=args.kl_beta,
+        device=device,
+        log_with=None,  # 可设置"wandb"等
+        project_kwargs={"logging_dir": "./logs"},
+    )
+
+    # 7. 初始化PPO Trainer
+    ppo_trainer = PPOTrainer(
+        config=ppo_config,
+        model=model,
+        ref_model=ref_model,
+        tokenizer=tokenizer,
+        dataset=dataset,
+        data_collator=DataCollatorWithPadding(tokenizer)
+    )
+
+    # 8. 定义response长度采样器（在min和max之间随机）
+    response_length_sampler = LengthSampler(args.min_response_len, args.max_response_len)
+
+    # 9. 训练循环
+    generation_kwargs = {
+        "min_length": -1,  # 不强制最小长度
+        "top_p": args.top_p,
+        "temperature": args.temperature,
+        "do_sample": True,
+        "pad_token_id": tokenizer.eos_token_id,
+        "max_new_tokens": args.max_new_tokens,
+    }
+
+    for epoch in range(args.max_epochs):
+        print(f"Epoch {epoch + 1}/{args.max_epochs}")
+
+        for batch_idx, batch in enumerate(ppo_trainer.dataloader):
+            # 获取query的input_ids
+            query_tensors = batch["input_ids"]
+
+            # 使用模型生成response
+            response_tensors = []
+            for query in query_tensors:
+                # query已经是tensor，添加batch维度
+                query = query.unsqueeze(0).to(device)
+
+                # 生成response（这里使用respond_to_batch工具函数）
+                response = respond_to_batch(
+                    model,
+                    query,
+                    length_sampler=response_length_sampler,
+                    **generation_kwargs
+                )
+                response_tensors.append(response.squeeze())
+
+            # 解码生成的文本
+            responses = [tokenizer.decode(r, skip_special_tokens=True) for r in response_tensors]
+            queries = [tokenizer.decode(q, skip_special_tokens=True) for q in query_tensors]
+
+            # 使用奖励模型计算奖励
+            # 这里我们结合query和response作为完整文本进行情感分析
+            full_texts = [q + " " + r for q, r in zip(queries, responses)]
+            rewards = reward_model.get_reward(full_texts)
+
+            # 转换为tensor
+            rewards = [torch.tensor(r, device=device) for r in rewards]
+
+            # 执行PPO更新步骤
+            stats = ppo_trainer.step(query_tensors, response_tensors, rewards)
+
+            # 打印统计信息
+            if batch_idx % 10 == 0:
+                print(f"Batch {batch_idx}, mean reward: {np.mean(rewards):.4f}")
+                print(f"Stats: {stats}")
+                print(f"Example query: {queries[0]}")
+                print(f"Example response: {responses[0]}")
+                print(f"Reward: {rewards[0].item():.4f}")
+                print("-" * 50)
 
-    # 保存模型
-    output_dir = Path(args.sft_model_name) / "ppo_trained"
-    output_dir.mkdir(exist_ok=True, parents=True)
-    agent.actor_critic.save_pretrained(output_dir)
-    agent.tokenizer.save_pretrained(output_dir)
-    print(f"Model saved to {output_dir}")
+        # 每个epoch保存一次模型
+        save_path = Path("ppo_models") / f"epoch_{epoch}"
+        ppo_trainer.save_pretrained(save_path)
+        tokenizer.save_pretrained(save_path)
+        print(f"Model saved to {save_path}")
 
 
 if __name__ == "__main__":
-    main()
+    main()

examples/tutorials/rlhf/gpt2_sst2_generation/step_2_train_model.py

@@ -0,0 +1,172 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+"""
+用sst的句子训练gpt2模型，让其随机生成一些评论。
+"""
+import argparse
+import os
+from pathlib import Path
+import platform
+
+if platform.system() in ("Windows", "Darwin"):
+    from project_settings import project_path, temp_directory
+else:
+    project_path = os.path.abspath("../../../")
+    project_path = Path(project_path)
+    temp_directory = Path("/root/autodl-tmp/OpenMiniMind/temp")
+
+from datasets import load_dataset
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, DataCollatorForLanguageModeling, Trainer, TrainingArguments
+from transformers import GPT2LMHeadModel
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model_name",
+        # default="openai-community/gpt2",
+        default=(project_path / "pretrained_models/openai-community/gpt2").as_posix(),
+        type=str
+    ),
+    parser.add_argument(
+        "--dataset_path",
+        default="stanfordnlp/sst2",
+        type=str
+    ),
+    parser.add_argument("--dataset_name", default=None, type=str),
+    parser.add_argument("--dataset_split", default=None, type=str),
+    parser.add_argument(
+        "--dataset_cache_dir",
+        default=(temp_directory / "hub_datasets").as_posix(),
+        type=str
+    ),
+    parser.add_argument(
+        "--model_cache_dir",
+        default=(temp_directory / "hub_models").as_posix(),
+        type=str
+    ),
+    parser.add_argument("--dataset_streaming", default=None, type=str),
+    parser.add_argument("--valid_dataset_size", default=1000, type=int),
+    parser.add_argument("--shuffle_buffer_size", default=5000, type=int),
+
+    parser.add_argument(
+        "--output_model_dir",
+        default=(project_path / "trained_models/gpt2-sst2-generation-epoch-3").as_posix(),
+        type=str
+    ),
+
+    parser.add_argument(
+        "--num_workers",
+        default=None if platform.system() in ("Windows", "Darwin") else os.cpu_count() // 2,
+        type=int
+    ),
+    parser.add_argument(
+        "--device",
+        default=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
+        type=int
+    ),
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = get_args()
+
+    model = AutoModelForCausalLM.from_pretrained(args.model_name)
+    model = model.to(args.device)
+    tokenizer = AutoTokenizer.from_pretrained(args.model_name)
+    tokenizer.pad_token = tokenizer.eos_token
+
+    dataset_dict = load_dataset(
+        path=args.dataset_path,
+        name=args.dataset_name,
+        split=args.dataset_split,
+        cache_dir=args.dataset_cache_dir,
+        # num_proc=args.num_workers if not args.dataset_streaming else None,
+        streaming=args.dataset_streaming,
+    )
+    train_dataset = dataset_dict["train"]
+    valid_dataset = dataset_dict["validation"]
+    # test_dataset = dataset_dict["test"]
+
+    def format_func(example):
+        sentence = example["sentence"]
+        sentence += tokenizer.eos_token
+        tokenized = tokenizer(sentence)
+        input_ids = tokenized["input_ids"]
+        attention_mask = tokenized["attention_mask"]
+        # print(input_ids)
+        # print(attention_mask)
+        result = {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+        }
+        return result
+
+    train_dataset = train_dataset.map(
+        format_func,
+        batched=False,
+        remove_columns=train_dataset.column_names,
+    )
+    valid_dataset = valid_dataset.map(
+        format_func,
+        batched=False,
+        remove_columns=valid_dataset.column_names,
+    )
+    print(f"train_dataset size: {len(train_dataset)}")
+    print(f"valid_dataset size: {len(valid_dataset)}")
+
+    train_dataset = train_dataset.filter(
+        function=lambda x: 5 < len(x["input_ids"]) < 1024
+    )
+    valid_dataset = valid_dataset.filter(
+        function=lambda x: 5 < len(x["input_ids"]) < 1024
+    )
+    print(f"train_dataset size: {len(train_dataset)}")
+    print(f"valid_dataset size: {len(valid_dataset)}")
+
+    data_collator = DataCollatorForLanguageModeling(
+        tokenizer,
+        mlm=False
+    )
+
+    training_args = TrainingArguments(
+        output_dir=args.output_model_dir,
+        # overwrite_output_dir=True,
+        num_train_epochs=3,
+        per_device_train_batch_size=16,
+        per_device_eval_batch_size=16,
+        eval_strategy="steps",
+        eval_steps=100,
+        save_strategy="steps",
+        save_steps=100,
+        save_total_limit=2,
+        logging_steps=100,
+        learning_rate=5e-5,
+        warmup_steps=500,
+        weight_decay=0.01,
+        fp16=torch.cuda.is_available(),
+        dataloader_num_workers=args.num_workers or 0,
+        remove_unused_columns=False,
+        load_best_model_at_end=False,
+        # metric_for_best_model="eval_loss",
+        # greater_is_better=False,
+    )
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        data_collator=data_collator,
+        train_dataset=train_dataset,
+        eval_dataset=valid_dataset,
+        tokenizer=tokenizer,
+    )
+
+    trainer.train()
+    trainer.save_model()
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/tutorials/rlhf/gpt2_sst2_generation/step_3_generation.py

@@ -0,0 +1,78 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import argparse
+import time
+
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
+
+from project_settings import project_path
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--pretrained_model_name_or_path",
+        default=(project_path / "trained_models/gpt2-sst2-generation-epoch-3"),
+        # default=(project_path / "trained_models/gpt2-sst2-ppo/checkpoint-150"),
+        # default=(project_path / "trained_models/gpt2-sst2-generation-epoch-3/checkpoint-5400"),
+        type=str
+    )
+    parser.add_argument(
+        "--max_new_tokens",
+        default=1024, # 8192, 128
+        type=int, help="最大生成长度（注意：并非模型实际长文本能力）"
+    )
+    parser.add_argument("--top_p", default=0.85, type=float, help="nucleus采样阈值（0-1）")
+    parser.add_argument("--temperature", default=0.85, type=float, help="生成温度，控制随机性（0-1，越大越随机）")
+
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = get_args()
+
+    if torch.cuda.is_available():
+        device = "cuda"
+    elif torch.backends.mps.is_available():
+        # device = "mps"
+        device = "cpu"
+    else:
+        device = "cpu"
+    print(f"device: {device}")
+
+    tokenizer = AutoTokenizer.from_pretrained(args.pretrained_model_name_or_path)
+    model = AutoModelForCausalLM.from_pretrained(args.pretrained_model_name_or_path)
+    model = model.eval().to(device)
+
+    tokenized = tokenizer(
+        # "this",
+        # "this is ",
+        "it 's ",
+        # "please listen ",
+        # "eldom has a movie",
+        # "thanks to scott 's charismatic",
+        return_tensors="pt"
+    )
+
+    streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=False)
+
+    generated_ids = model.generate(
+        inputs=tokenized["input_ids"], attention_mask=tokenized["attention_mask"],
+        max_new_tokens=args.max_new_tokens, do_sample=True, streamer=streamer,
+        pad_token_id=tokenizer.pad_token_id, eos_token_id=tokenizer.eos_token_id,
+        top_p=args.top_p, temperature=args.temperature, repetition_penalty=3.0,
+        early_stopping=True,
+    )
+    # response = tokenizer.decode(generated_ids[0][len(tokenized["input_ids"][0]):], skip_special_tokens=True)
+    response = tokenizer.decode(generated_ids[0], skip_special_tokens=False)
+    print(response)
+    # print(generated_ids)
+    print(f"count: {generated_ids.shape}")
+
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/tutorials/rlhf/gpt2_sst2_ppo/requirements.txt

@@ -0,0 +1,2 @@
+trl==0.16.1
+transformers==4.50.2

examples/tutorials/rlhf/gpt2_sst2_ppo/step_1_prepare_data.py

@@ -0,0 +1,58 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+"""
+或使用命令行
+pip install modelscope
+modelscope download \
+--model 'qgyd2021/gpt2-for-sequence-classification-sst2-reward' \
+--local_dir '/root/autodl-tmp/trained_models/Qwen3-8B-sft-deepspeed'
+
+python3 step_1_prepare_data.py \
+--repo_id qgyd2021/gpt2-for-sequence-classification-sst2-reward \
+--local_dir /root/autodl-tmp/OpenMiniMind/trained_models/gpt2-for-sequence-classification-sst2-reward
+
+python3 step_1_prepare_data.py \
+--repo_id qgyd2021/gpt2-sst2-generation-epoch-3 \
+--local_dir /root/autodl-tmp/OpenMiniMind/trained_models/gpt2-sst2-generation-epoch-3
+
+
+"""
+import argparse
+import os
+from pathlib import Path
+import platform
+
+if platform.system() in ("Windows", "Darwin"):
+    from project_settings import project_path, temp_directory
+else:
+    project_path = os.path.abspath("../../../")
+    project_path = Path(project_path)
+    temp_directory = Path("/root/autodl-tmp/OpenMiniMind/temp")
+
+from modelscope import snapshot_download
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--repo_id", default="qgyd2021/Qwen3-8B-sft-deepspeed", type=str)
+    parser.add_argument(
+        "--local_dir",
+        default=(temp_directory / "../trained_models/Qwen3-8B-sft-deepspeed").as_posix(),
+        type=str
+    )
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = get_args()
+
+    snapshot_download(
+        model_id=args.repo_id,
+        local_dir=args.local_dir,
+    )
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/tutorials/rlhf/gpt2_sst2_ppo/step_2_train_model_on_cpu.py

@@ -0,0 +1,217 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+"""
+PPO Training with TRL on SST-2 dataset
+基于 https://huggingface.co/docs/trl/v0.16.1/en/ppo_trainer 的实现
+
+
+（1）策略模型 policy 根据 queries 生成 query_response 和 logits
+（2）根据 logits 计算 logprob 概率，再索引出 response 对应的概率。
+（3）参考模型 ref_policy 根据 query_response 计算其 ref_logits
+（4）根据 ref_logits 计算 ref_logprob 概率，再索引出 response 对应的概率。
+（5）query_response 中生成的第一个 eos_token 开始后面的 token 都替换为 pad_token。
+（6）价值模型 value_model 计算 query_response 中 response 部分每个token的价值。
+（7）奖励模型 reward_model 计算 postprocessed_query_response 中最后一个非 pad_token 的奖励。
+（8）得到：
+（9）kl = logprobs - ref_logprobs
+non_score_reward = -args.kl_coef * kl
+advantages
+returns = advantages + values
+
+"""
+import argparse
+import os
+from pathlib import Path
+import platform
+
+import torch
+from datasets import load_dataset
+from transformers import (
+    AutoTokenizer, AutoModelForCausalLM, AutoModelForSequenceClassification,
+    GPT2LMHeadModel, GPT2ForSequenceClassification,
+    DataCollatorWithPadding
+)
+
+# 路径配置
+if platform.system() in ("Windows", "Darwin"):
+    from project_settings import project_path, temp_directory
+else:
+    project_path = Path(os.path.abspath("../../../"))
+    temp_directory = Path("/root/autodl-tmp/OpenMiniMind/temp")
+
+from trl import PPOTrainer, PPOConfig
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--sft_model_name", type=str,
+                        default=(project_path / "trained_models/gpt2-sst2-generation-epoch-3").as_posix())
+    parser.add_argument("--reward_model_name", type=str,
+                        default=(project_path / "trained_models/gpt2-for-sequence-classification-sst2-reward").as_posix())
+    parser.add_argument("--dataset_path", default="stanfordnlp/sst2", type=str)
+    parser.add_argument("--dataset_cache_dir",
+                        default=(temp_directory / "hub_datasets").as_posix(), type=str)
+    parser.add_argument("--model_cache_dir",
+                        default=(temp_directory / "hub_models").as_posix(), type=str)
+
+    # 训练参数
+
+    # 生成参数
+
+    parser.add_argument(
+        "--output_model_dir",
+        default=(project_path / "trained_models/gpt2-sst2-ppo").as_posix(),
+        type=str
+    ),
+
+    # 其他
+    parser.add_argument("--num_workers", default=0 if platform.system() == "Windows" else 2, type=int)
+    parser.add_argument("--device", default="cpu", type=str)
+
+    return parser.parse_args()
+
+
+def format_func(example, tokenizer):
+    sentence: str = example["sentence"]
+    # score: float = float(example["label"])
+    tokenized = tokenizer(sentence)
+    input_ids = tokenized["input_ids"]
+    attention_mask = tokenized["attention_mask"]
+    result = {
+        "input_ids": input_ids,
+        "attention_mask": attention_mask,
+    }
+    return result
+
+
+def token_truncate(example, tokenizer):
+    input_ids = example["input_ids"]
+    attention_mask = example["attention_mask"]
+    input_ids = input_ids[:3]
+    attention_mask = attention_mask[:3]
+    # text = tokenizer.decode(input_ids)
+    result = {
+        "input_ids": input_ids,
+        "attention_mask": attention_mask,
+        # "text": text,
+    }
+    return result
+
+
+def main():
+    args = get_args()
+
+    # 设备设置
+    device = torch.device(args.device)
+
+    # 1. 加载tokenizer
+    tokenizer = AutoTokenizer.from_pretrained(
+        args.sft_model_name,
+        padding_side="left",  # 对于生成任务很重要
+        cache_dir=args.model_cache_dir,
+    )
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+        tokenizer.pad_token_id = tokenizer.eos_token_id
+    print(f"eos_token: {tokenizer.eos_token}")
+    print(f"pad_token: {tokenizer.pad_token}")
+
+    model = AutoModelForCausalLM.from_pretrained(args.sft_model_name)
+    value_model = AutoModelForSequenceClassification.from_pretrained(
+        args.sft_model_name,
+        num_labels=1
+    )
+    value_model.transformer = model.transformer
+
+    ref_model = AutoModelForCausalLM.from_pretrained(args.sft_model_name)
+
+    reward_model = AutoModelForSequenceClassification.from_pretrained(
+        args.reward_model_name,
+        num_labels=1
+    )
+
+    dataset_dict = load_dataset(
+        path=args.dataset_path,
+        cache_dir=args.dataset_cache_dir,
+    )
+    train_dataset = dataset_dict["train"]
+    valid_dataset = dataset_dict["validation"]
+    # test_dataset = dataset_dict["test"]
+
+    train_dataset = train_dataset.map(
+        lambda example: format_func(example, tokenizer),
+        batched=False,
+        remove_columns=train_dataset.column_names,
+    )
+    valid_dataset = valid_dataset.map(
+        lambda example: format_func(example, tokenizer),
+        batched=False,
+        remove_columns=valid_dataset.column_names,
+    )
+    train_dataset = train_dataset.filter(
+        function=lambda x: len(x["input_ids"]) > 8
+    )
+    valid_dataset = valid_dataset.filter(
+        function=lambda x: len(x["input_ids"]) > 8
+    )
+    train_dataset = train_dataset.map(
+        lambda example: token_truncate(example, tokenizer),
+        batched=False,
+        remove_columns=train_dataset.column_names,
+    )
+    valid_dataset = valid_dataset.map(
+        lambda example: token_truncate(example, tokenizer),
+        batched=False,
+        remove_columns=valid_dataset.column_names,
+    )
+
+    ppo_config = PPOConfig(
+        output_dir=args.output_model_dir,
+        num_train_epochs=1,
+        eval_strategy="steps",
+        eval_steps=50,
+        save_strategy="steps",
+        save_steps=50,
+        save_total_limit=2,
+        logging_steps=50,
+        learning_rate=1e-5,
+        warmup_steps=50,
+        per_device_eval_batch_size=10,
+
+        num_mini_batches=2,
+        num_sample_generations=100,
+        # total_episodes=100000,  # 最多训练多少个样本。
+        response_length=64,
+        # stop_token=tokenizer.eos_token, # stop_token 和 stop_token_id 只设置一个。
+        stop_token_id=tokenizer.eos_token_id,
+        batch_size=16,
+
+        num_ppo_epochs=1,
+        whiten_rewards=True,
+        gamma=1.0,
+        lam=0.95,
+
+        dataset_num_proc=args.num_workers,
+    )
+
+    data_collator = DataCollatorWithPadding(tokenizer)
+
+    ppo_trainer = PPOTrainer(
+        args=ppo_config,
+        processing_class=tokenizer,
+        model=model,
+        ref_model=ref_model,
+        reward_model=reward_model,
+        train_dataset=train_dataset,
+        value_model=value_model,
+        data_collator=data_collator,
+        eval_dataset=valid_dataset,
+    )
+    ppo_trainer.train()
+    ppo_trainer.save_model()
+
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/tutorials/rlhf/gpt2_sst2_ppo/step_2_train_model_two_gpu.py

@@ -0,0 +1,201 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+"""
+PPO Training with TRL on SST-2 dataset
+基于 https://huggingface.co/docs/trl/v0.16.1/en/ppo_trainer 的实现
+
+
+双卡 V100
+
+"""
+import argparse
+import os
+from pathlib import Path
+import platform
+
+import torch
+from datasets import load_dataset
+from transformers import (
+    AutoTokenizer, AutoModelForCausalLM, AutoModelForSequenceClassification,
+    GPT2LMHeadModel, GPT2ForSequenceClassification,
+    DataCollatorWithPadding
+)
+
+# 路径配置
+if platform.system() in ("Windows", "Darwin"):
+    from project_settings import project_path, temp_directory
+else:
+    project_path = Path(os.path.abspath("../../../"))
+    temp_directory = Path("/root/autodl-tmp/OpenMiniMind/temp")
+
+from trl import PPOTrainer, PPOConfig
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--sft_model_name", type=str,
+                        default=(project_path / "trained_models/gpt2-sst2-generation-epoch-3").as_posix())
+    parser.add_argument("--reward_model_name", type=str,
+                        default=(project_path / "trained_models/gpt2-for-sequence-classification-sst2-reward").as_posix())
+    parser.add_argument("--dataset_path", default="stanfordnlp/sst2", type=str)
+    parser.add_argument("--dataset_cache_dir",
+                        default=(temp_directory / "hub_datasets").as_posix(), type=str)
+    parser.add_argument("--model_cache_dir",
+                        default=(temp_directory / "hub_models").as_posix(), type=str)
+
+    # 训练参数
+
+    # 生成参数
+
+    parser.add_argument(
+        "--output_model_dir",
+        default=(project_path / "trained_models/gpt2-sst2-ppo").as_posix(),
+        type=str
+    ),
+
+    # 其他
+    parser.add_argument("--num_workers", default=0 if platform.system() == "Windows" else 2, type=int)
+
+    return parser.parse_args()
+
+
+def format_func(example, tokenizer):
+    sentence: str = example["sentence"]
+    # score: float = float(example["label"])
+    tokenized = tokenizer(sentence)
+    input_ids = tokenized["input_ids"]
+    attention_mask = tokenized["attention_mask"]
+    result = {
+        "input_ids": input_ids,
+        "attention_mask": attention_mask,
+    }
+    return result
+
+
+def token_truncate(example, tokenizer):
+    input_ids = example["input_ids"]
+    attention_mask = example["attention_mask"]
+    input_ids = input_ids[:3]
+    attention_mask = attention_mask[:3]
+    # text = tokenizer.decode(input_ids)
+    result = {
+        "input_ids": input_ids,
+        "attention_mask": attention_mask,
+        # "text": text,
+    }
+    return result
+
+
+def main():
+    args = get_args()
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        args.sft_model_name,
+        padding_side="left",  # 对于生成任务很重要
+        cache_dir=args.model_cache_dir,
+    )
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+        tokenizer.pad_token_id = tokenizer.eos_token_id
+    print(f"eos_token: {tokenizer.eos_token}")
+    print(f"pad_token: {tokenizer.pad_token}")
+
+    model = AutoModelForCausalLM.from_pretrained(args.sft_model_name)
+    value_model = AutoModelForSequenceClassification.from_pretrained(
+        args.sft_model_name,
+        num_labels=1
+    )
+    value_model.transformer = model.transformer
+
+    ref_model = AutoModelForCausalLM.from_pretrained(args.sft_model_name)
+    reward_model = AutoModelForSequenceClassification.from_pretrained(
+        args.reward_model_name,
+        num_labels=1
+    )
+
+    dataset_dict = load_dataset(
+        path=args.dataset_path,
+        cache_dir=args.dataset_cache_dir,
+    )
+    train_dataset = dataset_dict["train"]
+    valid_dataset = dataset_dict["validation"]
+    # test_dataset = dataset_dict["test"]
+
+    train_dataset = train_dataset.map(
+        lambda example: format_func(example, tokenizer),
+        batched=False,
+        remove_columns=train_dataset.column_names,
+    )
+    valid_dataset = valid_dataset.map(
+        lambda example: format_func(example, tokenizer),
+        batched=False,
+        remove_columns=valid_dataset.column_names,
+    )
+    train_dataset = train_dataset.filter(
+        function=lambda x: len(x["input_ids"]) > 8
+    )
+    valid_dataset = valid_dataset.filter(
+        function=lambda x: len(x["input_ids"]) > 8
+    )
+    train_dataset = train_dataset.map(
+        lambda example: token_truncate(example, tokenizer),
+        batched=False,
+        remove_columns=train_dataset.column_names,
+    )
+    valid_dataset = valid_dataset.map(
+        lambda example: token_truncate(example, tokenizer),
+        batched=False,
+        remove_columns=valid_dataset.column_names,
+    )
+
+    ppo_config = PPOConfig(
+        output_dir=args.output_model_dir,
+        num_train_epochs=1,
+        eval_strategy="steps",
+        eval_steps=50,
+        save_strategy="steps",
+        save_steps=50,
+        save_total_limit=2,
+        logging_steps=50,
+        learning_rate=1e-5,
+        warmup_steps=50,
+        per_device_eval_batch_size=10,
+
+        num_mini_batches=2,
+        num_sample_generations=100,
+        # total_episodes=100000,  # 最多训练多少个样本。
+        response_length=1024,
+        # stop_token=tokenizer.eos_token, # stop_token 和 stop_token_id 只设置一个。
+        stop_token_id=tokenizer.eos_token_id,
+        batch_size=16,
+
+        num_ppo_epochs=1,
+        whiten_rewards=True,
+        gamma=1.0,
+        lam=0.95,
+
+        dataset_num_proc=args.num_workers,
+    )
+
+    data_collator = DataCollatorWithPadding(tokenizer)
+
+    ppo_trainer = PPOTrainer(
+        args=ppo_config,
+        processing_class=tokenizer,
+        model=model,
+        ref_model=ref_model,
+        reward_model=reward_model,
+        train_dataset=train_dataset,
+        value_model=value_model,
+        data_collator=data_collator,
+        eval_dataset=valid_dataset,
+    )
+
+    ppo_trainer.train()
+    ppo_trainer.save_model()
+
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/tutorials/rlhf/gpt2_sst2_ppo/step_3_generation.py

@@ -0,0 +1,77 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import argparse
+import time
+
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
+
+from project_settings import project_path
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--pretrained_model_name_or_path",
+        default=(project_path / "trained_models/gpt2-sst2-ppo/checkpoint-1250"),
+        # default=(project_path / "trained_models/gpt2-sst2-ppo/checkpoint-150"),
+        # default=(project_path / "trained_models/gpt2-sst2-generation-epoch-3/checkpoint-5400"),
+        type=str
+    )
+    parser.add_argument(
+        "--max_new_tokens",
+        default=512, # 8192, 128
+        type=int, help="最大生成长度（注意：并非模型实际长文本能力）"
+    )
+    parser.add_argument("--top_p", default=0.85, type=float, help="nucleus采样阈值（0-1）")
+    parser.add_argument("--temperature", default=0.85, type=float, help="生成温度，控制随机性（0-1，越大越随机）")
+
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = get_args()
+
+    if torch.cuda.is_available():
+        device = "cuda"
+    elif torch.backends.mps.is_available():
+        # device = "mps"
+        device = "cpu"
+    else:
+        device = "cpu"
+    print(f"device: {device}")
+
+    tokenizer = AutoTokenizer.from_pretrained(args.pretrained_model_name_or_path)
+    model = AutoModelForCausalLM.from_pretrained(args.pretrained_model_name_or_path)
+    model = model.eval().to(device)
+
+    tokenized = tokenizer(
+        # "this",
+        "this is ",
+        # "please listen ",
+        # "eldom has a movie",
+        # "thanks to scott 's charismatic",
+        return_tensors="pt"
+    )
+
+    streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=False)
+
+    generated_ids = model.generate(
+        inputs=tokenized["input_ids"], attention_mask=tokenized["attention_mask"],
+        max_new_tokens=args.max_new_tokens, do_sample=True, streamer=streamer,
+        pad_token_id=tokenizer.pad_token_id, eos_token_id=tokenizer.eos_token_id,
+        top_p=args.top_p, temperature=args.temperature, repetition_penalty=3.0,
+        early_stopping=True,
+    )
+    # response = tokenizer.decode(generated_ids[0][len(tokenized["input_ids"][0]):], skip_special_tokens=True)
+    response = tokenizer.decode(generated_ids[0], skip_special_tokens=False)
+    print(response)
+    # print(generated_ids)
+    print(f"count: {generated_ids.shape}")
+
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/tutorials/{dpo/ultrachat-sft/step_2_train_sft_model2.py → rlhf/gpt2_sst2_reward/step_2_train_model.py}

@@ -1,3 +1,5 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
 import argparse
 import os
 from pathlib import Path
@@ -11,26 +13,25 @@ else:
     temp_directory = Path("/root/autodl-tmp/OpenMiniMind/temp")
 
 from datasets import load_dataset
+import numpy as np
 import torch
-from torch.utils.data import DataLoader
-from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, DataCollatorForLanguageModeling
-from transformers.models.llama.modeling_llama import LlamaModel
-from transformers.models.llama.tokenization_llama import LlamaTokenizer
-from trl import SFTTrainer, SFTConfig, DataCollatorForCompletionOnlyLM
+from transformers import (AutoTokenizer, AutoModelForSequenceClassification,
+                          DataCollatorWithPadding,
+                          Trainer, TrainingArguments
+                          )
 
 
 def get_args():
     parser = argparse.ArgumentParser()
     parser.add_argument(
         "--model_name",
-        default=(project_path / "pretrained_models/jingyaogong/MiniMind2").as_posix() if platform.system() in ("Windows", "Darwin") else "jingyaogong/MiniMind2",
-        # default=(project_path / "pretrained_models/Qwen/Qwen2.5-0.5B").as_posix() if platform.system() in ("Windows", "Darwin") else "Qwen/Qwen2.5-0.5B",
+        # default="openai-community/gpt2",
+        default=(project_path / "pretrained_models/openai-community/gpt2").as_posix(),
         type=str
     ),
     parser.add_argument(
         "--dataset_path",
-        # default="HuggingFaceH4/ultrachat_200k",
-        default="miyuki2026/tutorials",
+        default="stanfordnlp/sst2",
         type=str
     ),
     parser.add_argument("--dataset_name", default=None, type=str),
@@ -45,15 +46,13 @@ def get_args():
         default=(temp_directory / "hub_models").as_posix(),
         type=str
     ),
-    parser.add_argument("--dataset_streaming", action="store_true"),
+    parser.add_argument("--dataset_streaming", default=None, type=str),
     parser.add_argument("--valid_dataset_size", default=1000, type=int),
     parser.add_argument("--shuffle_buffer_size", default=5000, type=int),
 
-    parser.add_argument("--max_seq_length", default=2048, type=int)
-
     parser.add_argument(
         "--output_model_dir",
-        default=(project_path / "trained_models/qwen2_5-0_5B-ultrachat-sft").as_posix(),
+        default=(project_path / "trained_models/gpt2-for-sequence-classification-sst2-reward").as_posix(),
         type=str
     ),
     parser.add_argument(
@@ -61,105 +60,110 @@ def get_args():
         default=None if platform.system() in ("Windows", "Darwin") else os.cpu_count() // 2,
         type=int
     ),
+    parser.add_argument(
+        "--device",
+        default=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
+        type=str
+    ),
     args = parser.parse_args()
     return args
 
 
+def format_func(example, tokenizer):
+    sentence: str = example["sentence"]
+    labels: float = float(example["label"])
+    tokenized = tokenizer(sentence)
+    input_ids = tokenized["input_ids"]
+    attention_mask = tokenized["attention_mask"]
+    result = {
+        "input_ids": input_ids,
+        "attention_mask": attention_mask,
+        "labels": labels,
+    }
+    return result
+
+
 def main():
     args = get_args()
 
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-    model = AutoModelForCausalLM.from_pretrained(args.model_name)
-    model = model.to(device)
     tokenizer = AutoTokenizer.from_pretrained(args.model_name)
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
         tokenizer.pad_token_id = tokenizer.eos_token_id
 
+    model = AutoModelForSequenceClassification.from_pretrained(
+        args.model_name,
+        num_labels=1,
+        pad_token_id=tokenizer.pad_token_id
+    )
+    print(f"model.num_labels: {model.num_labels}")
+
     dataset_dict = load_dataset(
         path=args.dataset_path,
         name=args.dataset_name,
-        data_dir="finetome-100k",
-        # split="train_sft",
-        # split="test_sft",
+        split=args.dataset_split,
         cache_dir=args.dataset_cache_dir,
-        # streaming=True,
+        # num_proc=args.num_workers if not args.dataset_streaming else None,
+        streaming=args.dataset_streaming,
     )
-    dataset = dataset_dict["train"]
-
-    if args.dataset_streaming:
-        valid_dataset = dataset.take(args.valid_dataset_size)
-        train_dataset = dataset.skip(args.valid_dataset_size)
-        train_dataset = train_dataset.shuffle(buffer_size=args.shuffle_buffer_size, seed=None)
-    else:
-        dataset = dataset.train_test_split(test_size=args.valid_dataset_size, seed=None)
-        train_dataset = dataset["train"]
-        valid_dataset = dataset["test"]
-    train_dataset = valid_dataset
-
-    response_template = "<|im_end|>\n<|im_start|>assistant"
-    instruction_template = "<|im_end|>\n<|im_start|>user"
-    data_collator = DataCollatorForCompletionOnlyLM(
-        response_template=response_template,
-        instruction_template=instruction_template,
-        tokenizer=tokenizer,
-        mlm=False,
+    train_dataset = dataset_dict["train"]
+    valid_dataset = dataset_dict["validation"]
+    # test_dataset = dataset_dict["test"]
+
+    train_dataset = train_dataset.map(
+        lambda example: format_func(example, tokenizer),
+        batched=False,
+        remove_columns=train_dataset.column_names,
     )
+    valid_dataset = valid_dataset.map(
+        lambda example: format_func(example, tokenizer),
+        batched=False,
+        remove_columns=valid_dataset.column_names,
+    )
+    train_dataset = train_dataset.filter(
+        function=lambda x: len(x["input_ids"]) > 6
+    )
+    valid_dataset = valid_dataset.filter(
+        function=lambda x: len(x["input_ids"]) > 6
+    )
+    data_collator = DataCollatorWithPadding(tokenizer)
 
-    def formatting_prompts_func(examples):
-        # print(examples)
-        formated_text: str = tokenizer.apply_chat_template(
-            conversation=examples["conversations"],
-            tokenize=False,
-            add_generation_prompt=False,
-        )
-        return formated_text
-
-    sft_config = SFTConfig(
+    training_args = TrainingArguments(
         output_dir=args.output_model_dir,
+        # overwrite_output_dir=True,
         num_train_epochs=1,
-        # per_device_train_batch_size=8,
-        # gradient_accumulation_steps=4,
+        per_device_train_batch_size=16,
+        per_device_eval_batch_size=16,
+        eval_strategy="steps",
+        eval_steps=200,
         save_strategy="steps",
-        save_steps=100,
+        save_steps=200,
         save_total_limit=2,
-        logging_steps=100,
-        learning_rate=2e-5,
-        warmup_steps=100,
-        lr_scheduler_type="cosine",
-        fp16=True if torch.cuda.is_available() else False,
-        gradient_checkpointing=False,
-        optim="adamw_torch",
-        report_to="none",
-        max_length=1024,
-        dataset_kwargs=dict(
-            skip_prepare_dataset=True
-        ),
+        logging_steps=200,
+        learning_rate=5e-5,
+        warmup_steps=200,
+        weight_decay=0.01,
+        fp16=torch.cuda.is_available(),
+        dataloader_num_workers=args.num_workers or 0,
+        remove_unused_columns=False,
+        load_best_model_at_end=True,
+        metric_for_best_model="eval_loss",
+        greater_is_better=False,
+        logging_dir=(Path(args.output_model_dir) / "logs").as_posix(),
     )
 
-    trainer = SFTTrainer(
+    trainer = Trainer(
         model=model,
-        args=sft_config,
+        args=training_args,
         data_collator=data_collator,
         train_dataset=train_dataset,
-        formatting_func=formatting_prompts_func,
+        eval_dataset=valid_dataset,
     )
 
-    # 开始训练
-    print("开始训练...")
     trainer.train()
-
-    # 保存模型
-    print(f"保存模型到: {args.output_model_dir}")
     trainer.save_model()
-    tokenizer.save_pretrained(args.output_model_dir)
-
-    print("训练完成!")
-
     return
 
 
 if __name__ == "__main__":
     main()
-

examples/tutorials/rlhf/gpt2_sst2_reward/step_3_test_model.py

@@ -0,0 +1,142 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import argparse
+import os
+from pathlib import Path
+import platform
+from typing import Any, Dict, List, Optional, Union, Tuple
+
+if platform.system() in ("Windows", "Darwin"):
+    from project_settings import project_path, temp_directory
+else:
+    project_path = os.path.abspath("../../../")
+    project_path = Path(project_path)
+    temp_directory = Path("/root/autodl-tmp/OpenMiniMind/temp")
+
+from datasets import load_dataset
+import numpy as np
+import torch
+import torch.nn as nn
+from transformers import (AutoTokenizer, AutoModelForSequenceClassification, GPT2ForSequenceClassification
+                          )
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model_name",
+        default=(project_path / "trained_models/gpt2-for-sequence-classification-sst2-reward").as_posix(),
+        type=str
+    ),
+    parser.add_argument(
+        "--dataset_path",
+        default="stanfordnlp/sst2",
+        type=str
+    ),
+    parser.add_argument("--dataset_name", default=None, type=str),
+    parser.add_argument("--dataset_split", default=None, type=str),
+    parser.add_argument(
+        "--dataset_cache_dir",
+        default=(temp_directory / "hub_datasets").as_posix(),
+        type=str
+    ),
+    parser.add_argument(
+        "--model_cache_dir",
+        default=(temp_directory / "hub_models").as_posix(),
+        type=str
+    ),
+    parser.add_argument("--dataset_streaming", default=None, type=str),
+
+    parser.add_argument(
+        "--num_workers",
+        default=None if platform.system() in ("Windows", "Darwin") else os.cpu_count() // 2,
+        type=int
+    ),
+    parser.add_argument(
+        "--device",
+        default=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
+        type=str
+    ),
+    args = parser.parse_args()
+    return args
+
+
+class RewardModelWrapper:
+    """
+    奖励模型包装器，用于计算生成文本的奖励
+    """
+
+    def __init__(self, reward_model_name, tokenizer, device):
+        self.device = device
+        self.tokenizer = tokenizer
+        # 加载你的GPT2RewardModel或标准模型
+        from transformers import GPT2ForSequenceClassification
+        self.model = GPT2ForSequenceClassification.from_pretrained(reward_model_name).to(device)
+        self.model.eval()
+
+    def get_reward(self, texts: List[str]) -> List[float]:
+        """
+        计算文本的奖励分数（SST-2情感分类）
+        """
+        inputs = self.tokenizer(texts, padding=True, truncation=True, return_tensors="pt").to(self.device)
+        with torch.no_grad():
+            outputs = self.model(**inputs)
+            # SST-2是二分类，取正面情感的logits作为奖励
+            logits = outputs.logits
+            # 如果模型输出是logits，使用softmax获取正面概率
+            probs = torch.softmax(logits, dim=-1)
+            # 假设标签1是正面
+            rewards = probs[:, 1].cpu().tolist()
+        return rewards
+
+
+def main():
+    args = get_args()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    tokenizer = AutoTokenizer.from_pretrained(args.model_name)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+        tokenizer.pad_token_id = tokenizer.eos_token_id
+
+    reward_model = AutoModelForSequenceClassification.from_pretrained(
+        args.model_name,
+    )
+    print(f"reward_model.num_labels: {reward_model.num_labels}")
+
+    dataset_dict = load_dataset(
+        path=args.dataset_path,
+        name=args.dataset_name,
+        split=args.dataset_split,
+        cache_dir=args.dataset_cache_dir,
+        # num_proc=args.num_workers if not args.dataset_streaming else None,
+        streaming=args.dataset_streaming,
+    )
+    # dataset = dataset_dict["train"]
+    dataset = dataset_dict["validation"]
+    # dataset = dataset_dict["test"]
+
+
+    for example in dataset:
+        sentence: str = example["sentence"]
+        score: float = float(example["label"])
+
+        outputs = tokenizer(
+            sentence,
+            return_tensors="pt"
+        )
+        input_ids = outputs["input_ids"]
+
+        with torch.no_grad():
+            rewards = reward_model.forward(input_ids)
+            logits = rewards.logits
+            logits = logits.detach().cpu().numpy()
+            reward = logits[0][0]
+        msg = f"reward: {reward}\nscore: {score}\nsentence: {sentence}\n"
+        print(msg)
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/tutorials/rlhf/gpt2_sst2_reward/step_4_test_model.py

@@ -0,0 +1,127 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import argparse
+import os
+from pathlib import Path
+import platform
+from typing import Any, Dict, List, Optional, Union, Tuple
+
+if platform.system() in ("Windows", "Darwin"):
+    from project_settings import project_path, temp_directory
+else:
+    project_path = os.path.abspath("../../../")
+    project_path = Path(project_path)
+    temp_directory = Path("/root/autodl-tmp/OpenMiniMind/temp")
+
+from datasets import load_dataset
+import numpy as np
+import torch
+import torch.nn as nn
+from transformers import (AutoTokenizer, AutoModelForSequenceClassification
+                          )
+from trl.trainer.utils import get_reward
+from transformers import GPT2ForSequenceClassification
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model_name",
+        default=(project_path / "trained_models/gpt2-for-sequence-classification-sst2-reward").as_posix(),
+        type=str
+    ),
+    parser.add_argument(
+        "--dataset_path",
+        default="stanfordnlp/sst2",
+        type=str
+    ),
+    parser.add_argument("--dataset_name", default=None, type=str),
+    parser.add_argument("--dataset_split", default=None, type=str),
+    parser.add_argument(
+        "--dataset_cache_dir",
+        default=(temp_directory / "hub_datasets").as_posix(),
+        type=str
+    ),
+    parser.add_argument(
+        "--model_cache_dir",
+        default=(temp_directory / "hub_models").as_posix(),
+        type=str
+    ),
+    parser.add_argument("--dataset_streaming", default=None, type=str),
+
+    parser.add_argument(
+        "--num_workers",
+        default=None if platform.system() in ("Windows", "Darwin") else os.cpu_count() // 2,
+        type=int
+    ),
+    parser.add_argument(
+        "--device",
+        default=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
+        type=str
+    ),
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = get_args()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    tokenizer = AutoTokenizer.from_pretrained(args.model_name)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+        tokenizer.pad_token_id = tokenizer.eos_token_id
+
+    reward_model = AutoModelForSequenceClassification.from_pretrained(args.model_name)
+
+    dataset_dict = load_dataset(
+        path=args.dataset_path,
+        name=args.dataset_name,
+        split=args.dataset_split,
+        cache_dir=args.dataset_cache_dir,
+        # num_proc=args.num_workers if not args.dataset_streaming else None,
+        streaming=args.dataset_streaming,
+    )
+    # dataset = dataset_dict["train"]
+    dataset = dataset_dict["validation"]
+    # dataset = dataset_dict["test"]
+
+    count = 0
+    batch_text = list()
+    for example in dataset:
+
+        sentence: str = example["sentence"]
+        # labels: int = int(example["label"])
+        batch_text.append(sentence)
+        count += 1
+        if count >= 4:
+            break
+
+    outputs = tokenizer(
+        batch_text,
+        padding=True,
+        truncation=True,
+        return_tensors="pt"
+    )
+    input_ids = outputs["input_ids"]
+    attention_mask = outputs["attention_mask"]
+
+    # last_token_idx = attention_mask.sum(dim=1) - 1
+    # print(last_token_idx)
+
+    reward_logits, score, sequence_lengths = get_reward(
+        model=reward_model,
+        query_responses=input_ids,
+        pad_token_id=tokenizer.pad_token_id,
+        context_length=0,
+    )
+    print(reward_logits)
+    print(score)
+    print(sequence_lengths)
+
+    return
+
+
+if __name__ == "__main__":
+    main()