app.py

import os
import gradio as gr
import pandas as pd
import random
import json
import time
import wandb
from utils import Info
import logging
from typing import Optional

wandb.login(key=os.environ['WANDB_KEY'])
run = wandb.init(project="pokemon-quiz", entity="yoon-gu")

with open('pokemon.json', 'r') as f:
    pokemons = json.load(f)
pokemons_types = sorted(set([t for poke in pokemons for t in poke['types']]))
df = pd.DataFrame(pokemons)
GEN_RANGE = {
    "모든 세대": [1, 1017],
    "1세대": [1, 151],
    "2세대": [152, 251],
    "3세대": [252, 386],
    "4세대": [387, 493],
    "5세대": [494, 649],
    "6세대": [650, 721],
    "7세대": [722, 809],
    "8세대": [810, 905],
    "9세대": [906, 1017]
}

QUESTION_TEMPLATE = {"question": "다음 포켓몬의 이름은 뭘까요?![]({img_url})", "answer": "{name}"}

def get_question_answer(pokemons_set, user):
    chosen = random.choice(pokemons_set)
    name = chosen['name']
    image_path = chosen['image_path']
    img_url = f"https://huggingface.co/spaces/yoon-gu/pokemon/resolve/main/{image_path}"
    q = QUESTION_TEMPLATE["question"].format(img_url=img_url)
    a = QUESTION_TEMPLATE['answer'].format(name=name)

    candidates = random.choices([poke['name'] for poke in pokemons_set], k=3)
    candidates.append(a)
    random.shuffle(candidates)
    infos[user].candidates = candidates
    return q, a

initial_info = {"done" : True,
                "score": 0, "count": 0,
                "best_score": 0, "best_time": float("inf"),
                "time": 0.0, "comment": "",
                "history": [],
                "candidates": ['1번', '2번', '3번', '4번']}

try:
    folder = run.use_artifact("settings:latest").download()
    with open(os.path.join(folder, "users.json"), "r") as f:
        USERS = json.load(f)
    with open(os.path.join(folder, "infos.json"), "r") as f:
        infos = json.load(f)
    for k, v in infos.items():
        infos[k] = Info(**v)
except Exception as e:
    logging.error(e)
    USERS = ["김서현", "김우주", "Anonymous"]
    infos = {user: Info(name=user) for user in USERS}

MD = """# 포켓몬 퀴즈 🦑
## 공부 방법 📚

아래 포켓몬 도감을 보고 공부하면 도움이 됩니다.
- https://huggingface.co/spaces/yoon-gu/pokemon

## 사용방법 📖

1. 사용자를 선택하세요.
2. 총 퀴즈 개수를 선택하세요.
3. 포켓몬 세대를 선택하세요.
4. 포켓몬 타입을 선택하세요.

## 점수판 📊
"""

