def parse_args()

in sagemaker_notebook_instance/containers/relationship_extraction/package/training.py [0:0]


def parse_args(sys_args):
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--learning-rate",
        type=float,
        default=0.0007
    )
    parser.add_argument(
        "--gpus",
        type=int,
        default=os.environ.get("SM_NUM_GPUS", 0)
    )
    parser.add_argument(
        "--model-dir",
        type=str,
        default=os.environ.get("SM_MODEL_DIR")
    )
    parser.add_argument(
        "--output-dir",
        type=str,
        default=os.environ.get("SM_OUTPUT_DATA_DIR")
    )
    parser.add_argument(
        "--train-data-dir",
        type=str,
        default=os.environ.get("SM_CHANNEL_TRAIN"),
    )
    parser.add_argument(
        "--test-data-dir",
        type=str,
        default=os.environ.get("SM_CHANNEL_TEST")
    )

    args, _ = parser.parse_known_args(sys_args)
    return args