Files
microsoft--unilm/PFPO/scripts/math_scale/construct_process_rm_sample_gd.py
T
2026-07-13 13:24:13 +08:00

222 lines
9.3 KiB
Python

import json
import argparse
from glob import glob
import os
import sys
import collections
from multiprocessing import Pool
from tqdm import tqdm
from functools import partial
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
from data.mathscale.util import mathscale_is_equiv
def counting_partial_response_value(preds, label):
assert isinstance(preds, list), preds
v = 0
for pred in preds:
res = mathscale_is_equiv(pred, label)[0]
if res:
v += 1
return v
def parse_value(v, binary: bool):
if binary:
return 1 if v > 0 else 0
return v
def _process_response_worker(item):
item_id, resp_id, prefix_id = item["prefix_id"].split("_")
prefix = item["prefix"]
if "mscale-v0.1" in item_id:
pass
elif "numina" not in item_id:
item_id = int(item_id)
if not item["label"]:
return item_id, {}
if item["pred"] == "":
return item_id, {}
if item["pred"] == "failed extracting answer from completion":
return item_id, {}
v = len([1 for res in item["res"] if res])
return item_id, {"v": v, "prefix": prefix}
def _process_trajectories_worker(item, binary: bool):
item_id, trajectories = item
outputs = {
"idx": item_id
}
trajectory_pairs = [(traj["v"], traj["prefix"]) for traj in trajectories]
prefix_vis = set()
values = []
prefixes = []
for v, prefix in trajectory_pairs:
if prefix in prefix_vis:
continue
values.append(parse_value(v, binary))
prefixes.append(prefix)
prefix_vis.add(prefix)
outputs[f"value"] = values
outputs[f"prefix"] = prefixes
return outputs
def merge_key(item, value):
assert isinstance(item, list)
if isinstance(value, list):
item = item + value
else:
item.append(value)
return item
def merge_seed_sampled_data(data, id_field="id"):
id2data = {}
for item in data:
if item[id_field] not in id2data:
id2data[item[id_field]] = item
continue
tmp = id2data[item[id_field]]
if isinstance(tmp["response"], str):
tmp["response"] = [tmp["response"]]
if not isinstance(tmp["res"], list):
tmp["res"] = [tmp["res"]]
if not isinstance(tmp["pred"], list):
tmp["pred"] = [tmp["pred"]]
tmp["response"] = merge_key(tmp["response"], item["response"])
tmp["res"] = merge_key(tmp["res"], item["res"])
tmp["pred"] = merge_key(tmp["pred"], item["pred"])
assert isinstance(tmp["pred"], list), tmp["pred"]
id2data[item[id_field]] = tmp
return list(id2data.values())
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--input_file", type=str)
parser.add_argument("--output_file", type=str)
parser.add_argument("--binary", default=False, action="store_true")
parser.add_argument("--num_workers", type=int, default=8)
args = parser.parse_args()
print("Collecting data...")
if os.path.exists(args.input_file):
data = json.load(open(args.input_file))
else:
data = []
for file in sorted(glob(args.input_file)):
if ".metrics" in file:
continue
print(file)
try:
sub_data = json.load(open(file))
except:
print(f"Warning: {file}l")
sub_data = [json.loads(line) for line in open(f"{file}l").readlines()]
data += sub_data
data = merge_seed_sampled_data(data)
item_id2partial_trajectories = collections.defaultdict(list)
missing = 0
with Pool(args.num_workers) as pool:
inputs = []
for item in data:
inputs.append(item)
for item_id, result in tqdm(pool.imap_unordered(_process_response_worker, inputs), total=len(inputs)):
if len(result) == 0:
missing += 1
continue
item_id2partial_trajectories[item_id].append(result)
shoot_cnt = collections.Counter()
for item_id, trajectories in item_id2partial_trajectories.items():
for traj in trajectories:
shoot_cnt[traj["v"]] += 1
print(shoot_cnt)
print(f"Missing {missing} items in the response data.")
outputs = []
cnt = collections.Counter()
with Pool(args.num_workers) as pool:
inputs = [(item_id, trajectories) for item_id, trajectories in item_id2partial_trajectories.items()]
_annotate = partial(_process_trajectories_worker, binary=args.binary)
for result in tqdm(pool.imap_unordered(_annotate, inputs), total=len(inputs)):
outputs.append(result)
cnt.update(result["value"])
print(cnt)
json.dump(outputs, open(args.output_file, "w"), indent=2)
if __name__ == '__main__':
main()
"""
>>> python scripts/math_scale/construct_process_rm_sample_gd.py \
--input_file "../msranlpintern/reward_modeling/experiments/mathstral.mathscale4o.sft.V100.tp2dp8.v2.0.s42/checkpoint-800/mathscale4o/split-512/train.500k.de_con.boxed.v1.0.n10.tem1.0.p0.9.upper0.7.r0.3.sample10.filter_same.prefix_completion.n3.tem1.0.p0.9.*-of-512.json" \
--output_file ../msranlpintern/reward_modeling/experiments/mathstral.mathscale4o.sft.V100.tp2dp8.v2.0.s42/checkpoint-800/mathscale4o/train.500k.de_con.boxed.v1.0.n10.tem1.0.p0.9.upper0.7.r0.3.sample10.filter_same.prefix_completion.n3.tem1.0.p0.9.process_rm_gd.binary.local.json \
--binary --num_workers 24
Counter({0: 1243140, 3: 916473, 2: 442815, 1: 435687})
Missing 74 items in the response data.
Counter({1: 1793986, 0: 1242775})
>>> python scripts/math_scale/construct_process_rm_sample_gd.py \
--input_file "../msranlpintern/share/models/llama3.1_8b_mathscale4o/model_lr1e-5_batch512_epochs3_gpus8_linearSchedule/mathscale4o/split-512/train.500k.de_con.boxed.v1.0.n10.tem1.0.p0.9.upper0.7.r0.3.sample10.filter_same.prefix_completion.n3.tem1.0.p0.9.*-of-512.json" \
--output_file ../msranlpintern/share/models/llama3.1_8b_mathscale4o/model_lr1e-5_batch512_epochs3_gpus8_linearSchedule/mathscale4o/train.500k.de_con.boxed.v1.0.n10.tem1.0.p0.9.upper0.7.r0.3.sample10.filter_same.prefix_completion.n3.tem1.0.p0.9.process_rm_gd.binary.local.json \
--binary --num_workers 24
Counter({0: 1103040, 3: 908168, 2: 435088, 1: 427143})
Missing 0 items in the response data.
Counter({1: 1769112, 0: 1102585})
>>> python scripts/math_scale/construct_process_rm_sample_gd.py \
--input_file "../msranlpintern/reward_modeling/experiments/mathstral.mathscale4o.sft.V100.tp2dp8.v2.0.s42/checkpoint-800/mathscale4o/split-512/train.500k.de_con.boxed.v1.0.n10.tem1.0.p0.9.upper0.7.r0.3.sample10.filter_same.prefix_completion.n3.tem1.0.p0.9.*-of-512.json" \
--output_file ../msranlpintern/reward_modeling/experiments/mathstral.mathscale4o.sft.V100.tp2dp8.v2.0.s42/checkpoint-800/mathscale4o/train.500k.de_con.boxed.v1.0.n10.tem1.0.p0.9.upper0.7.r0.3.sample10.filter_same.prefix_completion.n3.tem1.0.p0.9.process_rm_gd.local.json \
--num_workers 24
Counter({0: 1243140, 3: 916473, 2: 442815, 1: 435687})
Missing 74 items in the response data.
Counter({0: 1242779, 3: 915927, 2: 442573, 1: 435482})
>>> python scripts/math_scale/construct_process_rm_sample_gd.py \
--input_file "../msranlpintern/share/models/llama3.1_8b_mathscale4o/model_lr1e-5_batch512_epochs3_gpus8_linearSchedule/mathscale4o/split-512/train.500k.de_con.boxed.v1.0.n10.tem1.0.p0.9.upper0.7.r0.3.sample10.filter_same.prefix_completion.n3.tem1.0.p0.9.*-of-512.json" \
--output_file ../msranlpintern/share/models/llama3.1_8b_mathscale4o/model_lr1e-5_batch512_epochs3_gpus8_linearSchedule/mathscale4o/train.500k.de_con.boxed.v1.0.n10.tem1.0.p0.9.upper0.7.r0.3.sample10.filter_same.prefix_completion.n3.tem1.0.p0.9.process_rm_gd.local.json \
--num_workers 24
>>> python scripts/math_scale/construct_process_rm_sample_gd.py \
--input_file "../msranlpintern/reward_modeling/experiments/mathstral.mathscale4o.process-dpo.iter0.V100.tp8dp48.v2.2.fix.s42/checkpoint-600/mathscale4o/split-512/train.500k.de_con.boxed.v1.0.n10.tem1.0.p0.9.upper0.7.r0.3.sample10.filter_same.prefix_completion.n3.tem1.0.p0.9.*-of-512.json" \
--output_file ../msranlpintern/reward_modeling/experiments/mathstral.mathscale4o.process-dpo.iter0.V100.tp8dp48.v2.2.fix.s42/checkpoint-600/mathscale4o/train.500k.de_con.boxed.v1.0.n10.tem1.0.p0.9.upper0.7.r0.3.sample10.filter_same.prefix_completion.n3.tem1.0.p0.9.process_rm_gd.local.json \
--num_workers 24
Counter({0: 1013210, 3: 660151, 2: 306372, 1: 305458})
Missing 0 items in the response data.
Counter({0: 1012382, 3: 659126, 2: 305949, 1: 305153})
>>> python scripts/math_scale/construct_process_rm_sample_gd.py \
--input_file "../msranlpintern/reward_modeling/experiments/llama3.1.8b.mathscale4o.process-dpo.iter0.A100.dp8.v2.2.s42/checkpoint-1200/mathscale4o/split-512/train.500k.de_con.boxed.v1.0.n10.tem1.0.p0.9.upper0.7.r0.3.sample10.filter_same.prefix_completion.n3.tem1.0.p0.9.*-of-512.json" \
--output_file ../msranlpintern/reward_modeling/experiments/llama3.1.8b.mathscale4o.process-dpo.iter0.A100.dp8.v2.2.s42/checkpoint-1200/mathscale4o/train.500k.de_con.boxed.v1.0.n10.tem1.0.p0.9.upper0.7.r0.3.sample10.filter_same.prefix_completion.n3.tem1.0.p0.9.process_rm_gd.local.json \
--num_workers 24
Counter({0: 950765, 3: 748413, 2: 353782, 1: 336295})
Missing 0 items in the response data.
Counter({0: 950541, 3: 747952, 2: 353608, 1: 336155})
"""