222 lines
9.3 KiB
Python
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})
|
|
|
|
"""
|