frontend.py

from gradio import (
    update,
    Blocks,
    Row,
    Examples,
    Group,
    Markdown,
    Textbox,
    Button,
    Tabs,
    TabItem,
    DataFrame as GRDataFrame,
)
from ambiguous_search_backends import AmbiguousSearchBackend
from classic_search_backends import ClassicSearchBackend


def launch_frontend(
    ambiguous_search_backend: AmbiguousSearchBackend,
    classic_search_backend: ClassicSearchBackend,
):
    def on_ambiguous_search_click(query: str):
        result = ambiguous_search_backend.submit(query)
        return GRDataFrame.update(value=result, visible=True)

    def on_ambiguous_search_query_clear_click():
        return Textbox.update(value="")

    def on_classic_search_click(query: str):
        result = classic_search_backend.submit(query)
        return GRDataFrame.update(value=result, visible=True)

    with Blocks() as frontend:
        Markdown("## Himitsudogu Search")
        with Tabs():
            with TabItem("あいまい検索"):
                Markdown(
                    """
                自然言語処理モデル"Sentence BERT"を使って、ドラえもんのひみつ道具をあいまい検索。
                """
                )
                query_textbox = Textbox(
                    label="自分の言葉でひみつ道具を説明してください", max_lines=1
                )
                with Row():
                    submit_button = Button("検索", variant="primary")
                    clear_button = Button("クリア")
                result_table = GRDataFrame(visible=False)
                Examples(
                    examples=[
                        "ふりかけると水を色々な性質にできる",
                        "小さいカメラが飛ぶ",
                        "壁を通り抜けられるようにする輪",
                        "地図をいじって実際の町の建物などの位置を変える",
                        "歌声を綺麗にする",
                    ],
                    inputs=query_textbox,
                )

                submit_button.click(
                    on_ambiguous_search_click,
                    inputs=query_textbox,
                    outputs=result_table,
                )
                clear_button.click(
                    on_ambiguous_search_query_clear_click,
                    inputs=None,
                    outputs=query_textbox,
                )

            with TabItem("古典検索"):
                Markdown(
                    f"""
                キーワードから検索。
                検索可能なひみつ道具は全{classic_search_backend.get_num_items()}種類。
                """
                )
                query_textbox = Textbox(label="キーワード", max_lines=1)
                with Row():
                    submit_button = Button("検索", variant="primary")
                result_table = GRDataFrame(visible=False)

                submit_button.click(
                    on_classic_search_click,
                    inputs=query_textbox,
                    outputs=result_table,
                )

    frontend.launch()