File size: 785 Bytes
9ae8cf2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
#!/usr/bin/env python3
import argparse
from huggingface_hub import HfApi


def main(api, model_id):
    info = api.list_repo_refs(model_id)
    branches = set([b.name for b in info.branches]) - set(["main"])

    return list(branches)


if __name__ == "__main__":
    DESCRIPTION = """
    Simple utility to get all branches from a repo
    """
    parser = argparse.ArgumentParser(description=DESCRIPTION)
    parser.add_argument(
        "--model_id",
        type=str,
        help="The name of the model on the hub to retrieve the branches from. E.g. `gpt2` or `facebook/wav2vec2-base-960h`",
    )

    args = parser.parse_args()
    model_id = args.model_id
    api = HfApi()
    branches = main(api, model_id)

    if len(branches) > 0:
        print(f"{model_id}: {branches}")