Skip to content

Commit

Permalink
Add tool to list top models (#707)
Browse files Browse the repository at this point in the history
feat: add tool to list top models
  • Loading branch information
dacorvo committed Sep 30, 2024
1 parent fcda0f1 commit 1ef1e12
Showing 1 changed file with 41 additions and 0 deletions.
41 changes: 41 additions & 0 deletions tools/list_top_models.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,41 @@
import argparse
from typing import Optional

from huggingface_hub import HfApi


class ModelStats(HfApi):

class Sort:
DOWNLOADS = "downloads"
TRENDING = "trendingScore"

def __init__(
self,
limit: Optional[int] = 20,
sort: Optional[Sort] = Sort.TRENDING,
model_name: Optional[str] = None,
task: Optional[str] = None,
):
super().__init__()
self.models = list(self.list_models(filter=model_name, task=task, limit=limit, sort=sort, full=True))


def main():
parser = argparse.ArgumentParser()
parser.add_argument("--limit", type=int, default=20, help="The number of models to return.")
parser.add_argument(
"--sort", type=str, choices=["trending", "downloads"], default="trending", help="The models sorting criteria."
)
parser.add_argument("--task", type=str, help="An optional task to filter models.")
parser.add_argument("--model_name", type=str, help="An optional model/arch name to filter models.")
args = parser.parse_args()
stats = ModelStats(
limit=args.limit, sort=getattr(ModelStats.Sort, args.sort.upper()), model_name=args.model_name, task=args.task
)
for model in stats.models:
print(model.id, model.downloads)


if __name__ == "__main__":
main()

0 comments on commit 1ef1e12

Please sign in to comment.