# Copyright 2025 Sina Weibo Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """ Generate responses given a dataset of prompts """ import csv import ray import numpy as np import hydra import os import time from tabulate import tabulate from collections import Counter os.environ['NCCL_DEBUG'] = 'WARN' os.environ['TOKENIZERS_PARALLELISM'] = 'true' # os.environ['TORCH_COMPILE_DISABLE'] = '1' from verl.utils.model import compute_position_id_with_mask import pandas as pd from transformers import AutoTokenizer from verl import DataProto from verl.utils.fs import copy_local_path_from_hdfs from verl.workers.fsdp_workers import ActorRolloutRefWorker from verl.utils.hdfs_io import makedirs from verl.single_controller.ray import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup from rllm.rewards.rl_reward import rllm_reward_fn from rllm.rewards.math_utils.utils import extract_answer @hydra.main(config_path='config', config_name='generation', version_base=None) def main(config): start_time = time.time() from pprint import pprint from omegaconf import OmegaConf pprint(OmegaConf.to_container(config, resolve=True)) # resolve=True will eval symbol values OmegaConf.resolve(config) local_path = copy_local_path_from_hdfs(config.model.path) from verl.utils import hf_tokenizer tokenizer = hf_tokenizer(local_path) # Check if output file already exists if os.path.exists(config.data.output_path): print(f"Output file {config.data.output_path} already exists. Skipping generation and proceeding to evaluation.") if config.data.output_path.endswith('.parquet'): dataset = pd.read_parquet(config.data.output_path) elif config.data.output_path.endswith('.json'): dataset = pd.read_json(config.data.output_path, orient='records', lines=True) else: if config.rollout.temperature == 0.: assert config.data.n_samples == 1, 'When temperature=0, n_samples must be 1.' # read dataset. Note that the dataset should directly contain chat template format (e.g., a list of dictionary) dataset = pd.read_parquet(config.data.path) chat_lst = dataset[config.data.prompt_key].tolist() chat_lst = [chat.tolist() for chat in chat_lst] tokenizer.padding_side = 'left' if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token ray_cls_with_init = RayClassWithInitArgs(cls=ray.remote(ActorRolloutRefWorker), config=config, role='rollout') resource_pool = RayResourcePool(process_on_nodes=[config.trainer.n_gpus_per_node] * config.trainer.nnodes) wg = RayWorkerGroup(resource_pool=resource_pool, ray_cls_with_init=ray_cls_with_init) wg.init_model() total_samples = len(dataset) config_batch_size = config.data.batch_size dp_size = wg.world_size // config.rollout.tensor_model_parallel_size num_batch = (total_samples + config_batch_size - 1) // config_batch_size output_lst = [] # We'll reshape at the end print('len(dataset):', total_samples) print('wg.worker_names:', wg.worker_names) for batch_idx in range(num_batch): print(f'[{batch_idx+1}/{num_batch}] Start to process.') batch_chat_lst = chat_lst[batch_idx * config_batch_size:(batch_idx + 1) * config_batch_size] # Repeat the batch n_samples times using block repetition inside the batch repeated_chat_lst = [] for chat in batch_chat_lst: repeated_chat_lst.extend([chat] * config.data.n_samples) inputs = tokenizer.apply_chat_template(repeated_chat_lst, add_generation_prompt=True, padding=True, truncation=True, max_length=config.rollout.prompt_length, return_tensors='pt', return_dict=True, tokenize=True) input_ids = inputs['input_ids'] attention_mask = inputs['attention_mask'] position_ids = compute_position_id_with_mask(attention_mask) batch_dict = {'input_ids': input_ids, 'attention_mask': attention_mask, 'position_ids': position_ids} print(f'main_gen.py, input_ids.shape = {input_ids.shape}, content:', input_ids[:, -32:-26]) data = DataProto.from_dict(batch_dict) real_batch_size = data.batch['input_ids'].shape[0] original_dp_size = dp_size dp_size = wg.world_size # ray会先把数据分到每个worker，再每个tp group内收集，所以要保证总数能被worker数整除，该校验不应考虑tp if real_batch_size % dp_size != 0: dummy_data_size = dp_size - real_batch_size % dp_size dummy_data = data[:dummy_data_size] data = DataProto.concat([data, dummy_data]) print( f'dp_size {dp_size} is not divisible by real_batch_size {real_batch_size}, add {dummy_data_size} dummy data' ) dp_size = original_dp_size batch_size = data.batch['input_ids'].shape[0] assert batch_size % dp_size == 0, f'batch_size {batch_size} is not divisible by dp_size {dp_size}' print(f'[{batch_idx+1}/{num_batch}] Start to generate.') # Generate all samples at once print('ZHS batch len:', len(data.batch['input_ids'])) output = wg.generate_sequences(data) # Remove dummy data output = output[:real_batch_size] output_text = tokenizer.batch_decode(output.batch['input_ids'][:, -config.rollout.response_length:], skip_special_tokens=False) # Remove padding pad_token = tokenizer.pad_token output_text_unpad = [] for text in output_text: output_text_unpad.append(text.replace(pad_token, '')) output_lst.extend(output_text_unpad) # Reshape output_lst from (total_samples,) to (n_data, n_samples) total_generated = len(output_lst) n_data = total_generated // config.data.n_samples output_lst = np.array(output_lst).reshape(n_data, config.data.n_samples).tolist() # Add to the data frame dataset['responses'] = output_lst # add correctness field total_lst = compute_correctness(dataset, config.data.data_source_key) dataset['correctness'] = total_lst # Write to a new parquet output_dir = os.path.dirname(config.data.output_path) makedirs(output_dir, exist_ok=True) dataset.to_json(config.data.output_path, orient='records', force_ascii=False, lines=True) if 'correctness' not in dataset: total_lst = compute_correctness(dataset,config.data.data_source_key) dataset['correctness'] = total_lst dataset.to_json(config.data.output_path, orient='records', force_ascii=False, lines=True) print(f"Output file {config.data.output_path} doesn't have correctness field. Have computed each answer's correctness and saved.") output_dir = os.path.dirname(config.data.output_path) # Compute evaluation metrics prompts = dataset[config.data.prompt_key] responses = dataset['responses'] # Using the generated responses data_sources = dataset[config.data.data_source_key] reward_model_data = dataset[config.data.reward_model_key] # 统计长度均值和截断比例 output_lst = [str(r) for responses_this in list(responses) for r in responses_this] # print(output_lst) print(type(output_lst), type(output_lst[0])) unpad_tokenized = tokenizer(output_lst, add_special_tokens=False).input_ids len_response_tokens = [len(tokens) for tokens in unpad_tokenized] len_mean = np.mean(len_response_tokens) cutoff_ratio = sum([l == config.rollout.response_length for l in len_response_tokens]) / len(unpad_tokenized) print('length cutoff ratio:', cutoff_ratio) passes = 0 total = len(dataset) total_scores = [] conses = 0 for i in range(total): response_lst = responses[i] data_source = data_sources[i] prompt = prompts[i] reward_data = reward_model_data[i] reward_fn = select_reward_fn(data_source) ground_truth = reward_data['ground_truth'] score_lst = [] for r in response_lst: score = reward_fn(data_source, r, ground_truth) score_lst.append(score) max_score = np.max(score_lst) total_scores.append(score_lst) if max_score == 1: passes += 1 extracted_lst = [extract_answer(r) for r in response_lst] extracted_lst = [r for r in extracted_lst if r is not None] cons_answers = find_mode(extracted_lst) cons_response_lst = [r for r in response_lst if extract_answer(r) in cons_answers] is_cons_correct_list = list() for r in cons_response_lst: score = reward_fn(data_source, r, ground_truth) is_cons_correct_list.append(score) if any(is_cons_correct_list): conses += np.mean(is_cons_correct_list) n_samples = config.data.n_samples pass_at_n = passes / total pass_at_1 = np.mean(total_scores) cons_at_n = conses / total spent_time = time.time() - start_time spent_hours = spent_time / 60 / 60 # Save metrics to CSV csv_path = os.path.join(output_dir, f'pass_{spent_hours:.2f}h.csv') # Prepare the row data # Extract the dataset name from the path dataset_name = os.path.basename(config.data.path) row_data = { 'model_path': config.model.path, 'dataset': dataset_name, 'pass@1': pass_at_1, f'pass@{n_samples}': pass_at_n, f'cons@{n_samples}': cons_at_n, 'cutoff_raio': cutoff_ratio, 'mean_response_tokens': len_mean, 'run_hours': spent_hours } # Check if file exists file_exists = os.path.isfile(csv_path) # Write to CSV with open(csv_path, mode='a', newline='') as f: # 追加写，不会覆盖，所以没事 writer = csv.DictWriter(f, fieldnames=row_data.keys()) if not file_exists: writer.writeheader() writer.writerow(row_data) # Convert the row data into a list of lists format for tabulate table_data = [[k, v] for k, v in row_data.items()] # Print table print(tabulate(table_data, headers=['Metric', 'Value'], tablefmt='grid')) def compute_correctness(dataset, data_source_key): total_lst = list() for i in range(len(dataset)): row = dataset.iloc[i] prompt = row['prompt'] gt = row['reward_model']['ground_truth'] # print(gt) responses_this = row['responses'] true_false = [int(rllm_reward_fn(row[data_source_key], response, gt)) for response in responses_this] total_lst.append(true_false) return total_lst def find_mode(lst): if len(lst) == 0: return list() counter = Counter(lst) max_count = max(counter.values()) mode = [k for k, v in counter.items() if v == max_count] return mode # Add the select_reward_fn from main_eval.py def select_reward_fn(data_source): if data_source == 'lighteval/MATH': from verl.utils.reward_score import math return math.compute_score else: from rllm.rewards.rl_reward import rllm_reward_fn return rllm_reward_fn if __name__ == '__main__': main()