mirror of https://github.com/hpcaitech/ColossalAI
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
65 lines
2.3 KiB
65 lines
2.3 KiB
8 months ago
|
import argparse
|
||
|
|
import json
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
parser = argparse.ArgumentParser()
|
||
|
|
parser.add_argument(
|
||
|
|
"--data_source",
|
||
|
|
type=str,
|
||
|
|
required=True,
|
||
|
|
default=None,
|
||
|
|
help="The raw data file",
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"--to_verify_file",
|
||
|
|
type=str,
|
||
|
|
required=True,
|
||
|
|
default=None,
|
||
|
|
help="The file that contains the data to be verified",
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"--data_type",
|
||
|
|
type=str,
|
||
|
|
required=True,
|
||
|
|
default=None,
|
||
|
|
help="The data type",
|
||
|
|
)
|
||
|
|
args = parser.parse_args()
|
||
|
|
|
||
|
|
# Read data
|
||
|
|
data = []
|
||
|
|
with open(args.data_source, "r", encoding="utf8") as f:
|
||
|
|
for line in f.readlines():
|
||
|
|
data.append(json.loads(line))
|
||
|
|
to_verify_data = []
|
||
|
|
with open(args.to_verify_file, "r", encoding="utf8") as f:
|
||
|
|
for line in f.readlines():
|
||
|
|
to_verify_data.append(json.loads(line))
|
||
|
|
|
||
|
|
if args.data_type == "sft":
|
||
|
|
target_lable = [msg["content"].strip() for msg in data[0]["messages"] if msg["from"] == "assistant"]
|
||
|
|
target_negative_label = [msg["content"].strip() for msg in data[0]["messages"] if msg["from"] == "human"]
|
||
|
|
|
||
|
|
# Read to verify file
|
||
|
|
|
||
|
|
to_verify_lable = to_verify_data[0]["labels_decode"]
|
||
|
|
for label in target_lable:
|
||
|
|
assert any([label in s for s in to_verify_lable]), f"Label {label} not in target label {to_verify_lable}"
|
||
|
|
for label in target_negative_label:
|
||
|
|
assert all(
|
||
|
|
[label not in s for s in to_verify_lable]
|
||
|
|
), f"Negative label {label} in target label {to_verify_lable}"
|
||
|
|
elif args.data_type == "dpo":
|
||
|
|
chosen_lable = data[0]["chosen"][0]["content"].strip()
|
||
|
|
rejected_lable = data[0]["rejected"][0]["content"].strip()
|
||
|
|
|
||
|
|
# Read to verify file
|
||
|
|
to_verify_lable_chosen = to_verify_data[0]["chosen_label_decode"]
|
||
|
|
to_verify_lable_rejected = to_verify_data[0]["rejected_label_decode"]
|
||
|
|
assert any(
|
||
|
|
[chosen_lable in s for s in to_verify_lable_chosen]
|
||
|
|
), f"Chosen label {chosen_lable} not in target chosen label {to_verify_lable_chosen}"
|
||
|
|
assert any(
|
||
|
|
[rejected_lable in s for s in to_verify_lable_rejected]
|
||
|
|
), f"Rejected label {rejected_lable} not in target rejected label {to_verify_lable_chosen}"
|