assistant toggle

app.py

@@ -19,7 +19,12 @@ tokenizer = AutoTokenizer.from_pretrained(model_id)
 assistant_model = AutoModelForCausalLM.from_pretrained(assistant_id).to(torch_device)
 
 
-def run_generation(user_text, top_p, temperature, top_k, max_new_tokens):
+def run_generation(user_text, use_assistant, top_p, temperature, top_k, max_new_tokens):
+    if temperature == 0.0:
+        do_sample = False
+    else:
+        do_sample = True
+
     # Get the model and tokenizer, and tokenize the user text.
     model_inputs = tokenizer([user_text], return_tensors="pt").to(torch_device)
 
@@ -28,9 +33,10 @@ def run_generation(user_text, top_p, temperature, top_k, max_new_tokens):
     streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True)
     generate_kwargs = dict(
         model_inputs,
+        assistant_model=assistant_model if use_assistant else None,
         streamer=streamer,
         max_new_tokens=max_new_tokens,
-        do_sample=True,
+        do_sample=do_sample,
         top_p=top_p,
         temperature=float(temperature),
         top_k=top_k
@@ -53,34 +59,36 @@ def reset_textbox():
 with gr.Blocks() as demo:
     gr.Markdown(
         "# 🤗 Assisted Generation Demo\n"
-        f"Model: {model_id} (using INT8)\n"
+        f"Model: {model_id} (using INT8)\n\n"
         f"Assistant Model: {assistant_id}"
     )
 
     with gr.Row():
         with gr.Column(scale=4):
             user_text = gr.Textbox(
-                placeholder="Write an email about an alpaca that likes flan",
+                placeholder="Question: What is the meaning of life? Answer:",
                 label="User input"
             )
             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)
             max_new_tokens = gr.Slider(
-                minimum=1, maximum=1000, value=250, step=1, interactive=True, label="Max New Tokens",
+                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 (nucleus sampling)",
+                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",
             )
             temperature = gr.Slider(
-                minimum=0.1, maximum=5.0, value=0.8, step=0.1, interactive=True, label="Temperature",
+                minimum=0.0, maximum=2.0, value=0.0, step=0.1, interactive=True, label="Temperature (0.0 = Greedy)",
             )
 
-    user_text.submit(run_generation, [user_text, top_p, temperature, top_k, max_new_tokens], model_output)
-    button_submit.click(run_generation, [user_text, top_p, temperature, top_k, max_new_tokens], model_output)
+    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)
 
     demo.queue(max_size=32).launch(enable_queue=True)