278 lines
12 KiB
Python
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})
|
|
"""
|