wehub-resource-sync
65 lines
1.7 KiB
Python
65 lines
1.7 KiB
Python
import json
|
|
import argparse
|
|
from glob import glob
|
|
import os
|
|
from tqdm import tqdm
|
|
|
|
|
|
def extract_qa(completion: str, remove_suffix: bool = False):
|
|
q_s = completion.find("Created Question:")
|
|
if q_s == -1:
|
|
return "", ""
|
|
|
|
s_s = completion.find("Solution to the Created Question:")
|
|
if s_s == -1:
|
|
return "", ""
|
|
|
|
question = completion[q_s + len("Created Question:"):s_s].strip()
|
|
solution = completion[s_s + len("Solution to the Created Question:"):].strip()
|
|
|
|
if remove_suffix:
|
|
solution = solution.split("The answer is")[0].strip()
|
|
|
|
return question, solution
|
|
|
|
|
|
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()
|
|
cnt = 0
|
|
outputs = []
|
|
for item in tqdm(data):
|
|
question, solution = extract_qa(item["completion"])
|
|
if question == "" or solution == "":
|
|
continue
|
|
if question + solution in q_set:
|
|
cnt += 1
|
|
continue
|
|
q_set.add(question + solution)
|
|
item["question"] = question
|
|
item["solution"] = solution
|
|
item["solution_wo_suffix"] = remove_suffix(solution)
|
|
outputs.append(item)
|
|
|
|
print(len(outputs))
|
|
json.dump(outputs, open(args.output_file, "w"), indent=2)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|