Update app.py

app.py

@@ -17,8 +17,21 @@ generator = get_pretrained_models("7B", "tokenizer", local_rank, world_size)
 history = []
 simple_history = []
 
-def chat(user_input):
-    bot_response = get_output(generator, user_input)[0][len(user_input):]
+def chat(user_input, top_p, temperature, max_gen_len):
+    bot_response = get_output(
+        generator=generator, 
+        prompt=user_input,
+        max_gen_len=max_gen_len,
+        temperature=temperature,
+        top_p=top_p)
+
+    # remove the first phrase identical to user prompt
+    bot_response = bot_response[0][len(user_input):]
+    # trip the last phrase
+    try:
+        bot_response = bot_response[:bot_response.rfind(".")]
+    except:
+        pass
 
     history.append({
         "role": "user",
@@ -39,6 +52,9 @@ def chat(user_input):
         simple_history[-1] = current_pair
         yield simple_history
 
+def reset_textbox():
+    return gr.update(value='')
+
 with gr.Blocks(css = """#col_container {width: 95%; margin-left: auto; margin-right: auto;}
                 #chatbot {height: 400px; overflow: auto;}""") as demo:
     
@@ -46,7 +62,13 @@ with gr.Blocks(css = """#col_container {width: 95%; margin-left: auto; margin-ri
         gr.Markdown(f"## {TITLE}\n\n\n\n{ABSTRACT}")
         chatbot = gr.Chatbot(elem_id='chatbot')
         textbox = gr.Textbox(placeholder="Enter a prompt")
-    
-        textbox.submit(chat, textbox, chatbot)
+
+        with gr.Accordion("Parameters", open=False):
+            max_gen_len = gr.Slider(minimum=20, maximum=512, value=256, step=1, interactive=True, label="Max Genenration Length",)
+            top_p = gr.Slider(minimum=-0, maximum=1.0, value=1.0, step=0.05, interactive=True, label="Top-p (nucleus sampling)",)
+            temperature = gr.Slider(minimum=-0, maximum=5.0, value=1.0, step=0.1, interactive=True, label="Temperature",)
+        
+    textbox.submit(chat, [textbox, top_p, temperature, max_gen_len], chatbot)
+    textbox.submit(reset_textbox, [], [textbox])
 
 demo.queue(api_open=False).launch()