wehub-resource-sync
41 lines
1.1 KiB
Python
41 lines
1.1 KiB
Python
import json
|
|
import argparse
|
|
from glob import glob
|
|
import os
|
|
from tqdm import tqdm
|
|
|
|
|
|
def remove_suffix(solution: str):
|
|
return solution.split("The answer is")[0].strip()
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--input_file", type=str)
|
|
parser.add_argument("--output_file", type=str)
|
|
args = parser.parse_args()
|
|
|
|
if os.path.exists(args.input_file):
|
|
data = [json.loads(line) for line in open(args.input_file).readlines()]
|
|
else:
|
|
data = []
|
|
for file in glob(args.input_file):
|
|
data.extend([json.loads(line) for line in open(file).readlines()])
|
|
|
|
q_set = set()
|
|
outputs = []
|
|
for i, item in tqdm(enumerate(data)):
|
|
q_set.add(item["question"] + item["completion"])
|
|
item["solution"] = item["completion"]
|
|
item["solution_wo_suffix"] = remove_suffix(item["completion"])
|
|
item["id"] = i
|
|
outputs.append(item)
|
|
|
|
print(len(q_set))
|
|
print(len(outputs))
|
|
json.dump(outputs, open(args.output_file, "w"), indent=2)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|