refactor: add feedback function, update ui

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 
 import os
 import threading
+import arrow
 import time
 import argparse
 import logging
@@ -90,27 +91,17 @@ class SentencePieceStreamer(BaseStreamer):
     def end(self):
         self.ended = True
 
-def user(prompt, history):
-    logging.info(f"[user] prompt: {prompt}")
-    logging.debug(f"[user] history: {history}")
-    
-    res = ("", history + [[prompt, None]])
-    return res
-
-def bot(
-    history, 
-    do_sample,
+def generate(
+    prompt,
+    max_new_tokens,
     temperature,
     repetition_penalty,
+
+    do_sample,
     no_repeat_ngram_size,
-    max_new_tokens,
 ):
-    logging.info("[bot]")
-    logging.info(dict(locals()))
-    logging.debug(f"history: {history}")
-    
-    # TODO: modify `<br>` back to `\n` based on the original user prinpt
-    prompt = history[-1][0]
+    log = dict(locals())
+    logging.debug(log)
     
     tokens = sp.encode(prompt)
     input_ids = torch.tensor(tokens, dtype=torch.long).unsqueeze(0).to(model.device)
@@ -120,7 +111,6 @@ def bot(
     max_possilbe_new_tokens = model.config.max_position_embeddings - len(tokens)
     max_possilbe_new_tokens = min(max_possilbe_new_tokens, max_new_tokens)
     
-    
     thr = threading.Thread(target=model.generate, args=(), kwargs=dict(
         input_ids=input_ids,
         do_sample=do_sample,
@@ -137,87 +127,123 @@ def bot(
     ))
     thr.start()
     
-    history[-1][1] = ""
     while not streamer.ended:
-        history[-1][1] = streamer.generated_text
         time.sleep(0.05)
-        yield history
+        yield streamer.generated_text
     
     # TODO: optimize for final few tokens
-    history[-1][1] = streamer.generated_text
-    logging.info(f"generation: {history[-1][1]}")
-    yield history
+    gen = streamer.generated_text
+    log.update(dict(generation=gen, time=str(arrow.now("+09:00"))))
+    logging.info(log)
+    yield gen
+
+def process_feedback(
+    rating,
+    prompt,
+    generation,
+    
+    max_new_tokens,
+    temperature,
+    repetition_penalty,
+    do_sample,
+    no_repeat_ngram_size,
+):
+    log = dict(locals())
+    log["time"] = str(arrow.now("+09:00"))
+    logging.info(log)
 
 if gr_interface:
     gr_interface.close(verbose=False)
 
 with gr.Blocks() as gr_interface:
-    chatbot = gr.Chatbot(label="StableLM JP Alpha").style(height=500)
-    
-    # generation params
-    do_sample = gr.Checkbox(True, label="Do Sample", visible=False)
-    
     with gr.Row():
-        temperature = gr.Slider(0, 1, value=0.7, step=0.05, label="Temperature")
-        repetition_penalty = gr.Slider(1, 1.5, value=1.2, step=0.05, label="Repetition Penalty")
+        gr.Markdown("# 日本語 StableLM Pre-Alpha")
     with gr.Row():
-        no_repeat_ngram_size = gr.Slider(0, 10, value=5, step=1, label="No Repeat Ngram Size")
-        max_new_tokens = gr.Slider(
-            128, 
-            model.config.max_position_embeddings, 
-            value=128, step=64, label="Max New Tokens")
-    
-    # prompt
-    # TODO: add more options
-#     prompt_options = gr.Dropdown(
-#         choices=[
-#             "運が良かったのか悪かったのか日本に帰ってきたタイミングでコロナが猛威を振るい始め、",
-#             """[問題]に対する[答え]を[選択肢]の中から選んでください。
-
-# [問題]: ある場所の周辺地域を指す言葉は？
-# [選択肢]: [空, オレゴン州, 街, 歩道橋, 近辺]
-# [答え]: 近辺
-
-# [問題]: 若くて世間に慣れていないことを何という？
-# [選択肢]: [青っぽい, 若い, ベテラン, 生々しい, 玄人]
-# [答え]: """
-#         ], 
-#         label="Prompt Options", 
-#         info="Select 1 option for quick start",
-#         allow_custom_value=False,
-#     )
-    prompt = gr.Textbox(label="Prompt", info="Pro tip: press Enter to submit directly")
-    
-    
-#     def on_prompt_options_change(pmt_opts, pmt):
-#         return pmt_opts
-    
-#     prompt_options.change(on_prompt_options_change, [prompt_options, prompt], prompt)
-
+        gr.Markdown("Description about this page.　ホゲホゲ")
     with gr.Row():
-        submit = gr.Button("Submit")
-        stop = gr.Button("Stop")
-        clear = gr.Button("Clear")
-    
-    # event handling
-    user_io = [prompt, chatbot]
-    bot_inputs = [
-        chatbot, 
-        do_sample, 
-        temperature, 
-        repetition_penalty, 
-        no_repeat_ngram_size,
-        max_new_tokens,
-    ]
-    
-    submit_event = prompt.submit(user, user_io, user_io, queue=False)\
-        .then(bot, bot_inputs, chatbot, queue=True)
-    
-    submit_click_event = submit.click(user, user_io, user_io, queue=False)\
-        .then(bot, bot_inputs, chatbot, queue=True)
-    
-    stop.click(None, None, None, cancels=[submit_event, submit_click_event], queue=False)
-    clear.click(lambda: None, None, chatbot, queue=False)
+        
+        # left panel
+        with gr.Column(scale=1):
+            
+            # generation params
+            with gr.Box():
+                gr.Markdown("パ���メータ")
+                
+                # hidden default params
+                do_sample = gr.Checkbox(True, label="Do Sample", visible=False)
+                no_repeat_ngram_size = gr.Slider(0, 10, value=5, step=1, label="No Repeat Ngram Size", visible=False)
+                
+                # visible params
+                max_new_tokens = gr.Slider(
+                    128, 
+                    min(512, model.config.max_position_embeddings), 
+                    value=128, 
+                    step=128, 
+                    label="max tokens",
+                    info="生成するトークンの最大数を指定する",
+                )
+                temperature = gr.Slider(
+                    0, 1, value=0.7, step=0.05, label="temperature", 
+                    info="低い値は出力をより集中させて決定論的にする")
+                repetition_penalty = gr.Slider(
+                    1, 1.5, value=1.2, step=0.05, label="frequency penalty",
+                    info="高い値はAIが繰り返す可能性を減少させる")
+                
+                # grouping params for easier reference
+                gr_params = [
+                    max_new_tokens,
+                    temperature,
+                    repetition_penalty,
+                    
+                    do_sample,
+                    no_repeat_ngram_size,
+                ]
+                
+        # right panel
+        with gr.Column(scale=2):
+            # user input block
+            with gr.Box():
+                textbox_prompt = gr.Textbox(
+                    label="Human",
+                    placeholder="AIに続きを書いて欲しいプロンプト",
+                    interactive=True,
+                    lines=5,
+                    value=""
+                )
+            with gr.Box():
+                with gr.Row():
+                    btn_submit = gr.Button(value="実行", variant="primary")
+                    btn_stop = gr.Button(value="中止", variant="stop")
+
+            # model output block
+            with gr.Box():
+                textbox_generation = gr.Textbox(
+                    label="AI",
+                    lines=5,
+                    value=""
+                )
+            with gr.Box():
+                with gr.Row():
+                    rating_options = [
+                        "😫すごく悪い",
+                        "😞微妙",
+                        "😐アリ",
+                        "🙂合格",
+                        "😄すごく良い",
+                    ]
+                    btn_ratings = [gr.Button(value=v) for v in rating_options]
+                    
+            with gr.Box():
+                gr.Markdown("TODO：For more feedback link for google form")
+
+    # event handling 
+    inputs = [textbox_prompt] + gr_params
+    click_event = btn_submit.click(generate, inputs, textbox_generation, queue=True)
+    btn_stop.click(None, None, None, cancels=click_event, queue=False)
+    
+    for btn_rating in btn_ratings:
+        btn_rating.click(process_feedback, [btn_rating, textbox_prompt, textbox_generation] + gr_params, queue=False)
 
+    
 gr_interface.queue(max_size=32, concurrency_count=2)
 gr_interface.launch(server_port=args.port, share=args.make_public)