src/hub.py

import io
import json

import httpx
from huggingface_hub import HfFileSystem, ModelCard, hf_hub_url
from huggingface_hub.utils import build_hf_headers

import src.constants as constants


class Client:
    def __init__(self):
        self.client = httpx.AsyncClient(follow_redirects=True)

    async def get(self, url, headers=None, params=None):
        try:
            r = await self.client.get(url, headers=headers, params=params)
            r.raise_for_status()
        except httpx.HTTPError:
            return
        return r


client = Client()
fs = HfFileSystem()


def glob(path):
    paths = fs.glob(path)
    return paths


async def load_json_file(path):
    url = to_url(path)
    r = await client.get(url)
    if r is None:
        return
    return r.json()


async def load_jsonlines_file(path):
    url = to_url(path)
    r = await client.get(url, headers=build_hf_headers())
    if r is None:
        return
    f = io.StringIO(r.text)
    return [json.loads(line) for line in f]


def to_url(path):
    *repo_type, org_name, ds_name, filename = path.split("/", 3)
    repo_type = repo_type[0][:-1] if repo_type else None
    return hf_hub_url(repo_id=f"{org_name}/{ds_name}", filename=filename, repo_type=repo_type)


async def load_model_card(model_id):
    url = to_url(f"{model_id}/README.md")
    r = await client.get(url)
    if r is None:
        return
    return ModelCard(r.text, ignore_metadata_errors=True)


async def list_models(filtering=None):
    params = {}
    if filtering:
        params["filter"] = filtering
    r = await client.get(f"{constants.HF_API_URL}/models", params=params)
    if r is None:
        return
    return r.json()