content: import argparse
  import json
  import os
  import pickle
  from concurrent.futures import ProcessPoolExecutor, as_completed
  from copy import deepcopy
  from typing import Any, Dict, List, Optional, Tuple
  
  from rich.progress import track
  
  from evalplus.data import write_jsonl
  from tools.tsr.coverage_init import collect_coverage_info
  from tools.tsr.mutation_init import collect_mutation_info
  from tools.tsr.sample_init import collect_sample_info
  from tools.tsr.utils import get_problems, get_task_ids, to_path
  
  
  def global_util_init(dataset: str):
      global problems
      global task_ids
      global problem_count
      problems = get_problems(dataset)
      task_ids = get_task_ids(dataset)
      problem_count = len(problems)
  
  
  ###########################
  # Greedy Min Set Covering #
  ###########################
  
  
  def merge_set_cover(*args) -> Dict[str, List[str]]:
      merged_set_cover = {task_id: [] for task_id in task_ids}
      for set_cover_dict in args:
          for task_id, plus_tests in set_cover_dict.items():
              for plus_test in plus_tests:
                  if plus_test not in merged_set_cover[task_id]:
                      merged_set_cover[task_id].append(plus_test)
      return merged_set_cover
  
  
  def greedy_cover(
      task_id: str, tests: Dict[str, List[Any]], exclude_model: str
  ) -> Tuple[str, List[str]]:
      q, U = [], set()
      for test_name, test_cover in tests.items():
          cover_set = set()
          for model_path, i_code in test_cover:
              if exclude_model not in model_path:
                  cover_set.add((model_path, i_code))
          q.append((test_name, cover_set))
          U = U.union(cover_set)
      # Greedy algorithm for min set cover
      min_cover = []
      while len(U) > 0:
          max_uncover_set, max_test_name = {}, ""
          for test_name, cover_set in q:
              if len(cover_set) > len(max_uncover_set):
                  max_uncover_set = cover_set
                  max_test_name = test_name
          min_cover.append(max_test_name)
          U = U - max_uncover_set
          qq = []
          for test_name, cover_set in q:
              new_cover_set = U.intersection(cover_set)
              if len(new_cover_set) != 0:
                  qq.append((test_name, new_cover_set))
          q = qq
      return task_id, min_cover
  
  
  def parallel_greedy_cover(
      info_dict: Optional[Dict[str, Dict[str, List[Any]]]],
      exclude_model: str,
      type: str,
      **kwargs,
  ) -> Dict[str, List[str]]:
      plus_tests = {task_id: [] for task_id in task_ids}
  
      with ProcessPoolExecutor(max_workers=32) as executor:
          futures = []
          for task_id in task_ids:
              if type == "sample":
                  path_task_id = to_path(task_id)
                  sample_dir = kwargs["sample_dir"]
                  with open(os.path.join(sample_dir, f"{path_task_id}.pkl"), "rb") as f:
                      td = pickle.load(f)
                  args = (task_id, td, exclude_model)
              else:
                  args = (task_id, info_dict[task_id], exclude_model)
              futures.append(executor.submit(greedy_cover, *args))
          for future in track(as_completed(futures), f"min set cover :: {type}"):
              task_id, min_cover = future.result()
              plus_tests[task_id] = min_cover
  
      return plus_tests
  
  
  #####################
  # Collect Set Cover #
  #####################
  
  
  def get_coverage_set_cover(
      coverage_dir: str, exclude_model: str, dataset: str
  ) -> Dict[str, List[str]]:
      coverage_info_dict = collect_coverage_info(coverage_dir, dataset)
      return parallel_greedy_cover(coverage_info_dict, exclude_model, "coverage")
  
  
  def get_mutation_set_cover(
      mutation_dir: str, exclude_model: str, dataset: str
  ) -> Dict[str, List[str]]:
      mutation_info_dict = collect_mutation_info(
          os.path.join(mutation_dir, "eval_results.json"), dataset
      )
      return parallel_greedy_cover(mutation_info_dict, exclude_model, "mutation")
  
  
  def get_sample_set_cover(
      sample_dir: str, sample_eval_dir: str, exclude_model: str, dataset: str
  ) -> Dict[str, List[str]]:
      collect_sample_info(sample_dir, sample_eval_dir, dataset)
      return parallel_greedy_cover(None, exclude_model, "sample", sample_dir=sample_dir)
  
  
  #################
  # pass@1 greedy #
  #################
  
  
  def compute_avg_test(set_cover_info: Dict[str, List[str]]) -> float:
      sum_tests = sum(
          len(problems[task_id]["base_input"]) + len(set_cover_info[task_id])
          for task_id in task_ids
      )
      return sum_tests / problem_count
  
  
  def gen_report(set_cover_info: Dict[str, List[str]], sample_eval_dir: str, model: str):
      tsr_dict = {"ntests": compute_avg_test(set_cover_info), "pass@1": 0}
      model_path = os.path.join(sample_eval_dir, f"{model}_temp_0.0", "eval_results.json")
      with open(model_path, "r") as f:
          mdict = json.load(f)
      correct_cnt = 0
      for task_id in task_ids:
          legacy_task_id = task_id
          if legacy_task_id not in mdict["eval"]:
              legacy_task_id = legacy_task_id.replace("/", "_")
          if mdict["eval"][legacy_task_id]["base"][0][0] != "success":
              continue
          correct = True
          for plus_id in set_cover_info[task_id]:
              index = int(plus_id.split("_")[-1])
              if mdict["eval"][legacy_task_id]["plus"][0][1][index] == False:
                  correct = False
                  break
          if correct:
              correct_cnt += 1
      tsr_dict["pass@1"] = correct_cnt / problem_count
      return tsr_dict
  
  
  def dump_humaneval_plus_mini(set_cover_info: Dict[str, List[str]], mini_path: str):
      new_problems = []
      for task_id in task_ids:
          otask = problems[task_id]
          task = {
              "task_id": task_id,
              "prompt": otask["prompt"],
              "contract": otask["contract"],
              "canonical_solution": otask["canonical_solution"],
              "entry_point": otask["entry_point"],
              "base_input": otask["base_input"],
              "plus_input": [],
              "atol": otask["atol"],
          }
          for plus_test in set_cover_info[task_id]:
              index = int(plus_test.split("_")[-1])
              task["plus_input"].append(otask["plus_input"][index])
          new_problems.append(deepcopy(task))
      write_jsonl(os.path.join(mini_path, "HumanEvalPlus-Mini.jsonl"), new_problems)
  
  
  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__":
      parser = argparse.ArgumentParser()
      parser.add_argument("--model", required=True, type=str, help="Model for testing")
      parser.add_argument("--dataset", type=str, choices=["humaneval", "mbpp"])
      parser.add_argument(
          "--report_dir", type=str, help="Path to JSON report and cache files"
      )
      parser.add_argument(
          "--sample_eval_dir", type=str, help="Path to sample evaluation files"
      )
      parser.add_argument("--mini_path", type=str, help="Path to Mini Dataset")
      args = parser.parse_args()
  
      global_util_init(args.dataset)
      main(args)
  