def cli_main():
import argparse
parser = argparse.ArgumentParser(
description=
"Downloading/Decompressing code_search_net dataset(s) or Tree-Sitter Library(ies)"
)
parser.add_argument(
"--yaml_file",
"-f",
default='config/ruby',
type=str,
help="load {yaml_file}.yml for train",
)
args = parser.parse_args()
yaml_file = os.path.join(os.path.dirname(__file__),
'{}.yml'.format(args.yaml_file))
LOGGER.info('Load arguments in {}'.format(yaml_file))
args = load_yaml(yaml_file)
LOGGER.info(args)
if args['distributed_training']['distributed_init_method'] is None:
distributed_utils.infer_init_method(args)
if args['distributed_training']['distributed_init_method'] is not None:
# distributed training
if torch.cuda.device_count(
) > 1 and not args['distributed_training']['distributed_no_spawn']:
start_rank = args['distributed_training']['distributed_rank']
args['distributed_training'][
'distributed_rank'] = None # assign automatically
torch.multiprocessing.spawn(
fn=distributed_main,
args=(args, start_rank),
nprocs=torch.cuda.device_count(),
)
else:
distributed_main(args['distributed_training']['device_id'], args)
elif args['distributed_training']['distributed_world_size'] > 1:
# fallback for single node with multiple GPUs
assert args['distributed_training'][
'distributed_world_size'] <= torch.cuda.device_count()
port = random.randint(10000, 20000)
args['distributed_training'][
'distributed_init_method'] = 'tcp://localhost:{port}'.format(
port=port)
args['distributed_training'][
'distributed_rank'] = None # set based on device id
torch.multiprocessing.spawn(
fn=distributed_main,
args=(args, ),
nprocs=args['distributed_training']['distributed_world_size'],
)
else:
LOGGER.info('single GPU training...')
single_main(args)
def cli_main():
import argparse
parser = argparse.ArgumentParser(
description="Downloading/Decompressing CodeSearchNet dataset(s) or Tree-Sitter Library(ies)")
parser.add_argument(
"--yaml_file", "-f", type=str, help="load {yaml_file}.yml for train",
default='config/preprocess',
)
args = parser.parse_args()
LOGGER.info(args)
yaml_file = os.path.join(os.path.dirname(__file__), '{}.yml'.format(args.yaml_file))
LOGGER.info('Load arguments in {}'.format(yaml_file))
args = load_yaml(yaml_file)
LOGGER.info(args)
main(args)
def cli_main():
import argparse
parser = argparse.ArgumentParser(
description=
"Generating raw/bin CodeSearchNet dataset(s) for retrieval task")
parser.add_argument(
"--yaml_file",
"-f",
default='config/ruby',
type=str,
help="load {yaml_file}.yml for train",
)
args = parser.parse_args()
LOGGER.info(args)
yaml_file = os.path.join(os.path.dirname(__file__),
'{}.yml'.format(args.yaml_file))
LOGGER.info('Load arguments in {}'.format(yaml_file))
args = load_yaml(yaml_file)
LOGGER.info(args)
main(args)