with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            markdown = gr.Markdown(MD.format(content=''))
            leader_board = gr.DataFrame(wrap=True, row_count=10)
            with gr.Row():
                gr.LoginButton()
                gr.LogoutButton()
            # user = gr.Dropdown(USERS, value="Anonymous", label="사용자", info="당신은 누구신가요?", allow_custom_value=True)
            quiz_count = gr.Radio([10, 20, 30, 40, 50], value=10, label="총 퀴즈 개수", info="퀴즈를 몇 개 풀 예정인가요?")
        with gr.Column():
            with gr.Row():
                with gr.Accordion("See Details", open=False):
                    generation = gr.CheckboxGroup(
                                [f"{k}세대" for k in range(1, 10)],
                                value=[f"{k}세대" for k in range(1, 10)],
                                label="포켓몬 세대",
                                info="원하는 포켓몬 세대를 선택하세요."
                            )
                    poke_types = gr.CheckboxGroup(
                                pokemons_types, value=pokemons_types,
                                label="포켓몬 타입",
                                info="원하는 포켓몬 타입을 선택하세요."
                        )
            with gr.Row():
                play = gr.Button(value="퀴즈 시작", label="퀴즈 시작")
                skip = gr.Button(value="문제 넘어가기", label="문제 스킵")
            chatbot = gr.Chatbot(bubble_full_width=False,
                                 height=600,
                                 avatar_images=["https://huggingface.co/spaces/yoon-gu/pokemon/resolve/main/images/No_0001_이상해씨.png",
                                                "https://huggingface.co/spaces/yoon-gu/pokemon/resolve/main/images/No_0155_브케인.png"])
            with gr.Row():
                button1 = gr.Button(value="1번")
                button2 = gr.Button(value="2번")
            with gr.Row():
                button3 = gr.Button(value="3번")
                button4 = gr.Button(value="4번")

    def respond(message, chat_history, quiz_count, gen, types, profile: Optional[gr.OAuthProfile]):
        if profile:
            user = profile.name
        else:
            raise gr.Error("사용자 이름을 입력하세요.")
        message = message.strip()
        done = infos[user].done

        dfs = []
        for g in gen:
            start, end = GEN_RANGE[g]
            dfs.append(df[start-1:end])
        sdf = pd.concat(dfs)
        pokemons_set = sdf[sdf['types'].apply(lambda x: any([t in x for t in  types]))]
        pokemons_set = pokemons_set.to_dict("records")
        if done:
            infos[user].history = []
            if "퀴즈시작" == message.replace(" ", ""):
                q, a = get_question_answer(pokemons_set, user)
                bot_message = f"퀴즈를 시작합니다.\n{q}"
                infos[user].answer = a
                infos[user].done = False
                infos[user].score = 0
                infos[user].count = 0
                infos[user].time = time.time()
                infos[user].generations = gen
                infos[user].types = types
            else:
                bot_message = "퀴즈를 시작하고 싶으시면, **퀴즈 시작** 버튼을 누르세요."
        else:
            if infos[user].answer == message:
                q, a = get_question_answer(pokemons_set, user)
                infos[user].answer = a
                infos[user].score += 10
                infos[user].count += 1
                bot_message = f"🎉정답입니다! 다음 문제입니다.\n{q}\n- 현재 점수: {infos[user].score:3.1f}점\n- 소요 시간: {time.time() - infos[user].time:4.3f}초"

            elif "포기하기" == message.replace(" ", ""):
                bot_message = f"퀴즈를 강제 종료합니다."
                infos[user].done = True
            elif "문제 넘어가기" == message:
                infos[user].count += 1
                q, a = get_question_answer(pokemons_set, user)
                infos[user].answer = a
                bot_message = f"문제를 넘어갑니다. 다음문제입니다.\n{q}"
            else:
                hint1 = ""
                for i, y in enumerate(infos[user].answer):
                    if i < len(message):
                        if message[i] == y:
                            hint1 += y
                        else:
                            hint1 += "X"
                    else:
                        hint1 += "X"
                bot_message = f"***{message}***!? 🧐 다시 한번 생각해보세요.\n힌트: {hint1}"
                infos[user].score -= 1

            # gr.Info(f"{user}님, 총 {quiz_count}개 문제 중 {infos[user].count+1}번째 문제 진행 중입니다.")

            if quiz_count == infos[user].count:
                bot_message = f"모든 퀴즈를 다 풀었습니다. 점수는 {infos[user].score:3.1f}점 입니다."
                infos[user].done = True
                if infos[user].score >= infos[user].best_score:
                    infos[user].best_score = infos[user].score
                    infos[user].best_time = min(time.time() - infos[user].time, infos[user].best_time)
                    infos[user].best_time = round(infos[user].best_time, 2)
                    infos[user].comment = f"{gen}+{types}"
                infos[user].history = []
                with open("infos.json", "w") as f:
                    json.dump({k: v.model_dump() for k, v in infos.items()}, f, indent=4, ensure_ascii=False)
                with open("users.json", "w") as f:
                    json.dump(USERS, f, indent=4, ensure_ascii=False)
                artifact = wandb.Artifact("settings", type="History")
                artifact.add_file("infos.json")
                artifact.add_file("users.json")
                run.log_artifact(artifact)

        infos[user].history.append((message, bot_message))
        leader_board = sorted(infos.items(), key=lambda x: (x[1].best_score, -x[1].best_time), reverse=True)
        lbdf = pd.DataFrame([dict(**a[1].model_dump()) for a in leader_board])
        lbdf.rename(columns={'name': "이름", 'best_score': "최고점수", 'best_time': "시간기록", 'comment': "퀴즈유형"}, inplace=True)
        lbdf.index += 1
        md = lbdf[['이름', '최고점수', '시간기록', '퀴즈유형']]
        if infos[user].done:
            btn = gr.Button(value="퀴즈 시작", label="퀴즈 시작")
        else:
            btn = gr.Button(value="포기하기", label="퀴즈 중단")

        random_buttons = [gr.Button(value=c) for c in infos[user].candidates]
        return btn, infos[user].history, md, *random_buttons

    play.click(respond, 
               inputs=[play, chatbot, quiz_count, generation, poke_types],
               outputs=[play, chatbot, leader_board, button1, button2, button3, button4])
    skip.click(respond,
               inputs=[skip, chatbot, quiz_count, generation, poke_types],
               outputs=[play, chatbot, leader_board, button1, button2, button3, button4])
    
    for btn in [button1, button2, button3, button4]:
        btn.click(respond, 
                inputs=[btn, chatbot, quiz_count, generation, poke_types],
                outputs=[play, chatbot, leader_board, button1, button2, button3, button4])
    
    def update_table(profile: Optional[gr.OAuthProfile]):
        user = profile.name
        global USERS
        if user not in USERS:
            USERS += [user]
            infos[user] = Info(name=user)
        leader_board = sorted(infos.items(), key=lambda x: (x[1].best_score, -x[1].best_time), reverse=True)
        lbdf = pd.DataFrame([dict(**a[1].model_dump()) for a in leader_board])
        lbdf.rename(columns={'name': "이름", 'best_score': "최고점수", 'best_time': "시간기록", 'comment': "퀴즈유형"}, inplace=True)
        lbdf.index += 1
        md = lbdf[['이름', '최고점수', '시간기록', '퀴즈유형']]

        if infos[user].done:
            btn = gr.Button(value="퀴즈 시작", label="퀴즈 시작")
        else:
            btn = gr.Button(value="포기하기", label="퀴즈 중단")

        random_buttons = [gr.Button(value=c) for c in infos[user].candidates]
        return new_dropdown, btn, md, infos[user].history, *random_buttons

    demo.load(update_table,
              inputs=None,
              outputs=[play, leader_board, chatbot, button1, button2, button3, button4])

demo.queue(concurrency_count=3)
demo.launch()