Update app.py

app.py

@@ -33,7 +33,14 @@ import spaces
 load_dotenv()
 
 HF_API_KEY = os.getenv("HF_API_KEY")
-MODEL_NAME = "weblab-GENIAC/Tanuki-8B-dpo-v1.0"
+MODEL_NAME_MAP = {
+    "150m-instruct3": "llm-jp/llm-jp-3-150m-instruct3",
+    "440m-instruct3": "llm-jp/llm-jp-3-440m-instruct3",
+    "980m-instruct3": "llm-jp/llm-jp-3-980m-instruct3",
+    "1.8b-instruct3": "llm-jp/llm-jp-3-1.8b-instruct3",
+    "3.7b-instruct3": "llm-jp/llm-jp-3-3.7b-instruct3",
+    "13b-instruct3": "llm-jp/llm-jp-3-13b-instruct3",
+}
 
 quantization_config = BitsAndBytesConfig(
     load_in_4bit=True,
@@ -41,18 +48,25 @@ quantization_config = BitsAndBytesConfig(
     bnb_4bit_quant_type="nf4",
     bnb_4bit_use_double_quant=True,
 )
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME, quantization_config=quantization_config, device_map="auto", token=HF_API_KEY
-)
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, token=HF_API_KEY)
+MODELS = {
+    key: AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME, quantization_config=quantization_config, device_map="auto"
+    ) for key, value in MODEL_NAME_MAP.items()
+}
+TOKENIZERS = {
+    key: AutoTokenizer.from_pretrained(MODEL_NAME) for key, value in MODEL_NAME_MAP.items()
+    
+}
 
 print("Compiling model...")
-model = torch.compile(model)
+for key, model in MODELS:
+    MODELS[key] = torch.compile(model)
 print("Model compiled.")
 
 
-@spaces.GPU(duration=30)
+@spaces.GPU(duration=45)
 def generate(
+    model_name: str,
     message: str,
     history: list[tuple[str, str]],
     system_message: str,
@@ -74,12 +88,12 @@ def generate(
 
     messages.append({"role": "user", "content": message})
 
-    tokenized_input = tokenizer.apply_chat_template(
+    tokenized_input = TOKENIZERS[model_name].apply_chat_template(
         messages, add_generation_prompt=True, tokenize=True, return_tensors="pt"
     ).to(model.device)
 
     streamer = TextIteratorStreamer(
-        tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True
+        TOKENIZERS[model_name], timeout=10.0, skip_prompt=True, skip_special_tokens=True
     )
     generate_kwargs = dict(
         input_ids=tokenized_input,
@@ -91,7 +105,7 @@ def generate(
         top_p=float(top_p),
         num_beams=1,
     )
-    t = Thread(target=model.generate, kwargs=generate_kwargs)
+    t = Thread(target=MODELS[model_name].generate, kwargs=generate_kwargs)
     t.start()
 
     # 返す値を初期化
@@ -105,6 +119,7 @@ def generate(
 
 
 def respond(
+    model_name: str,
     message: str,
     history: list[tuple[str, str]],
     system_message: str,
@@ -115,6 +130,7 @@ def respond(
 ):
 
     for stream in generate(
+        model_name,
         message,
         history,
         system_message,
@@ -127,6 +143,7 @@ def respond(
 
 
 def retry(
+    model_name: str,
     history: list[tuple[str, str]],
     system_message: str,
     max_tokens: int,
@@ -140,6 +157,7 @@ def retry(
     history = history[:-1]
 
     for stream in generate(
+        model_name,
         user_message,
         history,
         system_message,
@@ -156,11 +174,13 @@ def demo():
 
         gr.Markdown(
             """\
-# weblab-GENIAC/Tanuki-8B-dpo-v1.0 デモ
-モデル: https://huggingface.co/weblab-GENIAC/Tanuki-8B-dpo-v1.0
+# llm-jp/llm-jp-3 instruct3 モデルデモ
+コレクション: https://huggingface.co/collections/llm-jp/llm-jp-3-fine-tuned-models-672c621db852a01eae939731
 """
         )
 
+        model_name_dropdown = gr.Dropdown(label="モデル", choices=list(MODELS.keys()), value=list(MODELS.keys())[0])
+
         chat_history = gr.Chatbot(value=[])
 
         with gr.Row():
@@ -183,7 +203,7 @@ def demo():
                     scale=2,
                 )
             gr.Markdown(
-                value="※ Tanuki は誤った情報を生成する可能性があります。"
+                value="※ 誤った情報を生成する可能性があります。"
             )
 
         with gr.Accordion(label="詳細設定", open=False):
@@ -195,7 +215,7 @@ def demo():
                 minimum=1, maximum=2048, value=256, step=1, label="Max new tokens"
             )
             temperature_slider = gr.Slider(
-                minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"
+                minimum=0.1, maximum=1.0, value=0.7, step=0.1, label="Temperature"
             )
             top_p_slider = gr.Slider(
                 minimum=0.1,
@@ -210,7 +230,6 @@ def demo():
 
         gr.Examples(
             examples=[
-                ["たぬきってなんですか？"],
                 ["情けは人の為ならずとはどういう意味ですか？"],
                 ["まどマギで一番可愛いのは誰？"],
             ],
@@ -218,22 +237,11 @@ def demo():
             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,
-            ],
-            outputs=[input_text, chat_history],
-        )
-        input_text.submit(
-            respond,
+        gr.on(
+            triggers=[start_btn.click, input_text.submit],
+            fn=respond,
             inputs=[
+                model_name_dropdown,
                 input_text,
                 chat_history,
                 system_prompt_text,
@@ -247,6 +255,7 @@ def demo():
         retry_btn.click(
             retry,
             inputs=[
+                model_name_dropdown,
                 chat_history,
                 system_prompt_text,
                 max_new_tokens_slider,