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}")
|