Update app.py

app.py

@@ -1,29 +1,125 @@
-import time
 import os
 import gradio as gr
 from text_generation import Client
 from conversation import get_default_conv_template
 from transformers import AutoTokenizer
+DESCRIPTION = """
+# Language Models for Taiwanese Culture
 
+Taiwan-LLaMa is a fine-tuned model specifically designed for traditional Chinese applications. It is built upon the LLaMa 2 architecture and includes a pretraining phase with over 5 billion tokens and fine-tuning with over 490k multi-turn conversational data in Traditional Chinese.
+
+## Key Features
+
+1. **Traditional Chinese Support**: The model is fine-tuned to understand and generate text in Traditional Chinese, making it suitable for Taiwanese culture and related applications.
+
+2. **Instruction-Tuned**: Further fine-tuned on conversational data to offer context-aware and instruction-following responses.
+
+3. **Performance on Vicuna Benchmark**: Taiwan-LLaMa's relative performance on Vicuna Benchmark is measured against models like GPT-4 and ChatGPT. It's particularly optimized for Taiwanese culture.
+
+4. **Flexible Customization**: Advanced options for controlling the model's behavior like system prompt, temperature, top-p, and top-k are available in the demo.
+
+## Model Versions
+
+Different versions of Taiwan-LLaMa are available:
+
+- **Taiwan-LLaMa v1.0 (This demo)**: Optimized for Taiwanese Culture
+- **Taiwan-LLaMa v0.9**: Partial instruction set
+- **Taiwan-LLaMa v0.0**: No Traditional Chinese pretraining
+
+The models can be accessed from the provided links in the Hugging Face repository.
+
+Try out the demo to interact with Taiwan-LLaMa and experience its capabilities in handling Traditional Chinese!
+"""
+
+LICENSE = """
+## Licenses
+
+- Code is licensed under Apache 2.0 License.
+- Models are licensed under the LLAMA 2 Community License.
+- By using this model, you agree to the terms and conditions specified in the license.
+- By using this demo, you agree to share your input utterances with us to improve the model.
+
+## Acknowledgements
+
+Taiwan-LLaMa project acknowledges the efforts of the [Meta LLaMa team](https://github.com/facebookresearch/llama) and [Vicuna team](https://github.com/lm-sys/FastChat) in democratizing large language models.
+"""
+
+DEFAULT_SYSTEM_PROMPT = "A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions. You are built by NTU Miulab by Yen-Ting Lin for research purpose."
 
 endpoint_url = os.environ.get("ENDPOINT_URL", "http://127.0.0.1:8080")
 client = Client(endpoint_url, timeout=120)
 eos_token = "</s>"
-max_prompt_length = 4096 - 512 - 10
+MAX_MAX_NEW_TOKENS = 1024
+DEFAULT_MAX_NEW_TOKENS = 1024
+
+max_prompt_length = 4096 - MAX_MAX_NEW_TOKENS - 10
 
 tokenizer = AutoTokenizer.from_pretrained("yentinglin/Taiwan-LLaMa-v1.0")
 
 with gr.Blocks() as demo:
+    gr.Markdown(DESCRIPTION)
+
     chatbot = gr.Chatbot()
