more settings

demo_watermark.py

@@ -261,7 +261,7 @@ def detect(input_text, args, device=None, tokenizer=None):
 
 def run_gradio(args, model=None, device=None, tokenizer=None):
 
-    generate_partial = partial(generate, model=model, device=None, tokenizer=tokenizer)
+    generate_partial = partial(generate, model=model, device=device, tokenizer=tokenizer)
     detect_partial = partial(detect, device=device, tokenizer=tokenizer)
 
     with gr.Blocks() as demo:
@@ -289,11 +289,13 @@ def run_gradio(args, model=None, device=None, tokenizer=None):
                         generation_seed = gr.Number(label="Generation Seed",value=args.generation_seed, interactive=True)
                     with gr.Row():
                         n_beams = gr.Dropdown(label="Number of Beams",choices=list(range(1,11,1)), value=args.n_beams, visible=(not args.use_sampling))
+                    with gr.Row():
+                        max_new_tokens = gr.Slider(label="Max Generated Tokens", minimum=10, maximum=1000, step=10, value=args.max_new_tokens)
 
                 with gr.Column(scale=1):
                     gr.Markdown(f"#### Watermarking Parameters")
                     with gr.Row():
-                        gamma = gr.Slider(label="gamma",minimum=0.1, maximum=0.9, step=0.1, value=args.gamma)
+                        gamma = gr.Slider(label="gamma",minimum=0.1, maximum=0.9, step=0.05, value=args.gamma)
                     with gr.Row():
                         delta = gr.Slider(label="delta",minimum=0.0, maximum=10.0, step=0.1, value=args.delta)
                     with gr.Row():
@@ -326,6 +328,7 @@ def run_gradio(args, model=None, device=None, tokenizer=None):
             elif value == "greedy":
                 return gr.update(visible=True)
         def update_n_beams(session_state, value): session_state.n_beams = int(value); return session_state
+        def update_max_new_tokens(session_state, value): session_state.max_new_tokens = int(value); return session_state
         def update_ignore_repeated_bigrams(session_state, value): session_state.ignore_repeated_bigrams = value; return session_state
         def update_normalizers(session_state, value): session_state.normalizers = value; return session_state
         
@@ -337,6 +340,7 @@ def run_gradio(args, model=None, device=None, tokenizer=None):
         sampling_temp.change(update_sampling_temp,inputs=[session_args, sampling_temp], outputs=[session_args])
         generation_seed.change(update_generation_seed,inputs=[session_args, generation_seed], outputs=[session_args])
         n_beams.change(update_n_beams,inputs=[session_args, n_beams], outputs=[session_args])
+        max_new_tokens.change(update_max_new_tokens,inputs=[session_args, max_new_tokens], outputs=[session_args])
 
         gamma.change(update_gamma,inputs=[session_args, gamma], outputs=[session_args])
         delta.change(update_delta,inputs=[session_args, delta], outputs=[session_args])
@@ -365,7 +369,7 @@ def run_gradio(args, model=None, device=None, tokenizer=None):
             truncation_warning = gr.Number(visible=False)
             def truncate_prompt(redecoded_input, truncation_warning, orig_prompt, args):
                 if truncation_warning:
-                    return redecoded_input + f"\n\n[Prompt was truncated before generation due to length...]"
+                    return redecoded_input + f"\n\n[Prompt was truncated before generation due to length...]", args
                 else: 
                     return orig_prompt, args
             
@@ -412,7 +416,7 @@ def main(args):
     if not args.skip_model_load:
         model, tokenizer, device = load_model(args)
     else:
-        model, tokenizer, device = None, None, []
+        model, tokenizer, device = None, None, None
 
     # Generate and detect, report to stdout
     if not args.skip_model_load:
@@ -442,7 +446,7 @@ def main(args):
         input_text = "In this work, we study watermarking of language model output. A watermark is a hidden pattern in text that is imperceptible to humans, while making the text algorithmically identifiable as synthetic. We propose an efficient watermark that makes synthetic text detectable from short spans of tokens (as few as 25 words), while false-positives (where human text is marked as machine-generated) are statistically improbable. The watermark detection algorithm can be made public, enabling third parties (e.g., social media platforms) to run it themselves, or it can be kept private and run behind an API.  We seek a watermark with the following properties:\n"
 
 
-        term_width = os.get_terminal_size()[0]
+        term_width = 80
         print("#"*term_width)
         print("Prompt:")
         print(input_text)