Update app.py

app.py

@@ -1358,8 +1358,7 @@ def create_free_form_generation_demo():
     presence_penalty = PRESENCE_PENALTY
 
     introduction = """
-## Free-form: 
-Put any context string (like few-shot prompts) and get the model to generate.
+### Free-form | Put any context string (like few-shot prompts)
     """
 
     with gr.Blocks() as demo_free_form:
@@ -1391,6 +1390,78 @@ Put any context string (like few-shot prompts) and get the model to generate.
 
 
 
+def create_file_upload_demo():
+    temperature = TEMPERATURE
+    frequence_penalty = FREQUENCE_PENALTY
+    presence_penalty = PRESENCE_PENALTY
+    max_tokens = MAX_TOKENS
+    demo_file_upload = gr.Interface(
+        batch_inference,
+        inputs=[
+            gr.File(file_count='single', file_types=['json']),
+            gr.Radio(["chat", "few-shot"], value='chat', label="Chat or Few-shot mode", info="Chat's output more user-friendly, Few-shot's output more consistent with few-shot patterns."),
+            gr.Number(value=temperature, label='Temperature', info="Higher -> more random"), 
+            gr.Number(value=max_tokens, label='Max tokens', info='Increase if want more generation'), 
+            gr.Number(value=frequence_penalty, label='Frequency penalty', info='> 0 encourage new tokens over repeated tokens'), 
+            gr.Number(value=presence_penalty, label='Presence penalty', info='> 0 encourage new tokens, < 0 encourage existing tokens'), 
+            gr.Textbox(value="<s>,</s>,<|im_start|>", label='Stop strings', info='Comma-separated string to stop generation only in FEW-SHOT mode', lines=1),
+            gr.Number(value=0, label='current_time', visible=False), 
+        ],
+        outputs=[
+            # "file",
+            gr.File(label="Generated file"),
+            # "json"
+            gr.JSON(label='Example outputs (display 2 samples)')
+        ],
+        description=FILE_UPLOAD_DESCRIPTION,
+        allow_flagging=False,
+        examples=[
+            ["upload_chat.json", "chat", 0.2, 1024, 0.5, 0, "<s>,</s>,<|im_start|>"],
+            ["upload_few_shot.json", "few-shot", 0.2, 128, 0.5, 0, "<s>,</s>,<|im_start|>,\\n"]
+        ],
+        cache_examples=False,
+    )
+    return demo_file_upload
+
+
+def create_chat_demo(title=None, description=None):
+    sys_prompt = SYSTEM_PROMPT_1
+    max_tokens = MAX_TOKENS
+    temperature = TEMPERATURE
+    frequence_penalty = FREQUENCE_PENALTY
+    presence_penalty = PRESENCE_PENALTY
+
+    demo_chat = gr.ChatInterface(
+        chat_response_stream_multiturn,
+        chatbot=ChatBot(
+            label=MODEL_NAME,
+            bubble_full_width=False,
+            latex_delimiters=[
+                { "left": "$", "right": "$", "display": False},
+                { "left": "$$", "right": "$$", "display": True},
+            ],
+            show_copy_button=True,
+        ),
+        textbox=gr.Textbox(placeholder='Type message', lines=4, max_lines=128, min_width=200),
+        submit_btn=gr.Button(value='Submit', variant="primary", scale=0),
+        # ! consider preventing the stop button
+        # stop_btn=None,
+        title=title,
+        description=description,
+        additional_inputs=[
+            gr.Number(value=temperature, label='Temperature (higher -> more random)'), 
+            gr.Number(value=max_tokens, label='Max generated tokens (increase if want more generation)'), 
+            gr.Number(value=frequence_penalty, label='Frequency penalty (> 0 encourage new tokens over repeated tokens)'), 
+            gr.Number(value=presence_penalty, label='Presence penalty (> 0 encourage new tokens, < 0 encourage existing tokens)'), 
+            gr.Number(value=0, label='current_time', visible=False), 
+            # ! Remove the system prompt textbox to avoid jailbreaking
+            gr.Textbox(value=sys_prompt, label='System prompt', lines=8, interactive=False)
+        ], 
+        examples=CHAT_EXAMPLES,
+        cache_examples=False
+    )
+    return demo_chat
+
 
 
 def launch_demo():
