features mostly in place

app.py

@@ -21,6 +21,7 @@ arg_dict = {
     'run_gradio': True, 
     'demo_public': False, 
     'model_name_or_path': 'facebook/opt-125m', 
+    # 'model_name_or_path': 'facebook/opt-2.7b', 
     'prompt_max_length': None, 
     'max_new_tokens': 200, 
     'generation_seed': 123, 
@@ -36,6 +37,7 @@ arg_dict = {
     'detection_z_threshold': 4.0, 
     'select_green_tokens': True,
     'skip_model_load': False,
+    'seed_separately': True,
 }
 
 args.__dict__.update(arg_dict)

demo_watermark.py

@@ -223,7 +223,10 @@ def generate(prompt, args, model=None, device=None, tokenizer=None):
 
     torch.manual_seed(args.generation_seed)
     output_without_watermark = generate_without_watermark(**tokd_input)
-    # torch.manual_seed(seed) # optional, but will not be the same again generally, unless delta==0.0, no-op watermark
+
+    # optional to seed before second generation, but will not be the same again generally, unless delta==0.0, no-op watermark
+    if args.seed_separately: 
+        torch.manual_seed(args.generation_seed)
     output_with_watermark = generate_with_watermark(**tokd_input)
 
     if args.is_decoder_only_model:
@@ -275,7 +278,52 @@ def run_gradio(args, model=None, device=None, tokenizer=None):
                 <img style="margin-top: 0em; margin-bottom: 0em" src="https://bit.ly/3gLdBN6" alt="Duplicate Space"></a>
                 <p/>
                 """)
+        # Construct state for parameters, define updates and toggles, and register event listeners
+        session_args = gr.State(value=args)
+
+        with gr.Tab("Generation"):
+            
+            with gr.Row():
+                prompt = gr.Textbox(label=f"Prompt", interactive=True)
+            with gr.Row():
+                generate_btn = gr.Button("Generate")
+            with gr.Row():
+                with gr.Column(scale=2):
+                    output_without_watermark = gr.Textbox(label="Output Without Watermark", interactive=False)
+                with gr.Column(scale=1):
+                    without_watermark_detection_result = gr.Textbox(label="Detection Result", interactive=False)
+            with gr.Row():
+                with gr.Column(scale=2):
+                    output_with_watermark = gr.Textbox(label="Output With Watermark", interactive=False)
+                with gr.Column(scale=1):
+                    with_watermark_detection_result = gr.Textbox(label="Detection Result", interactive=False)
+
+            redecoded_input = gr.Textbox(visible=False)
+            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...]", args
+                else: 
+                    return orig_prompt, args
+            
+            generate_btn.click(fn=generate_partial, inputs=[prompt,session_args], outputs=[redecoded_input, truncation_warning, output_without_watermark, output_with_watermark,session_args])
+            
+            # Show truncated version of prompt if truncation occurred
+            redecoded_input.change(fn=truncate_prompt, inputs=[redecoded_input,truncation_warning,prompt,session_args], outputs=[prompt,session_args])
+
+            # Call detection when the outputs of the generate function are updated.
+            output_without_watermark.change(fn=detect_partial, inputs=[output_without_watermark,session_args], outputs=[without_watermark_detection_result,session_args])
+            output_with_watermark.change(fn=detect_partial, inputs=[output_with_watermark,session_args], outputs=[with_watermark_detection_result,session_args])
         
+        with gr.Tab("Detector Only"):
+            with gr.Row():
+                detection_input = gr.Textbox(label="Text to Analyze", interactive=True)
+            with gr.Row():
+                detect_btn = gr.Button("Detect")
+            with gr.Row():
+                detection_result = gr.Textbox(label="Detection Result", interactive=False)
+            detect_btn.click(fn=detect_partial, inputs=[detection_input,session_args], outputs=[detection_result, session_args])
+
         # Parameter selection group
         with gr.Accordion("Advanced Settings",open=False):
             with gr.Row():
@@ -302,11 +350,29 @@ def run_gradio(args, model=None, device=None, tokenizer=None):
                         ignore_repeated_bigrams = gr.Checkbox(label="Ignore Bigram Repeats")
                     with gr.Row():
                         normalizers = gr.CheckboxGroup(label="Normalizations", choices=["unicode", "homoglyphs", "truecase"], value=args.normalizers)
+            gr.Markdown(f"_Note: sliders don't always update perfectly. Clicking on the bar or using the number window to the right can help._")
+            with gr.Accordion("Actual submitted parameters:",open=False):
+                current_parameters = gr.Textbox(label="submitted parameters", value=args)
+            with gr.Accordion("Legacy Settings",open=False):
+                with gr.Row():
+                    with gr.Column(scale=1):
+                        seed_separately = gr.Checkbox(label="Seed both generations separately", value=args.seed_separately)
+                    with gr.Column(scale=1):
+                        select_green_tokens = gr.Checkbox(label="Select 'greenlist' from partition", value=args.select_green_tokens)
+
         
-        # State manager
-        # Construct state for parameters, define updates and toggles, and register event listeners
-        session_args = gr.State(value=args)
+        with gr.Accordion("A note on model capability",open=False):
+            gr.Markdown(
+                """
+                The models that can be used in this demo are limited to those that are open source as well as fit on a single commodity GPU. In particular, there are few models above 10B parameters and way fewer trained using both Instruction finetuning or RLHF that are open source that we can use.
+                
+                Therefore, the model, in both it's un-watermarked (normal) and watermarked state, is not generally able to respond well to the kinds of prompts that a 100B+ Instruction and RLHF tuned model such as ChatGPT, Claude, or Bard is.
+                
+                We suggest you try prompts that give the model a few sentences and then allow it to 'continue' the prompt, as these weaker models are more capable in this simpler language modeling setting.
+                """
+                )
 
+        # State manager logic
         def update_sampling_temp(session_state, value): session_state.sampling_temp = float(value); return session_state
         def update_generation_seed(session_state, value): session_state.generation_seed = int(value); return session_state
         def update_gamma(session_state, value): session_state.gamma = float(value); return session_state
@@ -331,76 +397,33 @@ def run_gradio(args, model=None, device=None, tokenizer=None):
         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
+        def update_seed_separately(session_state, value): session_state.seed_separately = value; return session_state
+        def update_select_green_tokens(session_state, value): session_state.select_green_tokens = value; return session_state
         
-        decoding.change(update_decoding,inputs=[session_args, decoding], outputs=[session_args])
         decoding.change(toggle_sampling_vis,inputs=[decoding], outputs=[sampling_temp])
         decoding.change(toggle_sampling_vis,inputs=[decoding], outputs=[generation_seed])
         decoding.change(toggle_sampling_vis_inv,inputs=[decoding], outputs=[n_beams])
-
+        
+        decoding.change(update_decoding,inputs=[session_args, decoding], outputs=[session_args])
         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])
         ignore_repeated_bigrams.change(update_ignore_repeated_bigrams,inputs=[session_args, ignore_repeated_bigrams], outputs=[session_args])
         normalizers.change(update_normalizers,inputs=[session_args, normalizers], outputs=[session_args])
+        seed_separately.change(update_seed_separately,inputs=[session_args, seed_separately], outputs=[session_args])
+        select_green_tokens.change(update_select_green_tokens,inputs=[session_args, select_green_tokens], outputs=[session_args])
 
-        with gr.Tab("Generation"):
-            
-            with gr.Row():
-                prompt = gr.Textbox(label=f"Prompt", interactive=True)
-            with gr.Row():
-                generate_btn = gr.Button("Generate")
-            with gr.Row():
-                with gr.Column(scale=2):
-                    output_without_watermark = gr.Textbox(label="Output Without Watermark", interactive=False)
-                with gr.Column(scale=1):
-                    without_watermark_detection_result = gr.Textbox(label="Detection Result", interactive=False)
-            with gr.Row():
-                with gr.Column(scale=2):
-                    output_with_watermark = gr.Textbox(label="Output With Watermark", interactive=False)
-                with gr.Column(scale=1):
-                    with_watermark_detection_result = gr.Textbox(label="Detection Result", interactive=False)
-            
-            
-            redecoded_input = gr.Textbox(visible=False)
-            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...]", args
-                else: 
-                    return orig_prompt, args
-            
-            generate_btn.click(fn=generate_partial, inputs=[prompt,session_args], outputs=[redecoded_input, truncation_warning, output_without_watermark, output_with_watermark,session_args])
-            
-            # Show truncated version of prompt if truncation occurred
-            redecoded_input.change(fn=truncate_prompt, inputs=[redecoded_input,truncation_warning,prompt,session_args], outputs=[prompt,session_args])
+        generate_btn.click(lambda value: str(value), inputs=[session_args], outputs=[current_parameters])
+        detect_btn.click(lambda value: str(value), inputs=[session_args], outputs=[current_parameters])
 
-            # Call detection when the outputs of the generate function are updated.
-            output_without_watermark.change(fn=detect_partial, inputs=[output_without_watermark,session_args], outputs=[without_watermark_detection_result,session_args])
-            output_with_watermark.change(fn=detect_partial, inputs=[output_with_watermark,session_args], outputs=[with_watermark_detection_result,session_args])
+        # When the parameters change, also fire detection, since some detection params dont change the model output.
+        current_parameters.change(fn=detect_partial, inputs=[output_without_watermark,session_args], outputs=[without_watermark_detection_result,session_args])
+        current_parameters.change(fn=detect_partial, inputs=[output_with_watermark,session_args], outputs=[with_watermark_detection_result,session_args])
 
-        with gr.Tab("Detector Only"):
-            with gr.Row():
-                detection_input = gr.Textbox(label="Text to Analyze", interactive=True)
-            with gr.Row():
-                detect_btn = gr.Button("Detect")
-            with gr.Row():
-                detection_result = gr.Textbox(label="Detection Result", interactive=False)
-            detect_btn.click(fn=detect_partial, inputs=[detection_input,session_args], outputs=[detection_result, session_args])
-        
-        with gr.Accordion("A note on model capability",open=False):
-            gr.Markdown(
-                """
-                The models that can be used in this demo are limited to those that are open source as well as fit on a single commodity GPU. In particular, there are few models above 10B parameters and way fewer trained using both Instruction finetuning or RLHF that are open source that we can use.
-                
-                Therefore, the model, in both it's un-watermarked (normal) and watermarked state, is not generally able to respond well to the kinds of prompts that a 100B+ Instruction and RLHF tuned model such as ChatGPT, Claude, or Bard is.
-                
-                We suggest you try prompts that give the model a few sentences and then allow it to 'continue' the prompt, as these weaker models are more capable in this simpler language modeling setting.
-                """
-                )
+    demo.queue(concurrency_count=3)
 
     if args.demo_public:
         demo.launch(share=True) # exposes app to the internet via randomly generated link