Fix inference code

Author: RyanLiu112

src/prm_evaluation/genprm_inference.py

@@ -65,10 +65,14 @@ def execute(self, text):
 
 
 class GenPRM:
-    def __init__(self, model_path):
+    def __init__(self, model_path, tensor_parallel_size):
         # Load the model and tokenizer
         timestamped_print(f"Loading model from {model_path}", level="INFO")
-        self.model = LLM(model=model_path, gpu_memory_utilization=0.90, enable_chunked_prefill=True)
+        self.model = LLM(
+            model=model_path,
+            tensor_parallel_size=tensor_parallel_size,
+            enable_chunked_prefill=True
+        )
         self.tokenizer = AutoTokenizer.from_pretrained(model_path)
         timestamped_print(f"GenPRM loaded successfully", level="INFO")
 
@@ -301,16 +305,16 @@ def _single_inference(
             cur_time += 1
             new_prompts = []
             if output2.text.endswith('</output>\n'):
-                output2.text = cur_prompt + output2.text
+                output2.text = cur_prompts[0] + output2.text
                 out_nodes.append(output2)
             else:
                 if execute:
                     # execute the code
-                    code_output = code_executor.execute(cur_prompt + output2.text)
+                    code_output = code_executor.execute(cur_prompts[0] + output2.text)
                     code_content = f"[Code Output]\n\n```\n{code_output}\n```\n"
-                    new_prompts.append(cur_prompt + output2.text + code_content)
+                    new_prompts.append(cur_prompts[0] + output2.text + code_content)
                 else:
-                    new_prompts.append(cur_prompt + output2.text + '[Code Output]\n\n```\n')
+                    new_prompts.append(cur_prompts[0] + output2.text + '[Code Output]\n\n```\n')
 
             cur_prompts = new_prompts
 

src/prm_evaluation/prm_evaluate.py

@@ -7,7 +7,7 @@
 root_dir = os.path.abspath(os.path.join(current_dir, ".."))
 sys.path.append(root_dir)
 import argparse
-import os
+import json
 import random
 import time
 import threading
@@ -61,6 +61,9 @@ def parse_args():
     parser.add_argument("--analyze_template", type=str, default="<analyze>\nLet's analyze the Paragraph {cur_step} step by step: ")
     parser.add_argument("--verify_template", type=str, default="<verify>\nLet's use python code to find any potential error:\n```python\n")
     parser.add_argument("--output_template", type=str, default="<output>\n**Judgement**: $\\boxed")
+    parser.add_argument("--tensor_parallel_size", type=int, default=1)
+    parser.add_argument("--idd", type=int, default=1)
+
     return parser.parse_args()
 
 
@@ -73,7 +76,7 @@ def parse_args():
 
 #####################################################           model load with VLLM             ########################################################
 
-genprm = GenPRM(args.reward_name_or_path)
+genprm = GenPRM(args.reward_name_or_path, args.tensor_parallel_size)
 
 #####################################################           load splited dataset             ########################################################
 
@@ -91,6 +94,7 @@ def get_shuffled_folders(directory):
 for data_path in target_list:
     folder_name = os.path.basename(data_path)
     save_path = os.path.join(args.split_out, folder_name)
+
     if args.analyze:
         save_path += '_analyze'
     if args.verify:
@@ -125,11 +129,9 @@ def get_shuffled_folders(directory):
     thread.start()
     timestamped_print("Heartbeat thread started. Main thread continues...")
 
-    data = load_from_disk(os.path.join(args.data_path, folder_name))
-    timestamped_print(data)
-    data_new = data.to_list()
-
-    sample = deepcopy(data_new)[0]
+    with open(os.path.join(args.data_path, folder_name, 'sample.json'), 'r') as f:
+        data_new = json.load(f)
+    sample = deepcopy(data_new)
     data_input = sample['steps']
     data_input[0] = sample['problem'] + '\n' + data_input[0]
     if data_input and data_input[-1] == '':
@@ -143,11 +145,11 @@ def get_shuffled_folders(directory):
     else:
         message = {
             'conversation': [
-                {'role': 'system', 'content': 'You are a math teacher. Your task is to review and critique the paragraphs in solution directly. Output your judgement in the format of `boxed{Yes}` if the paragraph is correct, or `boxed{No}` if the paragraph is incorrect.'}
+                {'role': 'system', 'content': 'You are a math teacher. Your task is to review and critique the paragraphs in solution directly. Output your judgement in the format of `\\boxed{Yes}` if the paragraph is correct, or `\\boxed{No}` if the paragraph is incorrect.'}
             ]
         }
     for j1 in range(len(data_input)):
-        line = {'content': data_input[j1], 'role': 'user'}
+        line = {'role': 'user', 'content': data_input[j1]}
         message['conversation'].append(line)
         line = {'content': '', 'role': 'assistant'}
         message['conversation'].append(line)
@@ -192,12 +194,13 @@ def get_shuffled_folders(directory):
             step_scores.append(reward)
 
         end = time.perf_counter()
-        data_new[0]['time'] = end - start
-        data_new[0]['value'] = step_scores
-        data_new[0]['conversation'] = conversation
+        data_new['time'] = end - start
+        data_new['value'] = step_scores
+        data_new['conversation'] = conversation
+
         timestamped_print(type(data_new))
-        timestamped_print(type(Dataset.from_list(data_new)))
-        (Dataset.from_list(data_new)).save_to_disk(save_path)
+        with open(os.path.join(save_path, f'result_{args.idd}.json'), 'w') as f:
+            json.dump(data_new, f, indent=4)
         timestamped_print(f"dataset has been saved to: {save_path}")
     except Exception as e:
         traceback.print_exc()

src/utils/split_dataset.py

@@ -8,9 +8,10 @@
     --split_dir _data/split_input/ProcessBench
 """
 
-from datasets import load_dataset, Dataset
 import argparse
+import json
 import os
+from datasets import load_dataset, Dataset
 
 
 def export_all_splits(dataset_name, root_output_dir):
@@ -43,17 +44,11 @@ def process_split(dataset, split_name, root_output_dir):
     for idx, example in enumerate(dataset):
         example_dir = os.path.join(
             split_dir,
-            f"{split_name}_example_{idx:05d}"
+            f"{split_name}_{idx:03d}"
         )
-        create_single_example_dataset(example, dataset.info, example_dir)
-
-
-def create_single_example_dataset(example, info, output_dir):
-    """Create self-contained dataset for one example"""
-    os.makedirs(output_dir, exist_ok=True)
-    single_ds = Dataset.from_list([example])
-    # single_ds.info = info
-    single_ds.save_to_disk(output_dir)
+        os.makedirs(example_dir, exist_ok=True)
+        with open(os.path.join(example_dir, "sample.json"), "w") as f:
+            json.dump(example, f, ensure_ascii=False, indent=4)
 
 
 if __name__ == "__main__":