@@ -1486,64 +1557,11 @@ def launch_demo():
 
     if ENABLE_BATCH_INFER:
 
-        demo_file_upload = gr.Interface(
-            batch_inference,
-            inputs=[
-                gr.File(file_count='single', file_types=['json']),
-                gr.Radio(["chat", "few-shot"], value='chat', label="Chat or Few-shot mode", info="Chat's output more user-friendly, Few-shot's output more consistent with few-shot patterns."),
-                gr.Number(value=temperature, label='Temperature', info="Higher -> more random"), 
-                gr.Number(value=max_tokens, label='Max tokens', info='Increase if want more generation'), 
-                gr.Number(value=frequence_penalty, label='Frequency penalty', info='> 0 encourage new tokens over repeated tokens'), 
-                gr.Number(value=presence_penalty, label='Presence penalty', info='> 0 encourage new tokens, < 0 encourage existing tokens'), 
-                gr.Textbox(value="<s>,</s>,<|im_start|>", label='Stop strings', info='Comma-separated string to stop generation only in FEW-SHOT mode', lines=1),
-                gr.Number(value=0, label='current_time', visible=False), 
-            ],
-            outputs=[
-                # "file",
-                gr.File(label="Generated file"),
-                # "json"
-                gr.JSON(label='Example outputs (display 2 samples)')
-            ],
-            description=FILE_UPLOAD_DESCRIPTION,
-            allow_flagging=False,
-            examples=[
-                ["upload_chat.json", "chat", 0.2, 1024, 0.5, 0, "<s>,</s>,<|im_start|>"],
-                ["upload_few_shot.json", "few-shot", 0.2, 128, 0.5, 0, "<s>,</s>,<|im_start|>,\\n"]
-            ],
-            cache_examples=False,
-        )
+        demo_file_upload = create_file_upload_demo()
 
         demo_free_form = create_free_form_generation_demo()
         
-        demo_chat = gr.ChatInterface(
-            response_fn,
-            chatbot=ChatBot(
-                label=MODEL_NAME,
-                bubble_full_width=False,
-                latex_delimiters=[
-                    { "left": "$", "right": "$", "display": False},
-                    { "left": "$$", "right": "$$", "display": True},
-                ],
-                show_copy_button=True,
-            ),
-            textbox=gr.Textbox(placeholder='Type message', lines=4, max_lines=128, min_width=200),
-            submit_btn=gr.Button(value='Submit', variant="primary", scale=0),
-            # ! consider preventing the stop button
-            # stop_btn=None,
-            # title=f"{model_title}",
-            # description=f"{model_desc}",
-            additional_inputs=[
-                gr.Number(value=temperature, label='Temperature (higher -> more random)'), 
-                gr.Number(value=max_tokens, label='Max generated tokens (increase if want more generation)'), 
-                gr.Number(value=frequence_penalty, label='Frequency penalty (> 0 encourage new tokens over repeated tokens)'), 
-                gr.Number(value=presence_penalty, label='Presence penalty (> 0 encourage new tokens, < 0 encourage existing tokens)'), 
-                gr.Number(value=0, label='current_time', visible=False), 
-                # ! Remove the system prompt textbox to avoid jailbreaking
-                # gr.Textbox(value=sys_prompt, label='System prompt', lines=8)
-            ], 
-            examples=CHAT_EXAMPLES,
-            cache_examples=False
-        )
+        demo_chat = create_chat_demo()
         descriptions = model_desc
         if DISPLAY_MODEL_PATH:
             descriptions += f"<br> {path_markdown.format(model_path=model_path)}"
@@ -1554,80 +1572,33 @@ def launch_demo():
             title=f"{model_title}",
             description=descriptions,
         )
