From 1ef1e12a06547f456ec156b1d8d50d885ccf3610 Mon Sep 17 00:00:00 2001
From: David Corvoysier <david@huggingface.co>
Date: Mon, 30 Sep 2024 15:33:53 +0200
Subject: [PATCH] Add tool to list top models (#707)

feat: add tool to list top models
---
 tools/list_top_models.py | 41 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 41 insertions(+)
 create mode 100644 tools/list_top_models.py

diff --git a/tools/list_top_models.py b/tools/list_top_models.py
new file mode 100644
index 000000000..49d90898d
--- /dev/null
+++ b/tools/list_top_models.py
@@ -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()