app.py

import gradio as gr
from gradio.oauth import OAuthToken
from huggingface_hub import HfApi, whoami
from huggingface_hub.hf_api import DatasetInfo, ModelInfo, SpaceInfo


def get_token(oauth_token: OAuthToken):
    return oauth_token.token


def list_orgs(oauth_token: OAuthToken = None):
    if oauth_token is None:
        return []
    data = whoami(oauth_token.token)
    if data["auth"]["type"] == "oauth":
        organisations = [data["name"]] + [org["name"] for org in data["orgs"]]
    else:
        organisations = [
            entry["entity"]["name"]
            for entry in data["auth"]["accessToken"]["fineGrained"]["scoped"]
            if "repo.write" in entry["permissions"]
        ]
        organisations = [org for org in organisations if org != data["name"]]
        organisations = [data["name"]] + organisations
    return organisations


def get_org_dropdown(oauth_token: OAuthToken = None):
    orgs = list_orgs(oauth_token)
    return gr.Dropdown(
        label="Organization",
        choices=orgs,
        value=orgs[0] if orgs else None,
        allow_custom_value=True,
    )


def get_resource_type_dropdown(
    org_name: str, resource_type: str, oauth_token: OAuthToken = None
):
    if oauth_token is None:
        return ""
    if not all([org_name, resource_type]):
        return ""

    hf_client = HfApi(token=oauth_token.token)
    if resource_type == "model":
        resources = hf_client.list_models(author=org_name)
    elif resource_type == "dataset":
        resources = hf_client.list_datasets(author=org_name)
    elif resource_type == "space":
        resources = hf_client.list_spaces(author=org_name)
    else:
        resources = []

    resources = [resource.id for resource in resources]
    return gr.Dropdown(
        label="Resource",
        choices=resources,
        value=resources[0] if resources else None,
    )


def get_resource(resource_id: str, resource_type: str, oauth_token: OAuthToken = None):
    if oauth_token is None:
        return ""
    if not all([resource_id, resource_type]):
        return ""
    hf_client = HfApi(token=oauth_token.token)

    resource: ModelInfo | DatasetInfo | SpaceInfo = hf_client.repo_info(
        resource_id, repo_type=resource_type
    )
    return str(resource)


def delete_resource(
    resource_id: str, resource_type: str, oauth_token: OAuthToken = None
):
    if oauth_token is None:
        return []
    if not all([resource_id, resource_type]):
        return []
    hf_client = HfApi(token=oauth_token.token)

    hf_client.delete_repo(resource_id, repo_type=resource_type, missing_ok=True)

    return gr.Markdown(f"Deleted {resource_id}", visible=True)


with gr.Blocks(theme="Nymbo/Nymbo_Theme") as demo:
    gr.Markdown("# 🧹 🧼 🗑️ Hub Repo Cleaner 🗑️ 🧼 🧹")
    gr.LoginButton()
    org_name = gr.Dropdown(label="Organization", choices=list_orgs())
    resource_type = gr.Dropdown(
        label="Resource Type",
        choices=["dataset", "model", "space"],
        value="dataset",
    )
    resources = gr.Dropdown(label="Resources")
    resource_info = gr.Code(language="json", label="Resource Info")

    org_name.change(
        get_resource_type_dropdown,
        inputs=[org_name, resource_type],
        outputs=[resources],
    )
    resource_type.change(
        get_resource_type_dropdown,
        inputs=[org_name, resource_type],
        outputs=[resources],
    )
    resources.change(
        get_resource,
        inputs=[resources, resource_type],
        outputs=[resource_info],
    )

    markdown = gr.Markdown("message", label="Message", visible=False)
    button = gr.Button("Delete")
    button.click(
        delete_resource, inputs=[resources, resource_type], outputs=[markdown]
    ).then(
        fn=get_resource_type_dropdown,
        inputs=[org_name, resource_type],
        outputs=[resources],
    )
    demo.load(get_org_dropdown, outputs=[org_name])

demo.launch()