Update app.py

app.py

@@ -50,36 +50,43 @@ no_change_btn = gr.Button.update()
 enable_btn = gr.Button.update(interactive=True)
 disable_btn = gr.Button.update(interactive=False)
 
-enable_moderation = False
-concurrency_count = 10
-model_list_mode = 'reload'
-allow_running = 1
-
-midware_url = "http://119.8.43.169:8080/api/v1/chat/models"
-worker_addr = 'http://119.8.43.169:8080/api/v1/chat'
-chat_token = 'abc'
-ft_list_job_url = "http://119.8.43.169:8080/api/v1/job"
-ft_submit_job_url = "http://119.8.43.169:8080/api/v1/job"
-ft_remove_job_url = "http://119.8.43.169:8080/api/v1/job/"
-ft_console_log_url = "ws://119.8.43.169:8080/api/v1/log/"
-
+enable_moderation = True if os.environ.get('enable_moderation', default='False')=="True" else False
+concurrency_count = int(os.environ.get('concurrency_count', default='10'))
+model_list_mode = os.environ.get('model_list_mode', default='reload')
+
+midware_url = os.environ.get('midware_url', default='')
+preset_token = os.environ.get('preset_token', default='')
+worker_addr = os.environ.get('worker_addr', default='')
+
+allow_running = int(os.environ.get('allow_running', default='1'))
+ft_list_job_url = os.environ.get('ft_list_job_url', default='')
+ft_submit_job_url = os.environ.get('ft_submit_job_url', default='')
+ft_remove_job_url = os.environ.get('ft_remove_job_url', default='')
+ft_console_log_url = os.environ.get('ft_console_log_url', default='')
+
+dataset_sample = {
+    "english": {
+        "train": ["abcdef"],
+        "valid": ["zxcvbn"]
+    },
+    "cat": {
+        "train": ["aaaaaa"],
+        "valid": ["bbbbbb"]
+    }
+}
 
-# enable_moderation = True if os.environ.get('enable_moderation', default='False')=="True" else False
-# concurrency_count = int(os.environ.get('concurrency_count', default='10'))
-# model_list_mode = os.environ.get('model_list_mode', default='reload')
 
-# midware_url = os.environ.get('midware_url', default='')
-# chat_token = os.environ.get('chat_token', default='')
-# worker_addr = os.environ.get('worker_addr', default='')
+dataset_to_midware_name = {
+    "english": "english",
+    "cat": "cat",
+    "dog": "dog",
+    "bird": "bird"
+}
 
