Function main

tools/tsr/minimization.py:185–231

(flags)

def main(flags):
    coverage_dir = os.path.join(flags.report_dir, "coverage_cache")
    mutation_dir = os.path.join(flags.report_dir, "mutation_cache")
    sample_dir = os.path.join(flags.report_dir, "sample_cache")
    os.makedirs(flags.report_dir, exist_ok=True)

    exclude_model: str = flags.model
    if exclude_model.endswith("b"):  # format: model_name + parameter size
        exclude_model = "".join(exclude_model.split("-")[:-1])

    coverage_set_cover = get_coverage_set_cover(
        coverage_dir, exclude_model, flags.dataset
    )
    mutation_set_cover = get_mutation_set_cover(
        mutation_dir, exclude_model, flags.dataset
    )
    sample_set_cover = get_sample_set_cover(
        sample_dir, flags.sample_eval_dir, exclude_model, flags.dataset
    )
    merged_set_cover = merge_set_cover(
        coverage_set_cover, mutation_set_cover, sample_set_cover
    )

    if flags.model != "ALL":
        final_report = dict()
        # Stage 1: Coverage min set cover
        final_report["coverage"] = gen_report(
            coverage_set_cover, flags.sample_eval_dir, flags.model
        )
        # Stage 2: Mutation min set cover
        final_report["mutation"] = gen_report(
            mutation_set_cover, flags.sample_eval_dir, flags.model
        )
        # Stage 3: Sampling min set cover
        final_report["sample"] = gen_report(
            sample_set_cover, flags.sample_eval_dir, flags.model
        )
        # Stage 4: All
        final_report["full"] = gen_report(
            merged_set_cover, flags.sample_eval_dir, flags.model
        )
        with open(
            os.path.join(flags.report_dir, f"report_{flags.model}.json"), "w"
        ) as f:
            json.dump(final_report, f, indent=4)
    else:
        dump_humaneval_plus_mini(merged_set_cover, flags.mini_path)


if __name__ == "__main__":