from baselines import get_method_class
import argparse

def parse_args():
    parser = argparse.ArgumentParser(description="Merge batched test cases in save_dir into test_cases.json.")
    parser.add_argument("--method_name", type=str, default='GCG',
                        help="The name of the red teaming method")
    parser.add_argument("--save_dir", type=str, default='./test_cases',
                        help="The directory used for saving test cases")
    args = parser.parse_args()

    return args

def main():
    # ========== load arguments and config ========== #
    args = parse_args()
    print(args)
    
    method_name = args.method_name

    # ========== merge test cases ========== #
    print('Merging test cases')
    method_class = get_method_class(method_name)
    method_class.merge_test_cases(args.save_dir)

if __name__ == "__main__":
    main()