diff --git a/tests/test_reward_utils.py b/tests/test_reward_utils.py new file mode 100644 index 0000000..45dfbdf --- /dev/null +++ b/tests/test_reward_utils.py @@ -0,0 +1,43 @@ +import unittest + +from trainer.reward_utils import rep_penalty + + +class StubTokenizer: + def __init__(self, token_ids): + self.token_ids = token_ids + self.calls = [] + + def encode(self, text, add_special_tokens=True): + self.calls.append((text, add_special_tokens)) + return self.token_ids + + +class RepetitionPenaltyTests(unittest.TestCase): + def test_repeated_chinese_tokens_receive_penalty(self): + tokenizer = StubTokenizer([1, 2, 1, 2, 1, 2, 1, 2, 1, 2]) + + self.assertEqual(rep_penalty("你好你好你好你好你好", tokenizer), 0.5) + self.assertEqual(tokenizer.calls, [("你好你好你好你好你好", False)]) + + def test_repeated_english_tokens_receive_penalty_case_insensitively(self): + tokenizer = StubTokenizer([3, 3, 3, 3]) + + self.assertEqual(rep_penalty("Hello HELLO hello hello", tokenizer), 0.5) + self.assertEqual(tokenizer.calls, [("hello hello hello hello", False)]) + + def test_unique_and_short_sequences_have_no_penalty(self): + self.assertEqual(rep_penalty("unique", StubTokenizer([1, 2, 3, 4, 5])), 0.0) + self.assertEqual(rep_penalty("short", StubTokenizer([1, 2])), 0.0) + + def test_invalid_configuration_is_rejected(self): + tokenizer = StubTokenizer([]) + + with self.assertRaisesRegex(ValueError, "n must be greater than zero"): + rep_penalty("text", tokenizer, n=0) + with self.assertRaisesRegex(ValueError, "cap must be non-negative"): + rep_penalty("text", tokenizer, cap=-0.1) + + +if __name__ == "__main__": + unittest.main() diff --git a/trainer/reward_utils.py b/trainer/reward_utils.py new file mode 100644 index 0000000..a41007d --- /dev/null +++ b/trainer/reward_utils.py @@ -0,0 +1,19 @@ +def rep_penalty(text, tokenizer, n=3, cap=0.5): + r"""Return a capped repetition penalty over model-token n-grams. + + Token IDs keep the calculation consistent across languages. In particular, + contiguous CJK text must not be treated as one word, which is what happens + with a ``\w+``-based regular expression. + """ + if n <= 0: + raise ValueError("n must be greater than zero") + if cap < 0: + raise ValueError("cap must be non-negative") + + token_ids = tokenizer.encode(text.lower(), add_special_tokens=False) + grams = [tuple(token_ids[i:i + n]) for i in range(len(token_ids) - n + 1)] + if not grams: + return 0.0 + + repeated = len(grams) - len(set(grams)) + return min(cap, repeated * cap * 2 / len(grams)) diff --git a/trainer/train_agent.py b/trainer/train_agent.py index 261b00c..6d83dfa 100644 --- a/trainer/train_agent.py +++ b/trainer/train_agent.py @@ -26,16 +26,12 @@ from model.model_minimind import MiniMindConfig, MiniMindForCausalLM from dataset.lm_dataset import AgentRLDataset from trainer.trainer_utils import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model, LMForRewardModel from trainer.rollout_engine import create_rollout_engine, compute_per_token_logps +from trainer.reward_utils import rep_penalty warnings.filterwarnings('ignore') # ================================ 工具与 Reward = Start ================================ -def rep_penalty(text, n=3, cap=0.5): - toks = re.findall(r"\w+|[^\w\s]", text.lower()) - grams = [tuple(toks[i:i + n]) for i in range(len(toks) - n + 1)] - return min(cap, (len(grams) - len(set(grams))) * cap * 2 / len(grams)) if grams else 0.0 - # ======== 工具定义 ======== TOOLS = [ {"type": "function", "function": {"name": "calculate_math", "description": "计算数学表达式", "parameters": {"type": "object", "properties": {"expression": {"type": "string"}}, "required": ["expression"]}}}, @@ -220,7 +216,7 @@ def calculate_rewards(prompts, completions, gt_batch, tools_batch, num_gen, rewa messages = [{"role": role, "content": content.strip()} for role, content in matches] score = reward_model.get_score(messages, answer) reward += score # RM分 - reward -= rep_penalty(answer) + reward -= rep_penalty(answer, tokenizer) rewards[idx] = max(min(reward, 3.0), -3.0) # 总分Clip # -------- 有工具调用：执行结果奖励 -------- else: @@ -240,7 +236,7 @@ def calculate_rewards(prompts, completions, gt_batch, tools_batch, num_gen, rewa verified = validate_gt_in_text(final_text, gt) if gt else set() if gt: reward += 2.5 * len(verified) / len(gt) # GT分 if unfinished: reward -= 0.5 # 未完成扣分 - reward -= rep_penalty(final_text if final_text else answer) + reward -= rep_penalty(final_text if final_text else answer, tokenizer) rewards[idx] = max(min(reward, 3.0), -3.0) # 总分Clip return rewards diff --git a/trainer/train_grpo.py b/trainer/train_grpo.py index 76ea045..06f2399 100755 --- a/trainer/train_grpo.py +++ b/trainer/train_grpo.py @@ -24,16 +24,11 @@ from model.model_minimind import MiniMindConfig, MiniMindForCausalLM from dataset.lm_dataset import RLAIFDataset from trainer.trainer_utils import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model, LMForRewardModel from trainer.rollout_engine import create_rollout_engine +from trainer.reward_utils import rep_penalty warnings.filterwarnings('ignore') -def rep_penalty(text, n=3, cap=0.5): - toks = re.findall(r"\w+|[^\w\s]", text.lower()) - grams = [tuple(toks[i:i + n]) for i in range(len(toks) - n + 1)] - return min(cap, (len(grams) - len(set(grams))) * cap * 2 / len(grams)) if grams else 0.0 - - def calculate_rewards(prompts, responses, reward_model): rewards = torch.zeros(len(responses), device=args.device) @@ -57,7 +52,7 @@ def calculate_rewards(prompts, responses, reward_model): rewards[response_idx] += 1.0 if 20 <= len(thinking_content.strip()) <= 300 else -0.5 rewards[response_idx] += 0.25 if response.count('') == 1 else -0.25 answer = answer_content.strip() - rewards[response_idx] -= rep_penalty(answer) + rewards[response_idx] -= rep_penalty(answer, tokenizer) score = reward_model.get_score(messages, answer) reward_model_scores.append(score) @@ -332,4 +327,4 @@ if __name__ == "__main__": # ========== 9. 清理分布进程 ========== if dist.is_initialized(): dist.barrier() - dist.destroy_process_group() \ No newline at end of file + dist.destroy_process_group() diff --git a/trainer/train_ppo.py b/trainer/train_ppo.py index 07c7738..88fc8b3 100644 --- a/trainer/train_ppo.py +++ b/trainer/train_ppo.py @@ -22,16 +22,11 @@ from model.model_minimind import MiniMindConfig, MiniMindForCausalLM from dataset.lm_dataset import RLAIFDataset from trainer.trainer_utils import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model, LMForRewardModel from trainer.rollout_engine import create_rollout_engine +from trainer.reward_utils import rep_penalty warnings.filterwarnings('ignore') -def rep_penalty(text, n=3, cap=0.5): - toks = re.findall(r"\w+|[^\w\s]", text.lower()) - grams = [tuple(toks[i:i + n]) for i in range(len(toks) - n + 1)] - return min(cap, (len(grams) - len(set(grams))) * cap * 2 / len(grams)) if grams else 0.0 - - # 自定义的Critic模型，继承自MiniMindLM class CriticModel(MiniMindForCausalLM): def __init__(self, params): @@ -64,7 +59,7 @@ def calculate_rewards(prompts, responses, reward_model): rewards[i] += 1.0 if 20 <= len(thinking_content.strip()) <= 300 else -0.5 rewards[i] += 0.25 if response.count('') == 1 else -0.25 answer = answer_content.strip() - rewards[i] -= rep_penalty(answer) + rewards[i] -= rep_penalty(answer, tokenizer) score = reward_model.get_score(messages, answer) reward_model_scores.append(score) @@ -449,4 +444,4 @@ if __name__ == "__main__": # ========== 9. 清理分布进程 ========== if dist.is_initialized(): dist.barrier() - dist.destroy_process_group() \ No newline at end of file + dist.destroy_process_group()