chore: report conversation data

app.py

@@ -13,6 +13,10 @@ except:
 
     print("flash-attn installed.")
 
+import os
+import uuid
+import requests
+
 import torch
 from transformers import (
     AutoModelForCausalLM,
@@ -23,18 +27,16 @@ from transformers import (
 from threading import Thread
 
 import gradio as gr
+from dotenv import load_dotenv
 
-try:
-    import spaces
-except:
+import spaces
 
-    class spaces:
-        @staticmethod
-        def GPU(duration: int):
-            return lambda x: x
 
+load_dotenv()
 
 MODEL_NAME = "hatakeyama-llm-team/Tanuki-8B-Instruct"
+PREFERENCE_API_URL = os.getenv("PREFERENCE_API_URL")
+assert PREFERENCE_API_URL, "PREFERENCE_API_URL is not set"
 
 quantization_config = BitsAndBytesConfig(
     load_in_4bit=True,
@@ -47,19 +49,72 @@ model = AutoModelForCausalLM.from_pretrained(
 )
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 
-print(model.hf_device_map)
+print("Compiling model...")
+model = torch.compile(model)
+print("Model compiled.")
 
 
-@spaces.GPU(duration=30)
-def respond(
-    message,
+def send_report(
+    type: str,
+    data: dict,
+):
+    print(f"Sending report: {data}")
+    try:
+        res = requests.post(PREFERENCE_API_URL, json={"type": type, **data})
+        print(f"Report sent: {res.json()}")
+    except Exception as e:
+        print(f"Failed to send report: {e}")
+
+
+def send_reply(
+    reply_id: str,
+    parent_id: str,
+    role: str,
+    body: str,
+):
+
+    send_report(
+        "conversation",
+        {
+            "reply_id": reply_id,
+            "parent_id": parent_id,
+            "role": role,
+            "body": body,
+        },
+    )
+
+
+def send_score(
+    reply_id: str,
+    score: int,
+):
+    # print(f"Score: {score}, reply_id: {reply_id}")
+    send_report(
+        "score",
+        {
+            "reply_id": reply_id,
+            "score": score,
+        },
+    )
+
+
+def generate_unique_id():
+    return str(uuid.uuid4())
+
+
+@spaces.GPU(duration=45)
+def generate(
+    message: str,
     history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-    top_k,
+    system_message: str,
+    max_tokens: int,
+    temperature: float,
+    top_p: float,
+    top_k: int,
 ):
+    if not message or message.strip() == "":
+        return "", history
+
     messages = [{"role": "system", "content": system_message}]
 
     for val in history:
@@ -80,50 +135,234 @@ def respond(
     generate_kwargs = dict(
         input_ids=tokenized_input,
         streamer=streamer,
-        max_new_tokens=max_tokens,
+        max_new_tokens=int(max_tokens),
         do_sample=True,
-        temperature=temperature,
-        top_k=top_k,
-        top_p=top_p,
+        temperature=float(temperature),
+        top_k=int(top_k),
+        top_p=float(top_p),
         num_beams=1,
     )
     t = Thread(target=model.generate, kwargs=generate_kwargs)
     t.start()
 
+    # 返す値を初期化
     partial_message = ""
+
     for new_token in streamer:
         partial_message += new_token
-        yield partial_message
-
-
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(
-            value="以下は、タスクを説明する指示と、文脈のある入力の組み合わせです。要求を適切に満たす応答を書きなさい。",
-            label="システムプロンプト",
-        ),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p",
-        ),
-        gr.Slider(minimum=1, maximum=2000, value=200, step=10, label="Top-k"),
-    ],
-    examples=[
-        ["たぬきってなんですか？"],
-        ["情けは人の為ならずとはどういう意味ですか？"],
-        ["まどマギで一番可愛いのは誰？"],
-        ["明晰夢とはなんですか？"],
-        ["シュレディンガー方程式とシュレディンガーの猫はどのような関係がありますか？"],
-    ],
-    cache_examples=False,
-)
+        new_history = history + [(message, partial_message)]
+        # 入力テキストをクリアする
+        yield "", new_history
+
+
+def respond(
+    message: str,
+    history: list[tuple[str, str]],
+    system_message: str,
+    max_tokens: int,
+    temperature: float,
+    top_p: float,
+    top_k: int,
+    reply_ids: list[str],
+):
+    if len(reply_ids) == 0:
+        reply_ids = [generate_unique_id()]
+    last_reply_id = reply_ids[-1]
+    user_reply_id = generate_unique_id()
+    assistant_reply_id = generate_unique_id()
+
+    reply_ids.append(user_reply_id)
+    reply_ids.append(assistant_reply_id)
+
+    for stream in generate(
+        message,
+        history,
+        system_message,
+        max_tokens,
+        temperature,
+        top_p,
+        top_k,
+    ):
+        yield *stream, reply_ids
+
+    # 記録を取る
+    if len(reply_ids) == 3:
+        send_reply(reply_ids[0], "", "system", system_message)
+    send_reply(user_reply_id, last_reply_id, "user", message)
+    send_reply(assistant_reply_id, user_reply_id, "assistant", stream[1][-1][1])
+
+
+def retry(
+    history: list[tuple[str, str]],
+    system_message: str,
+    max_tokens: int,
+    temperature: float,
+    top_p: float,
+    top_k: int,
+    reply_ids: list[str],
+):
+    # 最後のメッセージを削除
+    last_conversation = history[-1]
+    user_message = last_conversation[0]
+    history = history[:-1]
+
+    user_reply_id = reply_ids[-2]
+    reply_ids = reply_ids[:-1]
+    assistant_reply_id = generate_unique_id()
+    reply_ids.append(assistant_reply_id)
+
+    for stream in generate(
+        user_message,
+        history,
+        system_message,
+        max_tokens,
+        temperature,
+        top_p,
+        top_k,
+    ):
+        yield *stream, reply_ids
+
+    # 記録を取る
+    send_reply(assistant_reply_id, user_reply_id, "assistant", stream[1][-1][1])
+
+
+def like_reponse(like_data: gr.LikeData, reply_ids: list[str]):
+    # print(like_data.index, like_data.value, like_data.liked)
+    assert isinstance(like_data.index, list)
+    # 評価を送信
+    send_score(reply_ids[like_data.index[0] + 1], 1 if like_data.liked else -1)
+
+
+def demo():
+    with gr.Blocks() as ui:
+
+        gr.Markdown(
+            """\
+# Tanuki 8B Instruct デモ
+モデル: https://huggingface.co/hatakeyama-llm-team/Tanuki-8B-Instruct
+
+アシスタントの回答が不適切だと思った場合は **低評価ボタンを押して低評価を送信**、同様に、回答が素晴らしいと思った場合は**高評価ボタンを押して高評価を送信**することで、モデルの改善に貢献できます。
+
+## 注意点
+**本デモに入力されたデータ・会話は匿名で全て記録されます**。これらのデータは Tanuki の学習に利用する可能性があります。そのため、**機密情報・個人情報を入力しないでください**。
+"""
+        )
+
+        reply_ids = gr.State(value=[generate_unique_id()])
+
+        chat_history = gr.Chatbot(value=[])
+
+        with gr.Row():
+            retry_btn = gr.Button(value="🔄 再生成", scale=1, size="sm")
+            clear_btn = gr.ClearButton(
+                components=[chat_history], value="🗑️ 削除", scale=1, size="sm"
+            )
+
+        with gr.Group():
+            with gr.Row():
+                input_text = gr.Textbox(
+                    value="",
+                    placeholder="質問を入力してください...",
+                    show_label=False,
+                    scale=8,
+                )
+                start_btn = gr.Button(
+                    value="送信",
+                    variant="primary",
+                    scale=1,
+                )
+            gr.Markdown(
+                value="※ 機密情報を入力しないでください。また、Tanuki は誤った情報を生成する可能性があります。"
+            )
+
+        with gr.Accordion(label="詳細設定", open=False):
+            system_prompt_text = gr.Textbox(
+                label="システムプロンプト",
+                value="以下は、タスクを説明する指示と、文脈のある入力の組み合わせです。要求を適切に満たす応答を書きなさい。",
+            )
+            max_new_tokens_slider = gr.Slider(
+                minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"
+            )
+            temperature_slider = gr.Slider(
+                minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"
+            )
+            top_p_slider = gr.Slider(
+                minimum=0.1,
+                maximum=1.0,
+                value=0.95,
+                step=0.05,
+                label="Top-p",
+            )
+            top_k_slider = gr.Slider(
+                minimum=1, maximum=2000, value=250, step=10, label="Top-k"
+            )
+
+        gr.Examples(
+            examples=[
+                ["たぬきってなんですか？"],
+                ["情けは人の為ならずとはどういう意味ですか？"],
+                ["まどマギで一番可愛いのは誰？"],
+                ["明晰夢とはなんですか？"],
+                [
+                    "シュレディンガー方程式とシュレディンガーの猫はどのような関係がありますか？"
+                ],
+            ],
+            inputs=[input_text],
+            cache_examples=False,
+        )
+
+        start_btn.click(
+            respond,
+            inputs=[
+                input_text,
+                chat_history,
+                system_prompt_text,
+                max_new_tokens_slider,
+                temperature_slider,
+                top_p_slider,
+                top_k_slider,
+                reply_ids,
+            ],
+            outputs=[input_text, chat_history, reply_ids],
+        )
+        input_text.submit(
+            respond,
+            inputs=[
+                input_text,
+                chat_history,
+                system_prompt_text,
+                max_new_tokens_slider,
+                temperature_slider,
+                top_p_slider,
+                top_k_slider,
+                reply_ids,
+            ],
+            outputs=[input_text, chat_history, reply_ids],
+        )
+        retry_btn.click(
+            retry,
+            inputs=[
+                chat_history,
+                system_prompt_text,
+                max_new_tokens_slider,
+                temperature_slider,
+                top_p_slider,
+                top_k_slider,
+                reply_ids,
+            ],
+            outputs=[input_text, chat_history, reply_ids],
+        )
+
+        # 評価されたら
+        chat_history.like(like_reponse, inputs=[reply_ids], outputs=None)
+
+        clear_btn.click(
+            lambda: [generate_unique_id()],  # system_message用のIDを生成
+            outputs=[reply_ids],
+        )
+
+    ui.launch()
 
 
 if __name__ == "__main__":
-    demo.launch()
+    demo()

requirements.txt

@@ -4,3 +4,5 @@ accelerate==0.30.1
 transformers==4.41.2
 spaces==0.28.3
 bitsandbytes==0.43.1
+dotenv
+requests