run.py

import gradio as gr
from huggingface_hub import list_models
from datasets import load_dataset
from type.dataset_type import TanukiPhase2AnnotationDataset

js = """
function blockEnter(event) {
    document.addEventListener('DOMContentLoaded', (event) => {
  const textbox = document.getElementById('answer');
  textbox.addEventListener('keydown', (e) => {
    if (e.key === 'Enter') {
      e.preventDefault();
    }
  });
});
}
"""

target_dataset = TanukiPhase2AnnotationDataset()

def load_css():
    with open("style.css", "r") as file:
        css_content = file.read()
    return css_content

def set_name(profile: gr.OAuthProfile | None) -> str:
    if profile is None:
        return "プライベートデータセット取得のためにログインしてください。"
    return f'{profile.username}さん、よろしくお願いいたします。'


def list_private_models(profile: gr.OAuthProfile | None, oauth_token: gr.OAuthToken | None) -> str:
    if oauth_token is None:
        return "プライベートデータセット取得のためにログインしてください。"
    models = [
        f"{model.id} ({'private' if model.private else 'public'})"
        for model in list_models(author=profile.username, token=oauth_token.token)
    ]
    return "選択中のデータセット:\n\n" + "\n - ".join(models)


def toggle_buttons(profile: gr.OAuthProfile | None):
    if profile is None:
        return gr.update(visible=True), gr.update(visible=False)
    return gr.update(visible=False), gr.update(visible=True)


def load_data(profile: gr.OAuthProfile | None, oauth_token: gr.OAuthToken | None) -> str:
    if oauth_token is None:
        return "Please log in to load your datasets."
    dataset = load_dataset("kanhatakeyama/OrcaJaMixtral8x22b")
    return dataset


def display_dataset(profile: gr.OAuthProfile | None, oauth_token: gr.OAuthToken | None):
    if profile is None:
        return gr.update(visible=True, value="ログインしてデータセットを表示してください。"), None, None
    dataset = load_data(profile, oauth_token=oauth_token.token)
    if dataset:
        # "train" スプリットの最初のデータを取得
        first_data = dataset['train'][0]
        question = first_data.get("question", "No question found")
        answer = first_data.get("answer", "No answer found")
        return gr.update(visible=False), gr.update(value=question, interactive=False), gr.update(value=answer)
    else:
        return gr.update(visible=True, value="データセットのロードに失敗しました。"), None, None

def switch_theme(theme):
    if theme == "Dark":
        return gr.themes.Default()
    else:
        return gr.themes.Monochrome()

theme_ = gr.State("Light")

with gr.Blocks(theme=theme_, css=load_css()) as demo:

    gr.Markdown("# 自動生成データセット アノテーション for Tanuki 2Phase")

    with gr.Tab("アノテーション"):

        def update_theme():
            new_theme = "Dark" if theme_.value == "Light" else "Light"
            theme_.value = new_theme
            return switch_theme(new_theme)

        with gr.Row(equal_height=True):

            login_btn = gr.LoginButton(visible=True, scale=1)
            logout_btn = gr.LogoutButton(visible=False, scale=1)

            profile = gr.Markdown()
            data_load_btn = gr.Button("データセットを読み込む")
            theme_button = gr.Button("Switch Theme")
            theme_button.click(fn=update_theme, outputs=None)

            demo.load(set_name, inputs=None, outputs=profile)
            demo.load(toggle_buttons, inputs=None, outputs=[login_btn, logout_btn])

        dataset_display = gr.Markdown(visible=False)
        question_text = gr.Textbox(label="質問: ", interactive=False)

        with gr.Row(equal_height=True):
            good_btn = gr.Button("良い")
            bad_btn = gr.Button("悪い")

        answer_text = gr.Textbox(label="回答: 改行はShift+Enterです。", elem_id="answer", interactive=True)

        data_load_btn.click(
            display_dataset,
            inputs=None,
            outputs=[dataset_display, question_text, answer_text],
        )

        def on_submit(answer_text):
            return f" {answer_text}"


    with gr.Tab("アノテ済みデータセット"):

        gr.Textbox("データセットID", lines=1, placeholder="データセットIDを入力してください。")


if __name__ == "__main__":
    demo.launch()