davanstrien HF Staff commited on
Commit
5fa4938
·
1 Parent(s): 8561b4b

small refactor

Browse files
Files changed (1) hide show
  1. app.py +17 -4
app.py CHANGED
@@ -4,13 +4,27 @@ from cachetools import TTLCache, cached
4
  from huggingface_hub import list_models
5
  from toolz import groupby
6
  from tqdm.auto import tqdm
 
 
 
 
 
 
 
 
 
 
7
 
8
 
9
  @cached(TTLCache(maxsize=10, ttl=60 * 60 * 3))
10
  def get_all_models():
11
  models = list(
12
  tqdm(
13
- iter(list_models(cardData=True, limit=None, sort="downloads", direction=-1))
 
 
 
 
14
  )
15
  )
16
  models = [model for model in models if model is not None]
@@ -87,6 +101,8 @@ def get_grandchildren(base_model):
87
 
88
  def return_models_for_base_model(base_model):
89
  models = grouped_by_base_model.get(base_model)
 
 
90
  # sort models by downloads
91
  models = sorted(models, key=lambda x: x.downloads, reverse=True)
92
  results = ""
@@ -139,9 +155,6 @@ def return_base_model_popularity(pipeline=None):
139
 
140
 
141
  def return_base_model_popularity_by_org(pipeline=None):
142
- referenced_base_models = [
143
- f"[`{model}`](https://huggingface.co/{model})" for model in base_models
144
- ]
145
  df_with_pipeline_info = pd.DataFrame(
146
  {"base_model": base_models, "pipeline": pipeline_tags}
147
  )
 
4
  from huggingface_hub import list_models
5
  from toolz import groupby
6
  from tqdm.auto import tqdm
7
+ from dotenv import load_dotenv
8
+ import os
9
+ from huggingface_hub import HfApi, hf_hub_url, login, snapshot_download
10
+
11
+ load_dotenv(override=True)
12
+
13
+ HF_TOKEN = os.getenv("HF_TOKEN")
14
+ USER_AGENT = os.getenv("USER_AGENT")
15
+ login(token=HF_TOKEN)
16
+ hf_api = HfApi(token=HF_TOKEN, user_agent=USER_AGENT)
17
 
18
 
19
  @cached(TTLCache(maxsize=10, ttl=60 * 60 * 3))
20
  def get_all_models():
21
  models = list(
22
  tqdm(
23
+ iter(
24
+ hf_api.list_models(
25
+ cardData=True, limit=None, sort="downloads", direction=-1
26
+ )
27
+ )
28
  )
29
  )
30
  models = [model for model in models if model is not None]
 
101
 
102
  def return_models_for_base_model(base_model):
103
  models = grouped_by_base_model.get(base_model)
104
+ if models is None:
105
+ return f"No models found for base model: {base_model}"
106
  # sort models by downloads
107
  models = sorted(models, key=lambda x: x.downloads, reverse=True)
108
  results = ""
 
155
 
156
 
157
  def return_base_model_popularity_by_org(pipeline=None):
 
 
 
158
  df_with_pipeline_info = pd.DataFrame(
159
  {"base_model": base_models, "pipeline": pipeline_tags}
160
  )