-    msg = gr.Textbox()
-    clear = gr.Button("Clear")
+    with gr.Row():
+        msg = gr.Textbox(
+            container=False,
+            show_label=False,
+            placeholder='Type a message...',
+            scale=10,
+        )
+        submit_button = gr.Button('Submit',
+                                  variant='primary',
+                                  scale=1,
+                                  min_width=0)
+
+    with gr.Row():
+        retry_button = gr.Button('🔄  Retry', variant='secondary')
+        undo_button = gr.Button('↩️ Undo', variant='secondary')
+        clear = gr.Button('🗑️  Clear', variant='secondary')
+
+    saved_input = gr.State()
+
+    with gr.Accordion(label='Advanced options', open=False):
+        system_prompt = gr.Textbox(label='System prompt',
+                                   value=DEFAULT_SYSTEM_PROMPT,
+                                   lines=6)
+        max_new_tokens = gr.Slider(
+            label='Max new tokens',
+            minimum=1,
+            maximum=MAX_MAX_NEW_TOKENS,
+            step=1,
+            value=DEFAULT_MAX_NEW_TOKENS,
+        )
+        temperature = gr.Slider(
+            label='Temperature',
+            minimum=0.1,
+            maximum=1.0,
+            step=0.1,
+            value=0.7,
+        )
+        top_p = gr.Slider(
+            label='Top-p (nucleus sampling)',
+            minimum=0.05,
+            maximum=1.0,
+            step=0.05,
+            value=0.9,
+        )
+        top_k = gr.Slider(
+            label='Top-k',
+            minimum=1,
+            maximum=1000,
+            step=1,
+            value=50,
+        )
 
     def user(user_message, history):
         return "", history + [[user_message, None]]
 
-    def bot(history):
+
+    def bot(history, max_new_tokens, temperature, top_p, top_k, system_prompt):
         conv = get_default_conv_template("vicuna").copy()
         roles = {"human": conv.roles[0], "gpt": conv.roles[1]}  # map human to USER and gpt to ASSISTANT
+        conv.system = system_prompt
         for user, bot in history:
             conv.append_message(roles['human'], user)
             conv.append_message(roles["gpt"], bot)
@@ -31,123 +127,105 @@ with gr.Blocks() as demo:
         prompt_tokens = tokenizer.encode(msg)
         length_of_prompt = len(prompt_tokens)
         if length_of_prompt > max_prompt_length:
-            msg = tokenizer.decode(prompt_tokens[-max_prompt_length+1:])
+            msg = tokenizer.decode(prompt_tokens[-max_prompt_length + 1:])
 
         history[-1][1] = ""
         for response in client.generate_stream(
                 msg,
-                max_new_tokens=512,
+                max_new_tokens=max_new_tokens,
+                temperature=temperature,
+                top_p=top_p,
+                top_k=top_k,
         ):
             if not response.token.special:
                 character = response.token.text
                 history[-1][1] += character
                 yield history
 
+    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
+        fn=bot,
+        inputs=[
+            chatbot,
+            max_new_tokens,
+            temperature,
+            top_p,
+            top_k,
+            system_prompt,
+        ],
+        outputs=chatbot
+    )
+    submit_button.click(
+        user, [msg, chatbot], [msg, chatbot], queue=False
+    ).then(
+        fn=bot,
+        inputs=[
+            chatbot,
+            max_new_tokens,
+            temperature,
+            top_p,
+            top_k,
+            system_prompt,
+        ],
+        outputs=chatbot
+    )
 
-    def generate_response(history, max_new_token=512, top_p=0.9, temperature=0.8, do_sample=True):
-        conv = get_default_conv_template("vicuna").copy()
-        roles = {"human": conv.roles[0], "gpt": conv.roles[1]}  # map human to USER and gpt to ASSISTANT
-        for user, bot in history:
-            conv.append_message(roles['human'], user)
-            conv.append_message(roles["gpt"], bot)
-        msg = conv.get_prompt()
 
