app.py

import spaces
import gradio as gr
from hfsearch import (HFSearchResult, search, update_filter, update_df, get_labels, get_valid_labels,
                      get_tags, get_subtag_categories, update_subtag_items, update_tags, update_subtags,
                      search_ref_repos, search_cols, DS_SIZE_CATEGORIES, SPACE_HARDWARES, SPACE_STAGES, TYPES_DESC)
from gradio_huggingfacehub_search import HuggingfaceHubSearch

CSS = """
.title { align-items: center; text-align: center; }
.info { align-items: center; text-align: center; }
"""

with gr.Blocks(theme="NoCrypt/miku", fill_width=True, css=CSS) as demo:
    gr.Markdown("# Search Hugging Face🤗", elem_classes="title")
    with gr.Column():
        search_result = gr.State(value=HFSearchResult())
        with gr.Tab("Basic Search"):
            with gr.Group():
                with gr.Row(equal_height=True):
                    repo_types = gr.CheckboxGroup(label="Repo type", choices=["model", "dataset", "space", "collection"], value=["model", "dataset", "space"])
                    filter_str = gr.Textbox(label="Filter", info="String(s) to filter repos", value="")
                    sort = gr.Dropdown(label="Sort", choices=["last_modified", "likes", "downloads", "trending_score"], value="likes")
                with gr.Accordion("Advanced", open=False):
                    ui_mode = gr.Radio(label="Mode", choices=["PC", "Phone"], value="Phone")
                    with gr.Row(equal_height=True):
                        search_str = gr.Textbox(label="Search", info="A string that will be contained in the returned repo ids", placeholder="bert", value="", lines=1)
                        #author = gr.Textbox(label="Author", info="The author (user or organization)", value="", lines=1)
                        author = HuggingfaceHubSearch(label="Author", placeholder="The author (user or organization)", search_type="user", sumbit_on_select=False)
                        followed = HuggingfaceHubSearch(label="Followed by", placeholder="Extract repos followed by this user", search_type="user", sumbit_on_select=False)
                    with gr.Column():
                        tags = gr.Textbox(label="Tags", info="Tag(s) to filter repos", value="")
                        with gr.Accordion("Tag input assistance", open=False):
                            with gr.Row(equal_height=True):
                                tag_item = gr.Dropdown(label="Item", choices=get_tags(), value=get_tags()[0], allow_custom_value=True, scale=4)
                                tag_btn = gr.Button("Add", scale=1)
                            with gr.Row(equal_height=True):
                                subtag_cat = gr.Dropdown(label="Category", choices=get_subtag_categories(), value=get_subtag_categories()[0], scale=2)
                                subtag_item = gr.Dropdown(label="Item", choices=[""], value="", allow_custom_value=True, scale=2)
                                subtug_btn = gr.Button("Add", scale=1)
                    with gr.Row(equal_height=True):
                        with gr.Column():
                            gated_status = gr.Radio(label="Gated status", choices=["gated", "non-gated", "all"], value="all")
                            appr_status = gr.CheckboxGroup(label="Approval method", choices=["auto", "manual"], value=["auto", "manual"])
                        limit = gr.Number(label="Limit", info="If 0, fetches all models", value=1000, step=1, minimum=0, maximum=10000000)
                    with gr.Tab("for Models"):
                        with gr.Column():
                            infer_status = gr.Radio(label="Inference status", choices=["warm", "cold", "frozen", "all"], value="all")
                            gr.Markdown("[About the Inference API status (Warm, Cold, Frozen)](https://huggingface.co./docs/api-inference/supported-models)", elem_classes="info")
                    #    with gr.Row(equal_height=True):
                    #        model_task = gr.Textbox(label="Task", info="String(s) of tasks models were designed for", placeholder="fill-mask", value="")
                    #        trained_dataset = gr.Textbox(label="Trained dataset", info="Trained dataset for a model", value="")
                    with gr.Tab("for Datasets"):
                        size_categories = gr.CheckboxGroup(label="Size categories", info="The size of the dataset", choices=DS_SIZE_CATEGORIES, value=[])
                    #    task_categories = gr.Textbox(label="Task categories", info="Identify datasets by the designed task", value="")
                    #    task_ids = gr.Textbox(label="Task IDs", info="Identify datasets by the specific task", value="")
                    #    language_creators = gr.Textbox(label="Language creators", info="Identify datasets with how the data was curated", value="")
                    #    language = gr.Textbox(label="Language", info="String(s) representing two-character language to filter datasets by", value="")
                    #    multilinguality = gr.Textbox(label="Multilinguality", info="String(s) representing a filter for datasets that contain multiple languages", value="")
                    with gr.Tab("for Spaces"):
                        with gr.Row(equal_height=True):
                            hardware = gr.CheckboxGroup(label="Specify hardware", choices=SPACE_HARDWARES, value=[])
                            stage = gr.CheckboxGroup(label="Specify stage", choices=SPACE_STAGES, value=[])
                    with gr.Accordion("More Advanced", open=False):
                        with gr.Row(equal_height=True):
                            sort_method = gr.Radio(label="Sort method", choices=["ascending order", "descending order"], value="ascending order")
                            fetch_detail = gr.CheckboxGroup(label="Fetch detail", choices=["Space Runtime"], value=["Space Runtime"])
                        with gr.Row(equal_height=True):
                            show_labels = gr.CheckboxGroup(label="Show items", choices=get_labels(), value=get_valid_labels())
                run_button = gr.Button("Search", variant="primary")
        with gr.Tab("Find Serverless Inference API enabled models"):
            with gr.Group():
                with gr.Row(equal_height=True):
                    infer_repo_types = gr.CheckboxGroup(label="Repo type", choices=["model", "dataset", "space"], value=["model"], visible=False)
                    with gr.Column():
                        infer_infer_status = gr.Radio(label="Inference status", choices=["warm", "cold", "frozen", "all"], value="warm")
                        gr.Markdown("[About the Inference API status (Warm, Cold, Frozen)](https://huggingface.co./docs/api-inference/supported-models)", elem_classes="info")
                    with gr.Column():
                        infer_gated_status = gr.Radio(label="Gated status", choices=["gated", "non-gated", "all"], value="all")
                        infer_appr_status = gr.CheckboxGroup(label="Approval method", choices=["auto", "manual"], value=["auto", "manual"])
                infer_run_button = gr.Button("Search", variant="primary")
        with gr.Tab("Find recommended repos"):
            with gr.Group():
                with gr.Row(equal_height=True):
                    #rec_repo_id = gr.Textbox(label="Repo ID", info="Input your favorite repo", value="")
                    rec_repo_id = HuggingfaceHubSearch(label="Repo ID", placeholder="Input your favorite Repo ID", search_type=["model", "dataset", "space"],
                                                       sumbit_on_select=False)
                    rec_repo_types = gr.CheckboxGroup(label="Repo type", choices=["model", "dataset", "space", "collection"], value=["model", "dataset", "space", "collection"])
                with gr.Row(equal_height=True):
                    rec_sort = gr.Dropdown(label="Sort", choices=["last_modified", "likes", "downloads", "downloads_all_time", "trending_score"], value="likes")
                    rec_limit = gr.Number(label="Limit", value=20, step=1, minimum=1, maximum=1000)
                with gr.Accordion("Advanced", open=False):
                    with gr.Row(equal_height=True):
                        rec_show_labels = gr.CheckboxGroup(label="Show items", choices=get_labels(), value=get_valid_labels())
            rec_run_button = gr.Button("Search", variant="primary")
        with gr.Tab("Find collections"):
            with gr.Group():
                with gr.Row(equal_height=True):
                    #rec_repo_id = gr.Textbox(label="Repo ID", info="Input your favorite repo", value="")
                    col_repo_id = HuggingfaceHubSearch(label="Repo ID or User ID", placeholder="Input your favorite Repo ID or User ID", search_type=["model", "dataset", "space", "user"],
                                                       sumbit_on_select=False)
                    col_sort = gr.Dropdown(label="Sort", choices=["last_modified", "trending_score", "likes"], value="likes")
                with gr.Row(equal_height=True):
                    col_repo_limit = gr.Number(label="Limit for repos", value=50, step=1, minimum=1, maximum=100)
                    col_user_limit = gr.Number(label="Limit for users", value=10, step=1, minimum=0, maximum=50)
                with gr.Accordion("Advanced", open=False):
                    col_show_labels = gr.CheckboxGroup(label="Show items", choices=get_labels(), value=get_valid_labels())
            col_run_button = gr.Button("Search", variant="primary")
        with gr.Group():
            with gr.Accordion("Filter", open=False):
                hide_labels = gr.CheckboxGroup(label="Hide items", choices=[], value=[], visible=False)
                with gr.Row(equal_height=True):
                    filter_item1 = gr.Dropdown(label="Filter item", choices=[""], value="", visible=False)
                    filter1 = gr.Dropdown(label="Filter", choices=[""], value="", allow_custom_value=True, visible=False)
                    filter_btn = gr.Button("Apply filter", variant="secondary", visible=False)
            types_desc = gr.Markdown(TYPES_DESC, elem_classes="info")
            result_df = gr.DataFrame(label="Results", type="pandas", value=None, interactive=False)

    run_button.click(search, [repo_types, sort, sort_method, filter_str, search_str, author, tags, infer_status, gated_status, appr_status,
                              size_categories, limit, hardware, stage, followed, fetch_detail, show_labels, ui_mode, search_result],
                     [result_df, hide_labels, search_result])\
    .success(update_filter, [filter_item1, search_result], [filter_item1, filter1, filter_btn, search_result], queue=False)
    infer_run_button.click(search, [infer_repo_types, sort, sort_method, filter_str, search_str, author, tags, infer_infer_status, infer_gated_status, infer_appr_status,
                                    size_categories, limit, hardware, stage, followed, fetch_detail, show_labels, ui_mode, search_result],
                           [result_df, hide_labels, search_result])\
    .success(update_filter, [filter_item1, search_result], [filter_item1, filter1, filter_btn, search_result], queue=False)
    gr.on(triggers=[hide_labels.change, filter_btn.click], fn=update_df, inputs=[hide_labels, filter_item1, filter1, search_result],
          outputs=[result_df, search_result], trigger_mode="once", queue=False, show_api=False)
    filter_item1.change(update_filter, [filter_item1, search_result], [filter_item1, filter1, filter_btn, search_result], queue=False, show_api=False)
    subtag_cat.change(update_subtag_items, [subtag_cat], [subtag_item], queue=False, show_api=False)
    subtug_btn.click(update_subtags, [tags, subtag_cat, subtag_item], [tags], queue=False, show_api=False)
    tag_btn.click(update_tags, [tags, tag_item], [tags], queue=False, show_api=False)
    gr.on(triggers=[rec_run_button.click, rec_repo_id.submit], fn=search_ref_repos,
          inputs=[rec_repo_id, rec_repo_types, rec_sort, rec_show_labels, rec_limit, ui_mode, search_result],
          outputs=[result_df, hide_labels, search_result])
    gr.on(triggers=[col_run_button.click, col_repo_id.submit], fn=search_cols,
          inputs=[col_repo_id, col_sort, col_show_labels, col_repo_limit, col_user_limit, ui_mode, search_result],
          outputs=[result_df, hide_labels, search_result])

demo.queue().launch(ssr_mode=False)