File size: 1,485 Bytes
d3a17fc 15bf463 b1dd47e 15bf463 b1dd47e 15bf463 b1dd47e 15bf463 b1dd47e 15bf463 b1dd47e d3a17fc b1dd47e |
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 32 33 34 35 36 37 38 39 40 41 42 |
from typing import Generator, Set
import requests
from bs4 import BeautifulSoup
from concurrent.futures import ThreadPoolExecutor, as_completed
SUPPORTED_MODEL_NAME_PAGES_FORMAT = "https://huggingface.co/models?pipeline_tag=text-generation&library=pytorch"
MAX_WORKERS = 10
def get_model_name(model_card: BeautifulSoup) -> str:
h4_class = "text-md truncate font-mono text-black dark:group-hover:text-yellow-500 group-hover:text-indigo-600"
h4_tag = model_card.find("h4", class_=h4_class)
return h4_tag.text
def get_page(page_index: int):
curr_page_url = f"{SUPPORTED_MODEL_NAME_PAGES_FORMAT}&p={page_index}"
response = requests.get(curr_page_url)
if response.status_code == 200:
soup = BeautifulSoup(response.content, "html.parser")
return soup
return None
def get_model_names(soup):
model_cards = soup.find_all("article", class_="overview-card-wrapper group", recursive=True)
return [get_model_name(model_card) for model_card in model_cards]
def generate_supported_model_names() -> 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())
|