| 37 | |
| 38 | |
| 39 | class RewardManager: |
| 40 | def __init__(self, tokenizer, num_examine) -> None: |
| 41 | self.tokenizer = tokenizer |
| 42 | self.num_examine = num_examine # the number of batches of decoded responses to print to the console |
| 43 | |
| 44 | def __call__(self, data: DataProto, return_dict: bool = False): |
| 45 | """We will expand this function gradually based on the available datasets""" |
| 46 | |
| 47 | # If there is rm score, we directly return rm score. Otherwise, we compute via rm_score_fn |
| 48 | if "rm_scores" in data.batch.keys(): |
| 49 | return data.batch["rm_scores"] |
| 50 | |
| 51 | reward_tensor = torch.zeros_like(data.batch["responses"], dtype=torch.float32) |
| 52 | |
| 53 | already_print_data_sources = {} |
| 54 | |
| 55 | for i in range(len(data)): |
| 56 | data_item = data[i] # DataProtoItem |
| 57 | |
| 58 | prompt_ids = data_item.batch["prompts"] |
| 59 | |
| 60 | prompt_length = prompt_ids.shape[-1] |
| 61 | |
| 62 | valid_prompt_length = data_item.batch["attention_mask"][:prompt_length].sum() |
| 63 | valid_prompt_ids = prompt_ids[-valid_prompt_length:] |
| 64 | |
| 65 | response_ids = data_item.batch["responses"] |
| 66 | valid_response_length = data_item.batch["attention_mask"][prompt_length:].sum() |
| 67 | valid_response_ids = response_ids[:valid_response_length] |
| 68 | |
| 69 | # decode |
| 70 | sequences = torch.cat((valid_prompt_ids, valid_response_ids)) |
| 71 | sequences_str = self.tokenizer.decode(sequences) |
| 72 | |
| 73 | ground_truth = data_item.non_tensor_batch["reward_model"]["ground_truth"] |
| 74 | |
| 75 | # select rm_score |
| 76 | data_source = data_item.non_tensor_batch["data_source"] |
| 77 | compute_score_fn = _select_rm_score_fn(data_source) |
| 78 | |
| 79 | score = compute_score_fn(solution_str=sequences_str, ground_truth=ground_truth) |
| 80 | reward_tensor[i, valid_response_length - 1] = score |
| 81 | |
| 82 | if data_source not in already_print_data_sources: |
| 83 | already_print_data_sources[data_source] = 0 |
| 84 | |
| 85 | if already_print_data_sources[data_source] < self.num_examine: |
| 86 | already_print_data_sources[data_source] += 1 |
| 87 | print(sequences_str) |
| 88 | |
| 89 | if return_dict: |
| 90 | return {"reward_tensor": reward_tensor} |
| 91 | else: |
| 92 | return reward_tensor |
| 93 | |
| 94 | |
| 95 | @hydra.main(config_path="config", config_name="ppo_trainer_split", version_base=None) |