parameterized min number of likes and downloads

supported_models.py

@@ -59,7 +59,12 @@ def get_page(page_index: int):
     return None
 
 
-def card_filter(model_card: Tag, model_name: str) -> bool:
+def card_filter(
+        model_card: Tag,
+        model_name: str,
+        min_number_of_downloads: int,
+        min_number_of_likes: int,
+) -> bool:
     if model_name in BLACKLISTED_MODEL_NAMES:
         return False
     numeric_contents = get_numeric_contents(model_card)
@@ -68,33 +73,55 @@ def card_filter(model_card: Tag, model_name: str) -> bool:
         # It means that he doesn't have any downloads/likes, so it's not a valid model card.
         return False
     number_of_downloads = convert_to_int(numeric_contents[0])
-    if number_of_downloads < MIN_NUMBER_OF_DOWNLOADS:
+    if number_of_downloads < min_number_of_downloads:
         return False
     number_of_likes = convert_to_int(numeric_contents[1])
-    if number_of_likes < MIN_NUMBER_OF_LIKES:
+    if number_of_likes < min_number_of_likes:
         return False
     return True
 
 
-def get_model_names(soup: BeautifulSoup):
+def get_model_names(
+        soup: BeautifulSoup,
+        min_number_of_downloads: int,
+        min_number_of_likes: int,
+) -> Generator[str, None, None]:
     model_cards: List[Tag] = soup.find_all("article", class_="overview-card-wrapper group", recursive=True)
     for model_card in model_cards:
         model_name = get_model_name(model_card)
-        if card_filter(model_card, model_name):
+        if card_filter(
+                model_card=model_card,
+                model_name=model_name,
+                min_number_of_downloads=min_number_of_downloads,
+                min_number_of_likes=min_number_of_likes
+        ):
             yield model_name
 
 
-def generate_supported_model_names() -> Generator[str, None, None]:
+def generate_supported_model_names(
+        min_number_of_downloads: int,
+        min_number_of_likes: int,
+) -> Generator[str, None, None]:
     with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
         future_to_index = {executor.submit(get_page, index): index for index in range(100)}
         for future in as_completed(future_to_index):
             soup = future.result()
             if soup:
-                yield from get_model_names(soup)
-
-
-def get_supported_model_names() -> Set[str]:
-    return set(generate_supported_model_names())
+                yield from get_model_names(
+                    soup=soup,
+                    min_number_of_downloads=min_number_of_downloads,
+                    min_number_of_likes=min_number_of_likes,
+                )
+
+
+def get_supported_model_names(
+        min_number_of_downloads: int = MIN_NUMBER_OF_DOWNLOADS,
+        min_number_of_likes: int = MIN_NUMBER_OF_LIKES,
+) -> Set[str]:
+    return set(generate_supported_model_names(
+        min_number_of_downloads=min_number_of_downloads,
+        min_number_of_likes=min_number_of_likes,
+    ))
 
 
 if __name__ == "__main__":