working demo

app.py

@@ -3,6 +3,7 @@ from threading import Thread
 import torch
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
+import time
 
 model_id = "EleutherAI/pythia-6.9b-deduped"
 assistant_id = "EleutherAI/pythia-70m-deduped"
@@ -19,7 +20,7 @@ tokenizer = AutoTokenizer.from_pretrained(model_id)
 assistant_model = AutoModelForCausalLM.from_pretrained(assistant_id).to(torch_device)
 
 
-def run_generation(user_text, use_assistant, top_p, temperature, top_k, max_new_tokens):
+def run_generation(user_text, use_assistant, temperature, max_new_tokens):
     if temperature == 0.0:
         do_sample = False
     else:
@@ -37,19 +38,21 @@ def run_generation(user_text, use_assistant, top_p, temperature, top_k, max_new_
         streamer=streamer,
         max_new_tokens=max_new_tokens,
         do_sample=do_sample,
-        top_p=top_p,
+        top_p=0.95,
         temperature=float(temperature),
-        top_k=top_k
+        top_k=50,
     )
     t = Thread(target=model.generate, kwargs=generate_kwargs)
+    start = time.time()
     t.start()
 
-    # Pull the generated text from the streamer, and update the model output.
+    # Pull the generated text from the streamer, and update the model output. Return the model output and time
+    # spent so far.
     model_output = ""
     for new_text in streamer:
         model_output += new_text
-        yield model_output
-    return model_output
+        yield [model_output, round(time.time() - start, 3)]
+    return [model_output, round(time.time() - start, 3)]
 
 
 def reset_textbox():
@@ -59,36 +62,34 @@ def reset_textbox():
 with gr.Blocks() as demo:
     gr.Markdown(
         "# 🤗 Assisted Generation Demo\n"
-        f"Model: {model_id} (using INT8)\n\n"
-        f"Assistant Model: {assistant_id}"
+        f"- Model: {model_id} (using INT8)\n"
+        f"- Assistant Model: {assistant_id}"
     )
 
     with gr.Row():
         with gr.Column(scale=4):
             user_text = gr.Textbox(
                 placeholder="Question: What is the meaning of life? Answer:",
-                label="User input"
+                label="Prompt"
             )
             model_output = gr.Textbox(label="Model output", lines=10, interactive=False)
             button_submit = gr.Button(value="Submit")
 
         with gr.Column(scale=1):
-            use_assistant = gr.Checkbox(label="Use Assistant", default=True)
+            gr.Markdown("### Generation Settings")
+            use_assistant = gr.Checkbox(label="Use Assisted Generation", value=True)
             max_new_tokens = gr.Slider(
-                minimum=1, maximum=500, value=250, step=1, interactive=True, label="Max New Tokens",
-            )
-            top_p = gr.Slider(
-                minimum=0.05, maximum=1.0, value=0.95, step=0.05, interactive=True, label="Top-p",
-            )
-            top_k = gr.Slider(
-                minimum=1, maximum=50, value=50, step=1, interactive=True, label="Top-k",
+                minimum=1, maximum=500, value=100, step=1, interactive=True, label="Max New Tokens",
             )
             temperature = gr.Slider(
                 minimum=0.0, maximum=2.0, value=0.0, step=0.1, interactive=True, label="Temperature (0.0 = Greedy)",
             )
+            gr.Markdown("### Generation time (seconds)")
+            generation_time = gr.Textbox(lines=1, interactive=False, show_label=False)
 
-    generate_inputs = [user_text, use_assistant, top_p, temperature, top_k, max_new_tokens]
-    user_text.submit(run_generation, generate_inputs, model_output)
-    button_submit.click(run_generation, generate_inputs, model_output)
+    generate_inputs = [user_text, use_assistant, temperature, max_new_tokens]
+    generate_outputs = [model_output, generation_time]
+    user_text.submit(run_generation, generate_inputs, generate_outputs)
+    button_submit.click(run_generation, generate_inputs, generate_outputs)
 
     demo.queue(max_size=32).launch(enable_queue=True)