a version with toggles

app.py

@@ -25,6 +25,7 @@ arg_dict = {
     'max_new_tokens': 200, 
     'generation_seed': 123, 
     'use_sampling': True, 
+    'n_beams': 1, 
     'sampling_temp': 0.7, 
     'use_gpu': True, 
     'seeding_scheme': 'markov_1', 
@@ -33,11 +34,11 @@ arg_dict = {
     'normalizers': '', 
     'ignore_repeated_bigrams': False, 
     'detection_z_threshold': 4.0, 
-    'select_green_tokens': True
+    'select_green_tokens': True,
+    'skip_model_load': False,
 }
 
 args.__dict__.update(arg_dict)
-print(args)
 
 from demo_watermark import main
 

demo_watermark.py

@@ -16,9 +16,13 @@
 
 import os
 import argparse
+from argparse import Namespace
 from pprint import pprint
 from functools import partial
 
+import numpy # for gradio hot reload
+import gradio as gr
+
 import torch
 
 from transformers import (AutoTokenizer,
@@ -45,8 +49,8 @@ def parse_args():
     parser.add_argument(
         "--run_gradio",
         type=str2bool,
-        default=False,
-        help="Whether to launch as a gradio demo.",
+        default=True,
+        help="Whether to launch as a gradio demo. Set to False if not installed and want to just run the stdout version.",
     )
     parser.add_argument(
         "--demo_public",
@@ -90,6 +94,12 @@ def parse_args():
         default=0.7,
         help="Sampling temperature to use when generating using multinomial sampling.",
     )
+    parser.add_argument(
+        "--n_beams",
+        type=int,
+        choices=[1,4,8],
+        help="Number of beams to use for beam search. 1 is normal greedy decoding",
+    )
     parser.add_argument(
         "--use_gpu",
         type=str2bool,
@@ -138,17 +148,21 @@ def parse_args():
         default=True,
         help="How to treat the permuation when selecting the greenlist tokens at each step. Legacy is (False) to pick the complement/reds first.",
     )
+    parser.add_argument(
+        "--skip_model_load",
+        type=str2bool,
+        default=False,
+        help="Skip the model loading to debug the interface.",
+    )
     args = parser.parse_args()
     return args
 
-
-def main(args): 
-
-    is_seq2seq_model = any([(model_type in args.model_name_or_path) for model_type in ["t5","T0"]])
-    is_decoder_only_model = any([(model_type in args.model_name_or_path) for model_type in ["gpt","opt","bloom"]])
-    if is_seq2seq_model:
+def load_model():
+    args.is_seq2seq_model = any([(model_type in args.model_name_or_path) for model_type in ["t5","T0"]])
+    args.is_decoder_only_model = any([(model_type in args.model_name_or_path) for model_type in ["gpt","opt","bloom"]])
+    if args.is_seq2seq_model:
         model = AutoModelForSeq2SeqLM.from_pretrained(args.model_name_or_path)
-    elif is_decoder_only_model:
+    elif args.is_decoder_only_model:
         model = AutoModelForCausalLM.from_pretrained(args.model_name_or_path)
     else:
         raise ValueError(f"Unknown model type: {args.model_name_or_path}")
@@ -161,213 +175,302 @@ def main(args):
     model.eval()
 
     tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path)
-    vocabulary = list(tokenizer.get_vocab().values())
-
-    def generate(prompt):
-
-        watermark_processor = WatermarkLogitsProcessor(vocab=vocabulary,
-                                                        gamma=args.gamma,
-                                                        delta=args.delta,
-                                                        seeding_scheme=args.seeding_scheme,
-                                                        select_green_tokens=args.select_green_tokens)
-
-        gen_kwargs = dict(max_new_tokens=args.max_new_tokens)
-
-        if args.use_sampling:
-            gen_kwargs.update(dict(
-                do_sample=True, 
-                top_k=0,
-                temperature=args.sampling_temp
-            ))
-        else:
-            gen_kwargs.update(dict(
-                num_beams=args.n_beams
-            ))
-
-        generate_without_watermark = partial(
-            model.generate,
-            **gen_kwargs
-        )
-        generate_with_watermark = partial(
-            model.generate,
-            logits_processor=LogitsProcessorList([watermark_processor]), 
-            **gen_kwargs
-        )
-        if args.prompt_max_length:
-            pass
-        elif hasattr(model.config,"max_position_embedding"):
-            args.prompt_max_length = model.config.max_position_embeddings-args.max_new_tokens
-        else:
-            args.prompt_max_length = 2048-args.max_new_tokens
-
-        tokd_input = tokenizer(prompt, return_tensors="pt", add_special_tokens=True, truncation=True, max_length=args.prompt_max_length).to(device)
-        truncation_warning = True if tokd_input["input_ids"].shape[-1] == args.prompt_max_length else False
-        redecoded_input = tokenizer.batch_decode(tokd_input["input_ids"], skip_special_tokens=True)[0]
-
-        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
-        output_with_watermark = generate_with_watermark(**tokd_input)
-
-        if is_decoder_only_model:
-            # need to isolate the newly generated tokens
-            output_without_watermark = output_without_watermark[:,tokd_input["input_ids"].shape[-1]:]
-            output_with_watermark = output_with_watermark[:,tokd_input["input_ids"].shape[-1]:]
-
-        decoded_output_without_watermark = tokenizer.batch_decode(output_without_watermark, skip_special_tokens=True)[0]
-        decoded_output_with_watermark = tokenizer.batch_decode(output_with_watermark, skip_special_tokens=True)[0]
-
-        return (redecoded_input,
-                int(truncation_warning),
-                decoded_output_without_watermark, 
-                decoded_output_with_watermark) 
-                # decoded_output_with_watermark)
-
-    def detect(input_text):
-        watermark_detector = WatermarkDetector(vocab=list(tokenizer.get_vocab().values()),
-                                            gamma=args.gamma,
-                                            seeding_scheme=args.seeding_scheme,
-                                            device=device,
-                                            tokenizer=tokenizer,
-                                            z_threshold=args.detection_z_threshold,
-                                            normalizers=(args.normalizers.split(",") if args.normalizers else []),
-                                            ignore_repeated_bigrams=args.ignore_repeated_bigrams,
-                                            select_green_tokens=args.select_green_tokens)
-        if len(input_text)-1 > watermark_detector.min_prefix_len:
-            score_dict = watermark_detector.detect(input_text)
-            output_str = (f"Detection result @ {watermark_detector.z_threshold}:\n"
-                            f"{score_dict}")
-        else:
-            output_str = (f"Error: string not long enough to compute watermark presence.")
-        return output_str
 
-    # Generate and detect, report to stdout
+    return model, tokenizer, device
 
-    # input_text = (
-    # "The diamondback terrapin or simply terrapin (Malaclemys terrapin) is a "
-    # "species of turtle native to the brackish coastal tidal marshes of the "
-    # "Northeastern and southern United States, and in Bermuda.[6] It belongs "
-    # "to the monotypic genus Malaclemys. It has one of the largest ranges of "
-    # "all turtles in North America, stretching as far south as the Florida Keys "
-    # "and as far north as Cape Cod.[7] The name 'terrapin' is derived from the "
-    # "Algonquian word torope.[8] It applies to Malaclemys terrapin in both "
-    # "British English and American English. The name originally was used by "
-    # "early European settlers in North America to describe these brackish-water "
-    # "turtles that inhabited neither freshwater habitats nor the sea. It retains "
-    # "this primary meaning in American English.[8] In British English, however, "
-    # "other semi-aquatic turtle species, such as the red-eared slider, might "
-    # "also be called terrapins. The common name refers to the diamond pattern "
-    # "on top of its shell (carapace), but the overall pattern and coloration "
-    # "vary greatly. The shell is usually wider at the back than in the front, "
-    # "and from above it appears wedge-shaped. The shell coloring can vary "
-    # "from brown to grey, and its body color can be grey, brown, yellow, "
-    # "or white. All have a unique pattern of wiggly, black markings or spots "
-    # "on their body and head. The diamondback terrapin has large webbed "
-    # "feet.[9] The species is"
-    # )
-
-    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 = 80
-    print("#"*term_width)
-    print("Prompt:")
-    print(input_text)
-
-    _, _, decoded_output_without_watermark, decoded_output_with_watermark = generate(input_text)
-    without_watermark_detection_result = detect(decoded_output_without_watermark)
-    with_watermark_detection_result = detect(decoded_output_with_watermark)
-
-    print("#"*term_width)
-    print("Output without watermark:")
-    print(decoded_output_without_watermark)
-    print("-"*term_width)
-    print(f"Detection result @ {args.detection_z_threshold}:")
-    pprint(without_watermark_detection_result)
-    print("-"*term_width)
-
-    print("#"*term_width)
-    print("Output with watermark:")
-    print(decoded_output_with_watermark)
-    print("-"*term_width)
-    print(f"Detection result @ {args.detection_z_threshold}:")
-    pprint(with_watermark_detection_result)
-    print("-"*term_width)
+def generate(prompt, args, model=None, tokenizer=None):
 
-    # Launch the app to generate and detect interactively (implements the hf space demo)
+    print(f"Generating with {args}")
 
-    if args.run_gradio:
-        import gradio as gr
-    
-        with gr.Blocks() as demo:
-            gr.Markdown("## Demo for ['A Watermark for Large Language Models'](https://arxiv.org/abs/2301.10226)")
-            # gr.HTML("""
-            #         <p>For faster inference without waiting in queue, you may duplicate the space and upgrade to GPU in settings.
-            #         <br/>
-            #         <a href="https://huggingface.co/spaces/tomg-group-umd/pez-dispenser?duplicate=true">
-            #         <img style="margin-top: 0em; margin-bottom: 0em" src="https://bit.ly/3gLdBN6" alt="Duplicate Space"></a>
-            #         <p/>
-            #         """)
-            gr.Markdown(f"#### Generation and Watermarking Parameters:\n\n{args.__dict__}")
-
-            with gr.Tab("Generation"):
-                with gr.Row():
-                    prompt = gr.Textbox(label=f"Prompt (max {args.prompt_max_length} tokens)", 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):
-                    if truncation_warning:
-                        return redecoded_input + f"\n\n[Prompt was truncated before generation due to length...]"
-                    else: 
-                        return orig_prompt
+    watermark_processor = WatermarkLogitsProcessor(vocab=list(tokenizer.get_vocab().values()),
+                                                    gamma=args.gamma,
+                                                    delta=args.delta,
+                                                    seeding_scheme=args.seeding_scheme,
+                                                    select_green_tokens=args.select_green_tokens)
+
+    gen_kwargs = dict(max_new_tokens=args.max_new_tokens)
+
+    if args.use_sampling:
+        gen_kwargs.update(dict(
+            do_sample=True, 
+            top_k=0,
+            temperature=args.sampling_temp
+        ))
+    else:
+        gen_kwargs.update(dict(
+            num_beams=args.n_beams
+        ))
+
+    generate_without_watermark = partial(
+        model.generate,
+        **gen_kwargs
+    )
+    generate_with_watermark = partial(
+        model.generate,
+        logits_processor=LogitsProcessorList([watermark_processor]), 
+        **gen_kwargs
+    )
+    if args.prompt_max_length:
+        pass
+    elif hasattr(model.config,"max_position_embedding"):
+        args.prompt_max_length = model.config.max_position_embeddings-args.max_new_tokens
+    else:
+        args.prompt_max_length = 2048-args.max_new_tokens
+
+    tokd_input = tokenizer(prompt, return_tensors="pt", add_special_tokens=True, truncation=True, max_length=args.prompt_max_length).to(device)
+    truncation_warning = True if tokd_input["input_ids"].shape[-1] == args.prompt_max_length else False
+    redecoded_input = tokenizer.batch_decode(tokd_input["input_ids"], skip_special_tokens=True)[0]
+
+    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
+    output_with_watermark = generate_with_watermark(**tokd_input)
+
+    if args.is_decoder_only_model:
+        # need to isolate the newly generated tokens
+        output_without_watermark = output_without_watermark[:,tokd_input["input_ids"].shape[-1]:]
+        output_with_watermark = output_with_watermark[:,tokd_input["input_ids"].shape[-1]:]
+
+    decoded_output_without_watermark = tokenizer.batch_decode(output_without_watermark, skip_special_tokens=True)[0]
+    decoded_output_with_watermark = tokenizer.batch_decode(output_with_watermark, skip_special_tokens=True)[0]
+
+    return (redecoded_input,
+            int(truncation_warning),
+            decoded_output_without_watermark, 
+            decoded_output_with_watermark,
+            args) 
+            # decoded_output_with_watermark)
+
+def detect(input_text, args, device=None, tokenizer=None):
+    watermark_detector = WatermarkDetector(vocab=list(tokenizer.get_vocab().values()),
+                                        gamma=args.gamma,
+                                        seeding_scheme=args.seeding_scheme,
+                                        device=device,
+                                        tokenizer=tokenizer,
+                                        z_threshold=args.detection_z_threshold,
+                                        normalizers=args.normalizers,
+                                        ignore_repeated_bigrams=args.ignore_repeated_bigrams,
+                                        select_green_tokens=args.select_green_tokens)
+    if len(input_text)-1 > watermark_detector.min_prefix_len:
+        score_dict = watermark_detector.detect(input_text)
+        output_str = (f"Detection result @ {watermark_detector.z_threshold}:\n"
+                        f"{score_dict}")
+    else:
+        output_str = (f"Error: string not long enough to compute watermark presence.")
+    return output_str, args
+
+def run_gradio(args, model=None, device=None, tokenizer=None):
+
+    generate_partial = partial(generate, model=model, tokenizer=tokenizer)
+    detect_partial = partial(detect, device=device, tokenizer=tokenizer)
+
+    with gr.Blocks() as demo:
+
+        # Top section, greeting and instructions
+        gr.Markdown("## Demo for ['A Watermark for Large Language Models'](https://arxiv.org/abs/2301.10226)")
+        gr.HTML("""
+                <p>For faster inference without waiting in queue, you may duplicate the space and upgrade to GPU in settings.
+                <br/>
+                <a href="https://huggingface.co/spaces/tomg-group-umd/lm-watermarking?duplicate=true">
+                <img style="margin-top: 0em; margin-bottom: 0em" src="https://bit.ly/3gLdBN6" alt="Duplicate Space"></a>
+                <p/>
+                """)
+        
+        # Parameter selection group
+        with gr.Accordion("Advanced Settings",open=False):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    gr.Markdown(f"#### Generation Parameters")
+                    with gr.Row():
+                        decoding = gr.Radio(label="Decoding Method",choices=["multinomial", "greedy"], value=("multinomial" if args.use_sampling else "greedy"))
+                    with gr.Row():
+                        sampling_temp = gr.Slider(label="Sampling Temperature", minimum=0.1, maximum=1.0, step=0.1, value=args.sampling_temp, visible=True)
+                    with gr.Row():
+                        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.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)
+                    with gr.Row():
+                        delta = gr.Slider(label="delta",minimum=0.0, maximum=10.0, step=0.1, value=args.delta)
+                    with gr.Row():
+                        ignore_repeated_bigrams = gr.Checkbox(label="Ignore Bigram Repeats")
+                    with gr.Row():
+                        normalizers = gr.CheckboxGroup(label="Normalizations", choices=["unicode", "homoglyphs", "truecase"], value=args.normalizers)
+        
+        # State manager
+        # Construct state for parameters, define updates and toggles, and register event listeners
+        session_args = gr.State(value=args)
+
+        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
+        def update_delta(session_state, value): session_state.delta = float(value); return session_state
+        def update_decoding(session_state, value):
+            if value == "multinomial":
+                session_state.use_sampling = True
+            elif value == "greedy":
+                session_state.use_sampling = False
+            return session_state
+        def toggle_sampling_vis(value):
+            if value == "multinomial":
+                return gr.update(visible=True)
+            elif value == "greedy":
+                return gr.update(visible=False)
+        def toggle_sampling_vis_inv(value):
+            if value == "multinomial":
+                return gr.update(visible=False)
+            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_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
+        
+        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])
+
+        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])
+
+        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])
+
+        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...]"
+                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])
+        
+        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.
                 
