def add_TBranT_config_args(parser):
    """TBranT arguments"""

    group = parser.add_argument_group('model_TBranT', 'TBranT model configuration')
    group.add_argument('--head_num', type=int, default=1, help='Number of heads of Transformer Encoder in T-BranT')
    group.add_argument('--layer_num', type=int, default=1, help='Number of Transformer Encoder layers in T-BranT')
    group.add_argument('--tree_gate', default=False, action='store_true', 
        help='Concatenate the tree features and candidate variable features.')
    group.add_argument('--graph', default=False, action='store_true', 
        help='Use the history search tree to obtain global tree features.')
    
    return parser