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

278 lines
12 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 majority_voting_frequency(preds):
assert isinstance(preds, list)
if isinstance(preds[0], list):
tmp = []
for pred in preds:
tmp.append(str(sorted(pred)))
tmp = collections.Counter(tmp)
tmp = sorted(tmp.items(), key=lambda x: x[1], reverse=True)
sorted_preds = [(eval(pred), fre) for pred, fre in tmp]
elif isinstance(preds[0], str):
tmp = collections.Counter(preds)
tmp = sorted(tmp.items(), key=lambda x: x[1], reverse=True)
sorted_preds = [(pred, fre) for pred, fre in tmp]
else:
raise ValueError(f"Unknown type {type(preds[0])}")
return sorted_preds
def extract_sc_mul_labels(response_data, id_field: str):
id2sc_preds = {}
for item in response_data:
if not item["response"]:
continue
if isinstance(item["response"], list):
preds = item["pred"]
else:
preds = [item["pred"]]
assert len(preds) > 1, f"Self-consistency requires at least 2 predictions, but got {len(preds)}."
sorted_preds = majority_voting_frequency(preds)
item_id = item[id_field] if isinstance(item[id_field], str) else str(item[id_field])
id2sc_preds[item_id] = sorted_preds
return id2sc_preds
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_init(id2sc_preds):
global _id2sc_preds
_id2sc_preds = id2sc_preds
def _process_response_worker(item, top_p: float = 0.0):
item_id, resp_id, prefix_id = item["prefix_id"].split("_")
prefix = item["prefix"]
# item_id = int(item_id)
if item_id not in _id2sc_preds:
return item_id, {}
if item["pred"] == "":
return item_id, {}
if item["pred"] == "failed extracting answer from completion":
return item_id, {}
sc_label, sc_freq = _id2sc_preds[item_id][0]
tot_num = sum([fre for _, fre in _id2sc_preds[item_id]])
if sc_freq / tot_num < top_p:
return item_id, {}
v = counting_partial_response_value(item["pred"], sc_label)
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_set = set()
values = []
prefixes = []
for v, prefix in trajectory_pairs:
if prefix in prefix_set:
continue
values.append(parse_value(v, binary))
prefixes.append(prefix)
prefix_set.add(prefix)
outputs["value"] = values
outputs["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("--response_file_for_sc", type=str)
parser.add_argument("--response_id_field", type=str, default="idx")
parser.add_argument("--binary", default=False, action="store_true")
parser.add_argument("--num_workers", type=int, default=8)
parser.add_argument("--top_p", type=float, default=0.0)
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 glob(args.input_file):
if ".metrics" in file:
continue
print(file)
data += json.load(open(file))
print("Collecting response data...")
if os.path.exists(args.response_file_for_sc):
response_data = json.load(open(args.response_file_for_sc))
else:
response_data = []
for file in glob(args.response_file_for_sc):
if ".metrics" in file:
continue
print(file)
response_data += json.load(open(file))
response_data = merge_seed_sampled_data(response_data, args.response_id_field)
last_round_id2sc_preds = extract_sc_mul_labels(response_data, args.response_id_field)
num_candidate_cnt = collections.Counter()
for item_id, preds in last_round_id2sc_preds.items():
num_candidate_cnt[len(preds)] += 1
print(num_candidate_cnt)
del response_data
item_id2partial_trajectories = collections.defaultdict(list)
missing = 0
with Pool(args.num_workers, initializer=_process_response_init, initargs=(last_round_id2sc_preds,)) as pool:
inputs = []
for item in data:
inputs.append(item)
_annotate = partial(_process_response_worker, top_p=args.top_p)
for item_id, result in tqdm(pool.imap_unordered(_annotate, 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)
parent_dir = os.path.dirname(args.output_file)
if not os.path.exists(parent_dir):
os.makedirs(parent_dir, exist_ok=True)
json.dump(outputs, open(args.output_file, "w"), indent=2)
if __name__ == '__main__':
main()
"""
>>> python scripts/math_scale/construct_process_rm_sample_sc.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_by_sc.azure.json \
--response_file_for_sc "../msranlpintern/reward_modeling/experiments/mathstral.mathscale4o.process-dpo.iter0.V100.tp8dp48.v2.2.fix.s42/checkpoint-600/mathscale4o/500k-split-*-of-20/train.500k.de_con.boxed.v1.0.*-of-20.0shot.n10.tem1.0.p0.9.?-of-8.json" \
--response_id_field id --num_workers 128
Counter({3: 889966, 0: 537841, 2: 459487, 1: 397897})
Missing 0 items in the response data.
>>> python scripts/math_scale/construct_process_rm_sample_sc.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.iter0-pdpo-sc.azure.json \
--response_file_for_sc "../msranlpintern/reward_modeling/experiments/mathstral.mathscale4o.process-dpo.iter0.V100.tp8dp48.v2.2.fix.s42/checkpoint-600/mathscale4o/500k-split-*-of-20/train.500k.de_con.boxed.v1.0.*-of-20.0shot.n90.tem1.0.p0.9.json" \
--response_id_field id --num_workers 128
>>> python scripts/math_scale/construct_process_rm_sample_sc.py \
--input_file "${OUTPUT_PREFIX_PATH}/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 ${OUTPUT_PREFIX_PATH}/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.iter0-pdpo-sc-p0.6.azure.json \
--response_file_for_sc "${OUTPUT_PREFIX_PATH}/experiments/mathstral.mathscale4o.process-dpo.iter0.V100.tp8dp48.v2.2.fix.s42/checkpoint-600/mathscale4o/500k-split-*-of-20/train.500k.de_con.boxed.v1.0.*-of-20.0shot.n90.tem1.0.p0.9.json" \
--response_id_field id --num_workers 128 --top_p 0.6
Counter({3: 945267, 2: 408413, 0: 308107, 1: 279042})
Missing 1097360 items in the response data.
>>> python scripts/math_scale/construct_process_rm_sample_sc.py \
--input_file "${MODEL_PREFIX_PATH}/mathstral-7B-v0.1/mathscale4o/split-256/train.500k.de_con.boxed.v1.0.n10.tem1.0.p0.9.upper0.7.r0.3.sample16.filter_same.*-of-4.prefix_completion.n3.tem1.0.p0.9.*-of-256.json" \
--output_file ${MODEL_PREFIX_PATH}/mathstral-7B-v0.1/mathscale4o/train.500k.de_con.boxed.v1.0.n10.tem1.0.p0.9.upper0.7.r0.3.sample16.filter_same.prefix_completion.n3.tem1.0.p0.9.process_rm.sc-10-p0.0.azure.json \
--response_file_for_sc "${MODEL_PREFIX_PATH}/mathstral-7B-v0.1/mathscale4o/500k-split-*-of-20/train.500k.de_con.boxed.v1.0.*-of-20.0shot.n10.tem1.0.p0.9.*-of-64.*json" --response_id_field id --num_workers 128
Counter({3: 1869170, 0: 1528083, 2: 859527, 1: 787261})
Missing 24189 items in the response data.
Counter({3: 1862945, 0: 1521858, 2: 855962, 1: 783862})
>>> python scripts/math_scale/construct_process_rm_sample_sc.py \
--input_file "${OUTPUT_PREFIX_PATH}/experiments/mathstral.mathscale4o.process-dpo.iter0.V100.tp8dp48.v2.2.fix.s42/checkpoint-600/numina/completion-split-128/cot.de_con.n8.tem1.0.p1.0.s0.upper0.7.r0.3.sample32.filter_same.{split_id}-of-4.n3.tem1.0.p1.0.*-of-128.s0.json" \
--output_file ${OUTPUT_PREFIX_PATH}/experiments/mathstral.mathscale4o.process-dpo.iter0.V100.tp8dp48.v2.2.fix.s42/checkpoint-600/numina/cot.de_con.n8.tem1.0.p1.0.s0.upper0.7.r0.3.sample32.filter_same.process_rm.iter0-pdpo-sc-p0.5.v0.0.{split_id}-of-4.azure.json \
--response_file_for_sc "${OUTPUT_PREFIX_PATH}/experiments/mathstral.mathscale4o.process-dpo.iter0.V100.tp8dp48.v2.2.fix.s42/checkpoint-600/numina/830k-split-*-of-10/cot.de_con.*-of-10.n8.tem1.0.p1.0.*-of-*.s*.json" --response_id_field id --num_workers 128 --top_p 0.5
Counter({3: 803168, 2: 454386, 1: 260676, 0: 159454})
Missing 2783152 items in the response data.
Counter({3: 772040, 2: 436348, 1: 250464, 0: 153477})
"""