Update app.py

app.py

@@ -3,19 +3,22 @@ import time
 import torch
 import gradio as gr
 
-from strings import TITLE, ABSTRACT 
+from strings import TITLE, ABSTRACT, EXAMPLES
 from gen import get_pretrained_models, get_output
 
-os.environ["RANK"] = "0"
-os.environ["WORLD_SIZE"] = "1"
-os.environ["MASTER_ADDR"] = "127.0.0.1"
-os.environ["MASTER_PORT"] = "50505"
-
 generator = get_pretrained_models("13B", "tokenizer")
 
 history = []
 
-def chat(user_input, top_p, temperature, max_gen_len, state_chatbot):
+def chat(
+    user_input, 
+    include_input,
+    truncate,    
+    top_p, 
+    temperature, 
+    max_gen_len, 
+    state_chatbot
+):
     bot_response = get_output(
         generator=generator, 
         prompt=user_input,
@@ -24,13 +27,16 @@ def chat(user_input, top_p, temperature, max_gen_len, state_chatbot):
         top_p=top_p)
 
     # remove the first phrase identical to user prompt
-    bot_response = bot_response[0][len(user_input):]
-    bot_response = bot_response.replace("\n", "<br><br>")
+    if not include_input:
+        bot_response = bot_response[0][len(user_input):]
+    bot_response = bot_response.replace("\n", "<br>")
+    
     # trip the last phrase
-    try:
-        bot_response = bot_response[:bot_response.rfind(".")]
-    except:
-        pass
+    if truncate:
+        try:
+            bot_response = bot_response[:bot_response.rfind(".")+1]
+        except:
+            pass
 
     history.append({
         "role": "user",
@@ -61,6 +67,14 @@ with gr.Blocks(css = """#col_container {width: 95%; margin-left: auto; margin-ri
                     
     with gr.Column(elem_id='col_container'):
         gr.Markdown(f"## {TITLE}\n\n\n\n{ABSTRACT}")
+
+        with gr.Accordion("Example prompts", open=False):
+            example_str = "\n"
+            for example in EXAMPLES:
+                example_str += f"- {example}\n"
+            
+            gr.Markdown(example_str)        
+        
         chatbot = gr.Chatbot(elem_id='chatbot')
         textbox = gr.Textbox(placeholder="Enter a prompt")
 
@@ -69,7 +83,11 @@ with gr.Blocks(css = """#col_container {width: 95%; margin-left: auto; margin-ri
             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, state_chatbot], [state_chatbot, chatbot])
+    textbox.submit(
+        chat, 
+        [textbox, include_input, truncate, top_p, temperature, max_gen_len, state_chatbot],
+        [state_chatbot, chatbot]
+    )    
     textbox.submit(reset_textbox, [], [textbox])
 
 demo.queue(api_open=False).launch()