Update app.py

app.py

@@ -5,7 +5,7 @@ except:
 
     print("Installing flash-attn...")
     subprocess.run(
-        "pip install flash-attn --no-build-isolation",
+        "uv install --system flash-attn --no-build-isolation",
         env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
         shell=True,
     )
@@ -35,10 +35,7 @@ import spaces
 load_dotenv()
 
 HF_API_KEY = os.getenv("HF_API_KEY")
-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"
-
+MODEL_NAME = "weblab-GENIAC/Tanuki-8B-dpo-v1.0"
 
 quantization_config = BitsAndBytesConfig(
     load_in_4bit=True,
@@ -56,55 +53,7 @@ model = torch.compile(model)
 print("Model compiled.")
 
 
-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)
+@spaces.GPU(duration=30)
 def generate(
     message: str,
     history: list[tuple[str, str]],
@@ -165,16 +114,7 @@ def respond(
     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,
@@ -185,13 +125,7 @@ def respond(
         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])
+        yield *stream
 
 
 def retry(
@@ -201,18 +135,12 @@ def retry(
     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,
@@ -222,17 +150,7 @@ def retry(
         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)
+        yield *stream
 
 
 def demo():
@@ -240,18 +158,11 @@ def demo():
 
         gr.Markdown(
             """\
-# Tanuki 8B Instruct デモ
-モデル: https://huggingface.co/hatakeyama-llm-team/Tanuki-8B-Instruct
-
-アシスタントの回答が不適切だと思った場合は **低評価ボタンを押して低評価を送信**、同様に、回答が素晴らしいと思った場合は**高評価ボタンを押して高評価を送信**することで、モデルの改善に貢献できます。
-
-## 注意点
-**本デモに入力されたデータ・会話は匿名で全て記録されます**。これらのデータは Tanuki の学習に利用する可能性があります。そのため、**機密情報・個人情報を入力しないでください**。
+# weblab-GENIAC/Tanuki-8B-dpo-v1.0 デモ
+モデル: https://huggingface.co/weblab-GENIAC/Tanuki-8B-dpo-v1.0
 """
         )
 
-        reply_ids = gr.State(value=[generate_unique_id()])
-
         chat_history = gr.Chatbot(value=[])
 
         with gr.Row():
@@ -304,10 +215,6 @@ def demo():
                 ["たぬきってなんですか？"],
                 ["情けは人の為ならずとはどういう意味ですか？"],
                 ["まどマギで一番可愛いのは誰？"],
-                ["明晰夢とはなんですか？"],
-                [
-                    "シュレディンガー方程式とシュレディンガーの猫はどのような関係があ��ますか？"
-                ],
             ],
             inputs=[input_text],
             cache_examples=False,
@@ -323,9 +230,8 @@ def demo():
                 temperature_slider,
                 top_p_slider,
                 top_k_slider,
-                reply_ids,
             ],
-            outputs=[input_text, chat_history, reply_ids],
+            outputs=[input_text, chat_history],
         )
         input_text.submit(
             respond,
@@ -337,9 +243,8 @@ def demo():
                 temperature_slider,
                 top_p_slider,
                 top_k_slider,
-                reply_ids,
             ],
-            outputs=[input_text, chat_history, reply_ids],
+            outputs=[input_text, chat_history],
         )
         retry_btn.click(
             retry,
@@ -350,17 +255,8 @@ def demo():
                 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],
+            outputs=[input_text, chat_history],
         )
 
     ui.launch()