Files
2026-07-13 12:49:17 +08:00

103 lines
3.6 KiB
Python

import argparse
import json
import os
from pathlib import Path
from dotenv import load_dotenv
from trae_selector.selector_evaluation import SelectorEvaluation
from trae_agent.utils.config import Config
_ = load_dotenv() # take environment variables
def main():
parser = argparse.ArgumentParser()
_ = parser.add_argument(
"--instances_path",
default="swe_bench/swebench-verified.json",
help="Path to instances JSON file",
)
_ = parser.add_argument("--candidate_path", required=True, help="Path to candidate patches")
_ = parser.add_argument("--result_path", required=True, help="Path to save results")
_ = parser.add_argument(
"--num_candidate", type=int, default=10, help="The number of candidate patches"
)
_ = parser.add_argument("--max_workers", type=int, default=10, help="Max number of workers")
_ = parser.add_argument(
"--group_size", type=int, default=10, help="Group size of candidate patches"
)
_ = parser.add_argument(
"--max_retry", type=int, default=3, help="Max retry times of LLM responses"
)
_ = parser.add_argument(
"--max_turn", type=int, default=50, help="Max turn times of Selector Agent"
)
_ = parser.add_argument("--majority_voting", action=argparse.BooleanOptionalAction)
_ = parser.add_argument(
"--config_file", type=str, default="config.yaml", help="Path to config file"
)
_ = parser.add_argument("--model_name", type=str, default="default_model", help="Model name")
args = parser.parse_args()
args.log_path = os.path.join(args.result_path, "log")
args.output_path = os.path.join(args.result_path, "output")
args.patches_path = os.path.join(args.result_path, "patch")
args.statistics_path = os.path.join(args.result_path, "statistics")
[
os.makedirs(_)
for _ in [args.log_path, args.patches_path, args.output_path, args.statistics_path]
if not os.path.exists(_)
]
with open(args.instances_path, "r") as file:
instance_list = json.load(file)
config = Config.create(config_file=args.config_file)
if not config.models:
raise ValueError("No models found in config file.")
if args.model_name not in config.models:
raise ValueError(f"Model {args.model_name} not found in config file.")
llm_config = config.models[args.model_name]
llm_config.resolve_config_values()
candidate_dic = {}
with open(args.candidate_path, "r") as file:
for line in file.readlines():
candidate = json.loads(line.strip())
if "regressions" not in candidate:
candidate["regressions"] = []
for _ in range(len(candidate["patches"])):
candidate["regressions"].append([])
candidate_dic[candidate["instance_id"]] = candidate
tools_path = Path(__file__).parent / "trae_selector/tools"
try:
log_path = Path(args.log_path)
log_path.mkdir(parents=True, exist_ok=True)
except Exception:
print(f"Error creating log path for {args.log_path}")
exit()
evaluation = SelectorEvaluation(
llm_config,
args.num_candidate,
args.max_retry,
args.max_turn,
args.log_path,
args.output_path,
args.patches_path,
instance_list,
candidate_dic,
tools_path.as_posix(),
args.statistics_path,
args.group_size,
majority_voting=args.majority_voting,
)
# evaluation.run_one("astropy__astropy-14369")
evaluation.run_all(max_workers=args.max_workers)
if __name__ == "__main__":
main()