-                generate_btn.click(fn=generate, inputs=[prompt], outputs=[redecoded_input, truncation_warning, output_without_watermark, output_with_watermark])
+                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.
                 
-                # Show truncated version of prompt if truncation occurred
-                redecoded_input.change(fn=truncate_prompt, inputs=[redecoded_input,truncation_warning,prompt], outputs=[prompt])
-
-                # Call detection when the outputs of the generate function are updated.
-                output_without_watermark.change(fn=detect, inputs=output_without_watermark, outputs=without_watermark_detection_result)
-                output_with_watermark.change(fn=detect, inputs=output_with_watermark, outputs=with_watermark_detection_result)
-
-            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, inputs=detection_input, outputs=detection_result)
-            
-            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.
-                    """
-                    )
-
-        if args.demo_public:
-            demo.launch(share=True) # exposes app to the internet via randomly generated link
-        else:
-            demo.launch()
+                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.
+                """
+                )
+
+    if args.demo_public:
+        demo.launch(share=True) # exposes app to the internet via randomly generated link
+    else:
+        demo.launch()
+
+def main(args): 
+
+    # Initial arg processing and log
+    args.normalizers = (args.normalizers.split(",") if args.normalizers else [])
+    print(args)
+
+    if not args.skip_model_load:
+        model, tokenizer, device = load_model(args)
+    else:
+        model, tokenizer, device = None, None, []
+
+    # Generate and detect, report to stdout
+    if not args.skip_model_load:
+        # input_text = (
+        # "The diamondback terrapin or simply terrapin (Malaclemys terrapin) is a "
+        # "species of turtle native to the brackish coastal tidal marshes of the "
+        # "Northeastern and southern United States, and in Bermuda.[6] It belongs "
+        # "to the monotypic genus Malaclemys. It has one of the largest ranges of "
+        # "all turtles in North America, stretching as far south as the Florida Keys "
+        # "and as far north as Cape Cod.[7] The name 'terrapin' is derived from the "
+        # "Algonquian word torope.[8] It applies to Malaclemys terrapin in both "
+        # "British English and American English. The name originally was used by "
+        # "early European settlers in North America to describe these brackish-water "
+        # "turtles that inhabited neither freshwater habitats nor the sea. It retains "
+        # "this primary meaning in American English.[8] In British English, however, "
+        # "other semi-aquatic turtle species, such as the red-eared slider, might "
+        # "also be called terrapins. The common name refers to the diamond pattern "
+        # "on top of its shell (carapace), but the overall pattern and coloration "
+        # "vary greatly. The shell is usually wider at the back than in the front, "
+        # "and from above it appears wedge-shaped. The shell coloring can vary "
+        # "from brown to grey, and its body color can be grey, brown, yellow, "
+        # "or white. All have a unique pattern of wiggly, black markings or spots "
+        # "on their body and head. The diamondback terrapin has large webbed "
+        # "feet.[9] The species is"
+        # )
+
+        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]
+        print("#"*term_width)
+        print("Prompt:")
+        print(input_text)
+
+        _, _, decoded_output_without_watermark, decoded_output_with_watermark, _ = generate(input_text, args)
+        without_watermark_detection_result = detect(decoded_output_without_watermark, args)
+        with_watermark_detection_result = detect(decoded_output_with_watermark, args)
+
+        print("#"*term_width)
+        print("Output without watermark:")
+        print(decoded_output_without_watermark)
+        print("-"*term_width)
+        print(f"Detection result @ {args.detection_z_threshold}:")
+        pprint(without_watermark_detection_result)
+        print("-"*term_width)
+
+        print("#"*term_width)
+        print("Output with watermark:")
+        print(decoded_output_with_watermark)
+        print("-"*term_width)
+        print(f"Detection result @ {args.detection_z_threshold}:")
+        pprint(with_watermark_detection_result)
+        print("-"*term_width)
+
+
+    # Launch the app to generate and detect interactively (implements the hf space demo)
+    if args.run_gradio:
+        run_gradio(args, model=model, tokenizer=tokenizer, device=device)
 
     return