-        demo.title = MODEL_NAME
-        callback = None
-        with demo:
-            if DATA_SET_REPO_PATH != "":
-                try:
-                    from performance_plot import attach_plot_to_demo
-                    attach_plot_to_demo(demo)
-                except Exception as e:
-                    print(f'Fail to load DEMO plot: {str(e)}')
-            
-            gr.Markdown(cite_markdown)
-            # if DISPLAY_MODEL_PATH:
-            #     gr.Markdown(path_markdown.format(model_path=model_path))
-            
-            if ENABLE_AGREE_POPUP:
-                demo.load(None, None, None, _js=AGREE_POP_SCRIPTS)
-        
-        demo.queue(api_open=False)
-        # demo.launch(server_port=PORT, show_api=False, allowed_paths=["seal_logo.png"])
-        # demo.launch(show_api=False, allowed_paths=["seal_logo.png"])
     else:
         descriptions = model_desc
         if DISPLAY_MODEL_PATH:
             descriptions += f"<br> {path_markdown.format(model_path=model_path)}"
-        demo = gr.ChatInterface(
-            response_fn,
-            chatbot=ChatBot(
-                label=MODEL_NAME,
-                bubble_full_width=False,
-                latex_delimiters=[
-                    { "left": "$", "right": "$", "display": False},
-                    { "left": "$$", "right": "$$", "display": True},
-                ],
-                show_copy_button=True,
-            ),
-            textbox=gr.Textbox(placeholder='Type message', lines=4, max_lines=128, min_width=200),
-            submit_btn=gr.Button(value='Submit', variant="primary", scale=0),
-            # ! consider preventing the stop button
-            # stop_btn=None,
-            title=f"{model_title}",
-            description=descriptions,
-            additional_inputs=[
-                gr.Number(value=temperature, label='Temperature (higher -> more random)'), 
-                gr.Number(value=max_tokens, label='Max generated tokens (increase if want more generation)'), 
-                gr.Number(value=frequence_penalty, label='Frequency penalty (> 0 encourage new tokens over repeated tokens)'), 
-                gr.Number(value=presence_penalty, label='Presence penalty (> 0 encourage new tokens, < 0 encourage existing tokens)'), 
-                gr.Number(value=0, label='current_time', visible=False), 
-                # ! Remove the system prompt textbox to avoid jailbreaking
-                # gr.Textbox(value=sys_prompt, label='System prompt', lines=8)
-            ], 
-            examples=CHAT_EXAMPLES,
-            cache_examples=False
-        )
-        demo.title = MODEL_NAME
-        with demo:
-            gr.Markdown(cite_markdown)
-            if DISPLAY_MODEL_PATH:
-                gr.Markdown(path_markdown.format(model_path=model_path))
-            
-            if ENABLE_AGREE_POPUP:
-                demo.load(None, None, None, _js=AGREE_POP_SCRIPTS)
+
+        demo = create_chat_demo(title=f"{model_title}", description=descriptions)
+    demo.title = MODEL_NAME
+    
+    with demo:
+        if DATA_SET_REPO_PATH != "":
+            try:
+                from performance_plot import attach_plot_to_demo
+                attach_plot_to_demo(demo)
+            except Exception as e:
+                print(f'Fail to load DEMO plot: {str(e)}')
         
+        gr.Markdown(cite_markdown)
+        if DISPLAY_MODEL_PATH:
+            gr.Markdown(path_markdown.format(model_path=model_path))
+        
+        if ENABLE_AGREE_POPUP:
+            demo.load(None, None, None, _js=AGREE_POP_SCRIPTS)
 
-        demo.queue(api_open=False)
-        #demo.launch(server_port=PORT, show_api=False, allowed_paths=["seal_logo.png"])
-    # demo.launch(show_api=False, allowed_paths=["seal_logo.png"])
+    demo.queue(api_open=False)
     return demo
 
 
-# def main():
-
-#     launch()
-
-
 if __name__ == "__main__":
     demo = launch_demo()
     demo.launch(show_api=False, allowed_paths=["seal_logo.png"])