-# allow_running = int(os.environ.get('allow_running', default='1'))
-# ft_list_job_url = os.environ.get('ft_list_job_url', default='')
-# ft_submit_job_url = os.environ.get('ft_submit_job_url', default='')
-# ft_remove_job_url = os.environ.get('ft_remove_job_url', default='')
-# ft_console_log_url = os.environ.get('ft_console_log_url', default='')
+hps_keys = ["epochs", "train_batch_size", "eval_batch_size", "gradient_accumulation_steps", "learning_rate", "weight_decay", "model_max_length"]
 
 
-headers = {"User-Agent": "FastChat Client", "PRIVATE-TOKEN": chat_token}
+headers = {"User-Agent": "FastChat Client", "PRIVATE-TOKEN": preset_token}
 
 learn_more_md = """
 ### License
@@ -155,14 +162,20 @@ def get_conv_log_filename():
 
 
 def get_model_list(midware_url):
+    setted_model_order = {
+        "vicuna-7b-v1.5-16k": 10,
+        "vicuna-13b-v1.5": 90,
+    }
     try:
-        ret = requests.get(midware_url, headers={"PRIVATE-TOKEN": chat_token}, timeout=5)
-        models = ret.json()["data"]
+        ret = requests.get(midware_url, headers={"PRIVATE-TOKEN": preset_token}, timeout=5)
+        if "code" in ret.json() and "invalid" in ret.json()["code"]:
+            gr.Warning("Invalid preset token.")
+            models = ["CANNOT GET MODEL"]
+        else:
+            models = ret.json()["data"]
     except requests.exceptions.RequestException:
         models = ["CANNOT GET MODEL"]
-
-    priority = {k: f"___{i:02d}" for i, k in enumerate(model_info)}
-    models.sort(key=lambda x: priority.get(x, x))
+    models = sorted(models, key=lambda x: setted_model_order.get(x, 100))
     logger.info(f"Models: {models}")
     return models
 
@@ -202,47 +215,16 @@ def load_demo(url_params, request: gr.Request):
 
     return load_demo_single(models, url_params)
 
-
-def vote_last_response(state, vote_type, model_selector, request: gr.Request):
-    with open(get_conv_log_filename(), "a") as fout:
-        data = {
-            "tstamp": round(time.time(), 4),
-            "type": vote_type,
-            "model": model_selector,
-            "state": state.dict(),
-            "ip": request.client.host,
-        }
-        fout.write(json.dumps(data) + "\n")
-
-
-def upvote_last_response(state, model_selector, request: gr.Request):
-    logger.info(f"upvote. ip: {request.client.host}")
-    vote_last_response(state, "upvote", model_selector, request)
-    return ("",) + (disable_btn,) * 3
-
-
-def downvote_last_response(state, model_selector, request: gr.Request):
-    logger.info(f"downvote. ip: {request.client.host}")
-    vote_last_response(state, "downvote", model_selector, request)
-    return ("",) + (disable_btn,) * 3
-
-
-def flag_last_response(state, model_selector, request: gr.Request):
-    logger.info(f"flag. ip: {request.client.host}")
-    vote_last_response(state, "flag", model_selector, request)
-    return ("",) + (disable_btn,) * 3
-
-
 def regenerate(state, request: gr.Request):
     logger.info(f"regenerate. ip: {request.client.host}")
     state.conv.update_last_message(None)
-    return (state, state.to_gradio_chatbot(), "") + (disable_btn,) * 5
+    return (state, state.to_gradio_chatbot(), "") + (disable_btn,) * 2
 
 
 def clear_history(request: gr.Request):
     logger.info(f"clear_history. ip: {request.client.host}")
     state = None
-    return (state, [], "") + (disable_btn,) * 5
+    return (state, [], "") + (disable_btn,) * 2
 
 
 def add_text(state, model_selector, text, request: gr.Request):
@@ -254,12 +236,12 @@ def add_text(state, model_selector, text, request: gr.Request):
 
     if len(text) <= 0:
         state.skip_next = True
-        return (state, state.to_gradio_chatbot(), "") + (no_change_btn,) * 5
+        return (state, state.to_gradio_chatbot(), "") + (no_change_btn,) * 2
 
     if ip_expiration_dict[ip] < time.time():
         logger.info(f"inactive. ip: {request.client.host}. text: {text}")
         state.skip_next = True
-        return (state, state.to_gradio_chatbot(), INACTIVE_MSG) + (no_change_btn,) * 5
+        return (state, state.to_gradio_chatbot(), INACTIVE_MSG) + (no_change_btn,) * 2
 
     if enable_moderation:
         flagged = violates_moderation(text)
@@ -268,7 +250,7 @@ def add_text(state, model_selector, text, request: gr.Request):
             state.skip_next = True
             return (state, state.to_gradio_chatbot(), MODERATION_MSG) + (
                 no_change_btn,
-            ) * 5
+            ) * 2
 
     conv = state.conv
     if (len(conv.messages) - conv.offset) // 2 >= CONVERSATION_TURN_LIMIT:
@@ -276,12 +258,12 @@ def add_text(state, model_selector, text, request: gr.Request):
         state.skip_next = True
         return (state, state.to_gradio_chatbot(), CONVERSATION_LIMIT_MSG) + (
             no_change_btn,
-        ) * 5
+        ) * 2
 
     text = text[:INPUT_CHAR_LEN_LIMIT]  # Hard cut-off
     conv.append_message(conv.roles[0], text)
     conv.append_message(conv.roles[1], None)
-    return (state, state.to_gradio_chatbot(), "") + (disable_btn,) * 5
+    return (state, state.to_gradio_chatbot(), "") + (disable_btn,) * 2
 
 
 def post_process_code(code):
@@ -343,7 +325,7 @@ def bot_response(state, temperature, top_p, max_new_tokens, request: gr.Request)
     if state.skip_next:
         # This generate call is skipped due to invalid inputs
         state.skip_next = False
-        yield (state, state.to_gradio_chatbot()) + (no_change_btn,) * 5
+        yield (state, state.to_gradio_chatbot()) + (no_change_btn,) * 2
         return
 
     conv, model_name = state.conv, state.model_name
@@ -370,9 +352,6 @@ def bot_response(state, temperature, top_p, max_new_tokens, request: gr.Request)
             yield (
                 state,
                 state.to_gradio_chatbot(),
-                disable_btn,
-                disable_btn,
-                disable_btn,
                 enable_btn,
                 enable_btn,
             )
@@ -400,7 +379,7 @@ def bot_response(state, temperature, top_p, max_new_tokens, request: gr.Request)
         )
 
     conv.update_last_message("▌")
-    yield (state, state.to_gradio_chatbot()) + (disable_btn,) * 5
+    yield (state, state.to_gradio_chatbot()) + (disable_btn,) * 2
 
     try:
         for data in stream_iter:
@@ -410,14 +389,11 @@ def bot_response(state, temperature, top_p, max_new_tokens, request: gr.Request)
                     output = post_process_code(output)
                 output = str_filter(output)
                 conv.update_last_message(output + "▌")
-                yield (state, state.to_gradio_chatbot()) + (disable_btn,) * 5
+                yield (state, state.to_gradio_chatbot()) + (disable_btn,) * 2
             else:
                 output = data["text"] + f"\n\n(error_code: {data['error_code']})"
                 conv.update_last_message(output)
                 yield (state, state.to_gradio_chatbot()) + (
-                    disable_btn,
-                    disable_btn,
-                    disable_btn,
                     enable_btn,
                     enable_btn,
                 )
@@ -429,9 +405,6 @@ def bot_response(state, temperature, top_p, max_new_tokens, request: gr.Request)
             f"(error_code: {ErrorCode.GRADIO_REQUEST_ERROR}, {e})"
         )
         yield (state, state.to_gradio_chatbot()) + (
-            disable_btn,
-            disable_btn,
-            disable_btn,
             enable_btn,
             enable_btn,
         )
@@ -442,9 +415,6 @@ def bot_response(state, temperature, top_p, max_new_tokens, request: gr.Request)
             f"(error_code: {ErrorCode.GRADIO_STREAM_UNKNOWN_ERROR}, {e})"
         )
         yield (state, state.to_gradio_chatbot()) + (
-            disable_btn,
-            disable_btn,
-            disable_btn,
             enable_btn,
             enable_btn,
         )
@@ -452,7 +422,7 @@ def bot_response(state, temperature, top_p, max_new_tokens, request: gr.Request)
 
     # Delete "▌"
     conv.update_last_message(conv.messages[-1][-1][:-1])
-    yield (state, state.to_gradio_chatbot()) + (enable_btn,) * 5
+    yield (state, state.to_gradio_chatbot()) + (enable_btn,) * 2
 
     finish_tstamp = time.time()
     logger.info(f"{output}")
@@ -532,7 +502,6 @@ def build_single_model_ui(models, add_promotion_links=False):
     with gr.Column():
         with gr.Tab("🧠 模型对话 Dialog"):
             state = gr.State()
-
             with gr.Row(elem_id="model_selector_row"):
                 model_selector = gr.Dropdown(
                     choices=models,
@@ -560,12 +529,12 @@ def build_single_model_ui(models, add_promotion_links=False):
                     send_btn = gr.Button(value="Send", visible=False)
 
             with gr.Row(visible=False) as button_row:
-                upvote_btn = gr.Button(value="👍  Upvote", interactive=False)
-                downvote_btn = gr.Button(value="👎  Downvote", interactive=False)
-                flag_btn = gr.Button(value="⚠️  Flag", interactive=False)
                 regenerate_btn = gr.Button(value="🔄  Regenerate", interactive=False)
                 clear_btn = gr.Button(value="🗑️  Clear history", interactive=False)
-
+            gr.Examples(
+                examples=["如何变得富有？", "你能用Python写一段快速排序吗？", "How to be rich?", "Can you write a quicksort code in Python?"],
+                inputs=textbox,
+            )
             with gr.Accordion("Parameters", open=False, visible=False) as parameter_row:
                 temperature = gr.Slider(
                     minimum=0.0,
@@ -595,22 +564,7 @@ def build_single_model_ui(models, add_promotion_links=False):
             gr.Markdown(learn_more_md)
 
             # Register listeners
-            btn_list = [upvote_btn, downvote_btn, flag_btn, regenerate_btn, clear_btn]
-            upvote_btn.click(
-                upvote_last_response,
-                [state, model_selector],
-                [textbox, upvote_btn, downvote_btn, flag_btn],
-            )
-            downvote_btn.click(
-                downvote_last_response,
-                [state, model_selector],
-                [textbox, upvote_btn, downvote_btn, flag_btn],
-            )
-            flag_btn.click(
-                flag_last_response,
-                [state, model_selector],
-                [textbox, upvote_btn, downvote_btn, flag_btn],
-            )
+            btn_list = [regenerate_btn, clear_btn]
             regenerate_btn.click(regenerate, state, [state, chatbot, textbox] + btn_list).then(
                 bot_response,
                 [state, temperature, top_p, max_output_tokens],
@@ -634,8 +588,6 @@ def build_single_model_ui(models, add_promotion_links=False):
                 [state, temperature, top_p, max_output_tokens],
                 [state, chatbot] + btn_list,
             )
-        
-            
     return state, model_selector, chatbot, textbox, send_btn, button_row, parameter_row
 
 
@@ -643,16 +595,22 @@ def ft_get_job_data():
     running = 0
     res_lst = []
     try:
-        r = requests.get(ft_list_job_url, timeout=8)
+        r = requests.get(ft_list_job_url, headers={"PRIVATE-TOKEN": preset_token}, timeout=8)
+        if "code" in r.json() and "invalid" in r.json()["code"]:
+            gr.Warning("Invalid preset token.")
+            return res_lst, running
         for d in r.json():
             if isinstance(d['status'], str) and d['status'].lower() == "running":
-                running += 1 
-            res_lst.append([d['jobName'], d['username'], d['created_at'], d['model'], d['dataset'], d['status'], json.dumps(d['parameter'])])
+                running += 1
+            hps = dict()
+            for key in hps_keys:
+                if key in d['parameter']:
+                    hps[key] = d['parameter'][key]
+            res_lst.append([d['jobName'], d['username'], d['created_at'], d['model'], d['dataset'], d['status'], json.dumps(hps)])
         res_lst = sorted(res_lst,key=(lambda x:x[2]), reverse=True)
         res_lst = sorted(res_lst,key=(lambda x:x[5]), reverse=True)
     except requests.exceptions.RequestException:
         logger.info(f"Get job list fail")
-        res_lst.append([])
     return res_lst, running
 
 
@@ -668,7 +626,7 @@ def ft_cease_click(ft_console):
 def console_generator(addr, sleep_time):
     total_str = ""
     ws = websocket.WebSocket()
-    ws.connect(addr)
+    ws.connect(addr, header={"PRIVATE-TOKEN": preset_token})
     while True:
         try:
             new_str = ws.recv()
@@ -691,7 +649,7 @@ def ft_submit_click(ft_latest_running_cnt, ft_user_name, ft_model, ft_dataset_na
         res_lst, running = ft_get_job_data()
         return res_lst, running, no_change_textbox
     if ft_latest_running_cnt < int(allow_running):
-        midware_header = {'Content-Type': 'application/json'}
+        midware_header = {"FINETUNE-SECRET": ft_token, "PRIVATE-TOKEN": preset_token}
         hps_json = {
             "epochs": str(ft_epochs),
             "train_batch_size": str(ft_train_batch_size),
@@ -702,10 +660,9 @@ def ft_submit_click(ft_latest_running_cnt, ft_user_name, ft_model, ft_dataset_na
             "model_max_length": str(ft_model_max_length)
         }
         json_data = {
-            "dataset": ft_dataset_name,
+            "dataset": dataset_to_midware_name[ft_dataset_name],
             "model": ft_model,
             "parameter": hps_json,
-            "secret": ft_token,
             "username": ft_user_name
         }
         try:
@@ -737,7 +694,7 @@ def ft_show_click(ft_selected_row_data):
 def ft_remove_click(ft_selected_row_data, ft_token):
     status = ft_selected_row_data[5]
     if isinstance(status, str) and status.lower() == "running":
-        r = requests.delete(ft_remove_job_url + ft_selected_row_data[0], json={"secret": ft_token})
+        r = requests.delete(ft_remove_job_url + ft_selected_row_data[0], headers={'FINETUNE-SECRET': ft_token, "PRIVATE-TOKEN": preset_token})
         if r.status_code == 200:
             gr.Info("Remove success.")
         else:
@@ -760,6 +717,13 @@ def ft_jobs_info_select(ft_jobs_info, evt: gr.SelectData):
         return [selected_row, no_change_dropdown, no_change_dropdown, no_change_slider, no_change_textbox, no_change_textbox, no_change_textbox, no_change_textbox, no_change_textbox, no_change_textbox]
 
 
+def ft_dataset_preview_click(ft_dataset_name):
+    value = dataset_sample.get(ft_dataset_name, {})
+    return gr.JSON.update(value=value, visible=True)
+
+def ft_hide_dataset_click():
+    return gr.JSON.update(visible=False)
+
 def build_demo(models):
     with gr.Blocks(
         title="Vicuna Test",