Update app.py

app.py

@@ -1,53 +1,83 @@
 import gradio as gr
-import pyjosa
-from pyjosa.josa import Josa
+import pandas as pd
 import random
 import json
 
 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):
+    chosen = random.choice(pokemons_set)
+    name = chosen['name']
+    image_path = chosen['image_path']
+    answer.value = QUESTION_TEMPLATE['answer'].format(name=name)
+    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)
+    return q, a
+
+MD = """# 포켓몬 퀴즈
+"""
 with gr.Blocks() as demo:
     quiz_start = gr.State(value=False)
+    score = gr.State(value=0)
     answer = gr.State(value="")
-    chatbot = gr.Chatbot(bubble_full_width=False)
-    msg = gr.Textbox(value="퀴즈 시작")
-    clear = gr.ClearButton([msg, chatbot])
-    
-    
-    def respond(message, chat_history, request: gr.Request):
-        print(request.client.host)
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown(MD)
+        with gr.Column():
+            with gr.Row():
+                generation = gr.Dropdown(
+                            [f"{k}세대" for k in range(1, 10)] + ["모든 세대"], value="1세대", label="포켓몬 세대", info="원하는 포켓몬 세대를 선택하세요."
+                        )
+                poke_types = gr.Dropdown(
+                            pokemons_types, value="모든 타입", label="포켓몬 타입", info="원하는 포켓몬 타입을 선택하세요."
+                    )
+            chatbot = gr.Chatbot(bubble_full_width=False)
+            msg = gr.Textbox(value="퀴즈 시작")
+            clear = gr.ClearButton([msg, chatbot, quiz_start])
+
+    def respond(gen, types, message, chat_history, request: gr.Request):
+        start, end = GEN_RANGE[gen]
+        sdf = df[start:end]
+        pokemons_set = sdf[sdf['types'].apply(lambda x: (types in x)) | (types == "모든 타입")]
+        pokemons_set = pokemons_set.to_dict("records")
         if not quiz_start.value:
-            chosen = random.choice(pokemons)
-            name = chosen['name']
-            image_path = chosen['image_path']
-            answer.value = QUESTION_TEMPLATE['answer'].format(name=name)
-            img_url = f"https://huggingface.co/spaces/yoon-gu/pokemon/resolve/main/{image_path}"
             if "퀴즈 시작" == message:
-                bot_message = "퀴즈를 시작합니다.\n" + QUESTION_TEMPLATE["question"].format(img_url=img_url)
+                q, a = get_question_answer(pokemons_set)
+                bot_message = f"퀴즈를 시작합니다.\n{q}"
+                answer.value = a
                 quiz_start.value = True
             else:
                 bot_message = "퀴즈를 시작하고 싶으시면, **퀴즈 시작**이라고 말씀해주세요."
         else:
             if answer.value == message:
-                chosen = random.choice(pokemons)
-                name = chosen['name']
-                image_path = chosen['image_path']
-                answer.value = QUESTION_TEMPLATE['answer'].format(name=name)
-                img_url = f"https://huggingface.co/spaces/yoon-gu/pokemon/resolve/main/{image_path}"
-                bot_message = "🎉정답입니다! 다음 문제입니다." + QUESTION_TEMPLATE["question"].format(img_url=img_url)
+                q, a = get_question_answer(pokemons_set)
+                answer.value = a
+                bot_message = f"🎉정답입니다! 다음 문제입니다.\n{q}"
             else:
-                try:
-                    trial = Josa.get_full_string(message, "는")
-                except pyjosa.exceptions.NotHangleException:
-                    trial = f"{message}!?"
-                bot_message = f"***{trial[:-1]}***{trial[-1]} 정답일까요? 🧐 다시 한번 생각해보세요."
-        
+                bot_message = f"***{message}***!? 🧐 다시 한번 생각해보세요."
+
         chat_history.append((message, bot_message))
         return "", chat_history
     
-    msg.submit(respond, [msg, chatbot], [msg, chatbot])
+    msg.submit(respond, [generation, poke_types, msg, chatbot], [msg, chatbot])
 
 demo.queue(concurrency_count=3)
 demo.launch()