def main():
    """
    This script takes the completions from GPT4, the corresponding inputs for worsening code, and the original dataset, to construct a new combined dataset,
    which contains the solutions from the policy model, as well as the synthesized negative codes for each plausible solution.
    :return:
    """
    parser = argparse.ArgumentParser()
    parser.add_argument("--worsen_file", type=str,
                        help="The file contains completion from the teacher model for worsening code"
                             "The inputs for this file are generated by `pp_worsen_inputs.py` script.")
    parser.add_argument("--completion_file", type=str,
                        help="The file contains the completion for each query.")
    # parser.add_argument("--completion_response_field", type=str, default="completion")
    parser.add_argument("--completion_problem_id_field", type=str, default="problem_id")
    parser.add_argument("--output_file", type=str)
    args = parser.parse_args()

    worsen_codes = [json.loads(line) for line in open(args.worsen_file).readlines()]
    print(f"Total number of worsen codes: {len(worsen_codes)}")

    if os.path.exists(args.completion_file):
        data = json.load(open(args.completion_file))
    else:
        data = []
        for file in glob(args.completion_file):
            data += json.load(open(file))

    p_id2response = {item[args.completion_problem_id_field]: item for item in data}

    p_id2pairs = defaultdict(list)
    for item in worsen_codes:
        if item["completion"]:
            try:
                completion = json.loads(item["completion"])
            except:
                print(f"Json parsing error: {item['completion']}")
                continue
            neg_code = completion["incorrect_program"]
            # if isinstance(neg_code, dict):
            #     print(json.dumps(completion, indent=2))

            if not isinstance(neg_code, str):
                print(f"Bad format: {neg_code}.")
                continue

            if not neg_code.strip():
                continue

            p_id, pred_id = item["id"].split("_")
            pred_id = int(pred_id[3:])
            p_id = int(p_id)

            response = p_id2response[p_id]
            pred = response["pred"][pred_id]
            assert pred in item["prompt"]

            p_id2pairs[p_id].append((pred, neg_code))

    outputs = []
    num_pairs = 0
    for p_id, pairs in p_id2pairs.items():
        pos = []
        neg = []
        for pred, neg_code in pairs:
            pos.append(pred)
            neg.append(neg_code)
        assert len(pos) == len(neg)
        num_pairs += len(pos)
        outputs.append({
            "problem_id": p_id,
            "pos": pos,
            "neg": neg,
        })

    print(f"Total number of outputs: {len(outputs)}")
    print(f"Total number of pairs: {num_pairs}")
    json.dump(outputs, open(args.output_file, "w"), indent=2, ensure_ascii=False)

    This script takes the completions from GPT4, the corresponding inputs for worsening code, and the original dataset, to construct a new combined dataset,
    which contains the solutions from the policy model, as well as the synthesized negative codes for each plausible solution.
    :return:
    
PFPO/scripts/apps/worsen_gpt4_combine.py