-        for response in client.generate_stream(
-                msg,
-                max_new_tokens=max_new_token,
-                top_p=top_p,
-                temperature=temperature,
-                do_sample=do_sample,
-        ):
-            history[-1][1] = ""
-            # if not response.token.special:
-            character = response.token.text
-            history[-1][1] += character
-            print(history[-1][1])
-            time.sleep(0.05)
-            yield history
+    def delete_prev_fn(
+            history: list[tuple[str, str]]) -> tuple[list[tuple[str, str]], str]:
+        try:
+            message, _ = history.pop()
+        except IndexError:
+            message = ''
+        return history, message or ''
 
 
-    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
-        bot, chatbot, chatbot
+    def display_input(message: str,
+                      history: list[tuple[str, str]]) -> list[tuple[str, str]]:
+        history.append((message, ''))
+        return history
+
+    retry_button.click(
+        fn=delete_prev_fn,
+        inputs=chatbot,
+        outputs=[chatbot, saved_input],
+        api_name=False,
+        queue=False,
+    ).then(
+        fn=display_input,
+        inputs=[saved_input, chatbot],
+        outputs=chatbot,
+        api_name=False,
+        queue=False,
+    ).then(
+        fn=bot,
+        inputs=[
+            chatbot,
+            max_new_tokens,
+            temperature,
+            top_p,
+            top_k,
+            system_prompt,
+        ],
+        outputs=chatbot,
     )
+
+    undo_button.click(
+        fn=delete_prev_fn,
+        inputs=chatbot,
+        outputs=[chatbot, saved_input],
+        api_name=False,
+        queue=False,
+    ).then(
+        fn=lambda x: x,
+        inputs=[saved_input],
+        outputs=msg,
+        api_name=False,
+        queue=False,
+    )
+
     clear.click(lambda: None, None, chatbot, queue=False)
-    
-demo.queue()
-demo.launch()
 
-#
-# with gr.Blocks() as demo:
-#     chatbot = gr.Chatbot()
-#     with gr.Row():
-#         with gr.Column(scale=4):
-#             with gr.Column(scale=12):
-#                 user_input = gr.Textbox(
-#                     show_label=False,
-#                     placeholder="Shift + Enter傳送...",
-#                     lines=10).style(
-#                     container=False)
-#             with gr.Column(min_width=32, scale=1):
-#                 submitBtn = gr.Button("Submit", variant="primary")
-#         with gr.Column(scale=1):
-#             emptyBtn = gr.Button("Clear History")
-#             max_new_token = gr.Slider(
-#                 1,
-#                 1024,
-#                 value=128,
-#                 step=1.0,
-#                 label="Maximum New Token Length",
-#                 interactive=True)
-#             top_p = gr.Slider(0, 1, value=0.9, step=0.01,
-#                               label="Top P", interactive=True)
-#             temperature = gr.Slider(
-#                 0,
-#                 1,
-#                 value=0.5,
-#                 step=0.01,
-#                 label="Temperature",
-#                 interactive=True)
-#             top_k = gr.Slider(1, 40, value=40, step=1,
-#                               label="Top K", interactive=True)
-#             do_sample = gr.Checkbox(
-#                 value=True,
-#                 label="Do Sample",
-#                 info="use random sample strategy",
-#                 interactive=True)
-#             repetition_penalty = gr.Slider(
-#                 1.0,
-#                 3.0,
-#                 value=1.1,
-#                 step=0.1,
-#                 label="Repetition Penalty",
-#                 interactive=True)
-#
-#     params = [user_input, chatbot]
-#     predict_params = [
-#         chatbot,
-#         max_new_token,
-#         top_p,
-#         temperature,
-#         top_k,
-#         do_sample,
-#         repetition_penalty]
-#
-#     submitBtn.click(
-#         generate_response,
-#         [user_input, max_new_token, top_p, top_k, temperature, do_sample, repetition_penalty],
-#         [chatbot],
-#         queue=False
-#     )
-#
-#     user_input.submit(
-#         generate_response,
-#         [user_input, max_new_token, top_p, top_k, temperature, do_sample, repetition_penalty],
-#         [chatbot],
-#         queue=False
-#     )
-#
-#     submitBtn.click(lambda: None, [], [user_input])
-#
-#     emptyBtn.click(lambda: chatbot.reset(), outputs=[chatbot], show_progress=True)
-#
-# demo.launch()
+    gr.Markdown(LICENSE)
+
+demo.queue(max_size=128)
+demo.launch()