try text-generation 0.3.1 release

demo_watermark.py

@@ -16,7 +16,6 @@
 
 import os
 import argparse
-from argparse import Namespace
 from pprint import pprint
 from functools import partial
 
@@ -30,15 +29,25 @@ from transformers import (AutoTokenizer,
                           AutoModelForCausalLM,
                           LogitsProcessorList)
 
+# from local_tokenizers.tokenization_llama import LLaMATokenizer
+
+from transformers import GPT2TokenizerFast
+OPT_TOKENIZER = GPT2TokenizerFast
+
 from watermark_processor import WatermarkLogitsProcessor, WatermarkDetector
 
+
+# ALPACA_MODEL_NAME = "alpaca"
+# ALPACA_MODEL_TOKENIZER = LLaMATokenizer
+# ALPACA_TOKENIZER_PATH = "/cmlscratch/jkirchen/llama"
+
 # FIXME correct lengths for all models
 API_MODEL_MAP = {
-    "bigscience/bloom"        : {"max_length": 2048, "gamma": 0.5, "delta": 2.0},
-    "bigscience/bloomz"       : {"max_length": 2048, "gamma": 0.5, "delta": 2.0},
-    "google/flan-ul2"         : {"max_length": 2048, "gamma": 0.5, "delta": 2.0},
-    "google/flan-t5-xxl"      : {"max_length": 2048, "gamma": 0.5, "delta": 2.0},
-    "EleutherAI/gpt-neox-20b" : {"max_length": 2048, "gamma": 0.5, "delta": 2.0},
+    "bigscience/bloom"        : {"max_length": 1000, "gamma": 0.5, "delta": 2.0},
+    "bigscience/bloomz"       : {"max_length": 1000, "gamma": 0.5, "delta": 2.0},
+    "google/flan-ul2"         : {"max_length": 1000, "gamma": 0.5, "delta": 2.0},
+    "google/flan-t5-xxl"      : {"max_length": 1000, "gamma": 0.5, "delta": 2.0},
+    "EleutherAI/gpt-neox-20b" : {"max_length": 1000, "gamma": 0.5, "delta": 2.0},
 }
 
 def str2bool(v):
@@ -231,35 +240,29 @@ def generate_with_api(prompt, args):
     timeout_msg = "[Model API timeout error. Try reducing the max_new_tokens parameter or the prompt length.]"
     try:
         generation_params["watermark"] = False
-        output = client.generate(prompt, **generation_params)
-        output_text_without_watermark = output.generated_text
+        without_watermark_iterator = client.generate_stream(prompt, **generation_params)
     except ReadTimeout as e:
         print(e)
-        output_text_without_watermark = timeout_msg
+        without_watermark_iterator = (char for char in timeout_msg)
     try:
         generation_params["watermark"] = True
-        output = client.generate(prompt, **generation_params)
-        output_text_with_watermark = output.generated_text
+        with_watermark_iterator = client.generate_stream(prompt, **generation_params)
     except ReadTimeout as e:
         print(e)
-        output_text_with_watermark = timeout_msg
-    
-    return (output_text_without_watermark, 
-            output_text_with_watermark)
+        with_watermark_iterator = (char for char in timeout_msg)
 
+    all_without_words, all_with_words = "", ""
+    for without_word, with_word in zip(without_watermark_iterator, with_watermark_iterator):
+        all_without_words += without_word.token.text
+        all_with_words += with_word.token.text
+        yield all_without_words, all_with_words
 
-def generate(prompt, args, tokenizer, model=None, device=None):
-    """Instatiate the WatermarkLogitsProcessor according to the watermark parameters
-       and generate watermarked text by passing it to the generate method of the model
-       as a logits processor. """
-    
-    print(f"Generating with {args}")
+
+def check_prompt(prompt, args, tokenizer, model=None, device=None):
 
     # This applies to both the local and API model scenarios
-    if args.prompt_max_length:
-        pass
-    elif args.model_name_or_path in API_MODEL_MAP:
-        args.prompt_max_length = API_MODEL_MAP[args.model_name_or_path]["max_length"]-args.max_new_tokens
+    if args.model_name_or_path in API_MODEL_MAP:
+        args.prompt_max_length = API_MODEL_MAP[args.model_name_or_path]["max_length"]
     elif hasattr(model.config,"max_position_embedding"):
         args.prompt_max_length = model.config.max_position_embeddings-args.max_new_tokens
     else:
@@ -269,69 +272,77 @@ def generate(prompt, args, tokenizer, model=None, device=None):
     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]
 
-    if args.model_name_or_path in API_MODEL_MAP:
-        api_outputs = generate_with_api(prompt, args)
-        decoded_output_without_watermark = api_outputs[0]
-        decoded_output_with_watermark = api_outputs[1]
-        return (redecoded_input,
-                int(truncation_warning),
-                decoded_output_without_watermark, 
-                decoded_output_with_watermark,
-                args,
-                tokenizer)
-
+    return (redecoded_input,
+            int(truncation_warning),
+            args)
 
-    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
-        ))
+def generate(prompt, args, tokenizer, model=None, device=None):
+    """Instatiate the WatermarkLogitsProcessor according to the watermark parameters
+       and generate watermarked text by passing it to the generate method of the model
+       as a logits processor. """
 
-    generate_without_watermark = partial(
-        model.generate,
-        **gen_kwargs
-    )
-    generate_with_watermark = partial(
-        model.generate,
-        logits_processor=LogitsProcessorList([watermark_processor]), 
-        **gen_kwargs
-    )
+    print(f"Generating with {args}")
+    print(f"Prompt: {prompt}")
 
-    torch.manual_seed(args.generation_seed)
-    output_without_watermark = generate_without_watermark(**tokd_input)
+    if args.model_name_or_path in API_MODEL_MAP:
+        api_outputs = generate_with_api(prompt, args)
+        yield from api_outputs
+    else:
+        tokd_input = tokenizer(prompt, return_tensors="pt", add_special_tokens=True, truncation=True, max_length=args.prompt_max_length).to(device)
+
+        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
+        )
 
-    # 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)
+        output_without_watermark = generate_without_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]:]
+        # 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)
 
-    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]
+        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]:]
 
-    return (redecoded_input,
-            int(truncation_warning),
-            decoded_output_without_watermark, 
-            decoded_output_with_watermark,
-            args, 
-            tokenizer) 
+        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]
+
+        # mocking the API outputs in a whitespace split generator style
+        all_without_words, all_with_words = "", ""
+        for without_word, with_word in zip(decoded_output_without_watermark.split(), decoded_output_with_watermark.split()):
+            all_without_words += without_word + " "
+            all_with_words += with_word + " "
+            yield all_without_words, all_with_words
 
 
 def format_names(s):
@@ -348,7 +359,6 @@ def format_names(s):
 def list_format_scores(score_dict, detection_threshold):
     """Format the detection metrics into a gradio dataframe input format"""
     lst_2d = []
-    # lst_2d.append(["z-score threshold", f"{detection_threshold}"])
     for k,v in score_dict.items():
         if k=='green_fraction': 
             lst_2d.append([format_names(k), f"{v:.1%}"])
@@ -366,9 +376,10 @@ def list_format_scores(score_dict, detection_threshold):
         lst_2d.insert(-1,["z-score Threshold", f"{detection_threshold}"])
     return lst_2d
 
-def detect(input_text, args, tokenizer, device=None):
+def detect(input_text, args, tokenizer, device=None, return_green_token_mask=True):
     """Instantiate the WatermarkDetection object and call detect on
         the input text returning the scores and outcome of the test"""
+
     print(f"Detecting with {args}")
     print(f"Detection Tokenizer: {type(tokenizer)}")
 
@@ -381,14 +392,14 @@ def detect(input_text, args, tokenizer, device=None):
                                         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:
     error = False
+    green_token_mask = None
     if input_text == "":
         error = True
     else:
-        try: 
-            score_dict = watermark_detector.detect(input_text)
-            # output = str_format_scores(score_dict, watermark_detector.z_threshold)
+        try:
+            score_dict = watermark_detector.detect(input_text, return_green_token_mask=return_green_token_mask)
+            green_token_mask = score_dict.pop("green_token_mask", None)
             output = list_format_scores(score_dict, watermark_detector.z_threshold)
         except ValueError as e:
             print(e)
@@ -396,16 +407,41 @@ def detect(input_text, args, tokenizer, device=None):
     if error:
         output = [["Error","string too short to compute metrics"]]
         output += [["",""] for _ in range(6)]
-    return output, args, tokenizer
+
+    html_output = ""
+    if green_token_mask is not None:
+        # hack bc we need a fast tokenizer with charspan support
+        if "opt" in args.model_name_or_path:
+            tokenizer = OPT_TOKENIZER.from_pretrained(args.model_name_or_path)
+
+        tokens = tokenizer(input_text)
+        if tokens["input_ids"][0] == tokenizer.bos_token_id:
+            tokens["input_ids"] = tokens["input_ids"][1:] # ignore attention mask
+        skip = watermark_detector.min_prefix_len
+        charspans = [tokens.token_to_chars(i) for i in range(skip,len(tokens["input_ids"]))]
+        charspans = [cs for cs in charspans if cs is not None] # remove the special token spans
+
+        if len(charspans) != len(green_token_mask): breakpoint()
+        assert len(charspans) == len(green_token_mask)
+
+        tags = [(f'<span class="green">{input_text[cs.start:cs.end]}</span>' if m else f'<span class="red">{input_text[cs.start:cs.end]}</span>') for cs, m in zip(charspans, green_token_mask)]
+        html_output = f'<p>{" ".join(tags)}</p>'
+
+    return output, args, tokenizer, html_output
 
 def run_gradio(args, model=None, device=None, tokenizer=None):
     """Define and launch the gradio demo interface"""
-    # generate_partial = partial(generate, model=model, device=device, tokenizer=tokenizer)
-    # detect_partial = partial(detect, device=device, tokenizer=tokenizer)
+    check_prompt_partial = partial(check_prompt, model=model, device=device)
     generate_partial = partial(generate, model=model, device=device)
     detect_partial = partial(detect, device=device)
 
-    with gr.Blocks() as demo:
+
+    css = """
+    .green { color: black!important;line-height:1.9em; padding: 0.2em 0.2em; background: #ccffcc; border-radius:0.5rem;}
+    .red { color: black!important;line-height:1.9em; padding: 0.2em 0.2em; background: #ffad99; border-radius:0.5rem;}
+    """
+
+    with gr.Blocks(css=css) as demo:
         # Top section, greeting and instructions
         with gr.Row():
             with gr.Column(scale=9):
@@ -422,7 +458,6 @@ def run_gradio(args, model=None, device=None, tokenizer=None):
                 [![](https://badgen.net/badge/icon/GitHub?icon=github&label)](https://github.com/jwkirchenbauer/lm-watermarking)
                 """
                 )
-                # gr.Markdown(f"Language model: {args.model_name_or_path} {'(float16 mode)' if args.load_fp16 else ''}")
                 # if model_name_or_path at startup not one of the API models then add to dropdown 
                 all_models = sorted(list(set(list(API_MODEL_MAP.keys())+[args.model_name_or_path])))
                 model_selector = gr.Dropdown(
@@ -475,8 +510,10 @@ def run_gradio(args, model=None, device=None, tokenizer=None):
                         """
                         **[Generate & Detect]**: The first tab shows that the watermark can be embedded with 
                         negligible impact on text quality. You can try any prompt and compare the quality of 
-                        normal text (*Output Without Watermark*) to the watermarked text (*Output With Watermark*) below it.   
-                        Metrics on the right show that the watermark can be reliably detected.
+                        normal text (*Output Without Watermark*) to the watermarked text (*Output With Watermark*) below it. 
+                        You can also "see" the watermark by looking at the **Highlighted** tab where the tokens are 
+                        colored green or red depending on which list they are in.
+                        Metrics on the right show that the watermark can be reliably detected given a reasonably small number of tokens (25-50).
                         Detection is very efficient and does not use the language model or its parameters.
 
                         **[Detector Only]**: You can also copy-paste the watermarked text (or any other text) 
@@ -495,7 +532,6 @@ def run_gradio(args, model=None, device=None, tokenizer=None):
                         """
                     )
 
-
         with gr.Tab("Generate & Detect"):
             
             with gr.Row():
@@ -504,15 +540,19 @@ def run_gradio(args, model=None, device=None, tokenizer=None):
                 generate_btn = gr.Button("Generate")
             with gr.Row():
                 with gr.Column(scale=2):
-                    output_without_watermark = gr.Textbox(label="Output Without Watermark", interactive=False,lines=14,max_lines=14)
+                    with gr.Tab("Output Without Watermark (Raw Text)"):
+                        output_without_watermark = gr.Textbox(interactive=False,lines=14,max_lines=14)
+                    with gr.Tab("Highlighted"):
+                        html_without_watermark = gr.HTML(elem_id="html-without-watermark")
                 with gr.Column(scale=1):
-                    # without_watermark_detection_result = gr.Textbox(label="Detection Result", interactive=False,lines=14,max_lines=14)
                     without_watermark_detection_result = gr.Dataframe(headers=["Metric", "Value"], interactive=False,row_count=7,col_count=2)
             with gr.Row():
                 with gr.Column(scale=2):
-                    output_with_watermark = gr.Textbox(label="Output With Watermark", interactive=False,lines=14,max_lines=14)
+                    with gr.Tab("Output With Watermark (Raw Text)"):
+                        output_with_watermark = gr.Textbox(interactive=False,lines=14,max_lines=14)
+                    with gr.Tab("Highlighted"):
+                        html_with_watermark = gr.HTML(elem_id="html-with-watermark")
                 with gr.Column(scale=1):
-                    # with_watermark_detection_result = gr.Textbox(label="Detection Result", interactive=False,lines=14,max_lines=14)
                     with_watermark_detection_result = gr.Dataframe(headers=["Metric", "Value"],interactive=False,row_count=7,col_count=2)
 
             redecoded_input = gr.Textbox(visible=False)
@@ -528,7 +568,6 @@ def run_gradio(args, model=None, device=None, tokenizer=None):
                 with gr.Column(scale=2):
                     detection_input = gr.Textbox(label="Text to Analyze", interactive=True,lines=14,max_lines=14)
                 with gr.Column(scale=1):
-                    # detection_result = gr.Textbox(label="Detection Result", interactive=False,lines=14,max_lines=14)
                     detection_result = gr.Dataframe(headers=["Metric", "Value"], interactive=False,row_count=7,col_count=2)
             with gr.Row():
                     detect_btn = gr.Button("Detect")
@@ -562,7 +601,6 @@ 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)
-            # with gr.Accordion("Actual submitted parameters:",open=False):
             with gr.Row():
                 gr.Markdown(f"_Note: sliders don't always update perfectly. Clicking on the bar or using the number window to the right can help. Window below shows the current settings._")
             with gr.Row():
@@ -657,15 +695,14 @@ def run_gradio(args, model=None, device=None, tokenizer=None):
                 <p/>
                 """)
         
-        # Register main generation tab click, outputing generations as well as a the encoded+redecoded+potentially truncated prompt and flag
-        generate_btn.click(fn=generate_partial, inputs=[prompt,session_args,session_tokenizer], outputs=[redecoded_input, truncation_warning, output_without_watermark, output_with_watermark,session_args,session_tokenizer])
+        # Register main generation tab click, outputing generations as well as a the encoded+redecoded+potentially truncated prompt and flag, then call detection
+        generate_btn.click(fn=check_prompt_partial, inputs=[prompt,session_args,session_tokenizer], outputs=[redecoded_input, truncation_warning, session_args]).success(
+                           fn=generate_partial, inputs=[redecoded_input,session_args,session_tokenizer], outputs=[output_without_watermark, output_with_watermark]).success(
+                           fn=detect_partial, inputs=[output_without_watermark,session_args,session_tokenizer], outputs=[without_watermark_detection_result,session_args,session_tokenizer,html_without_watermark]).success(
+                           fn=detect_partial, inputs=[output_with_watermark,session_args,session_tokenizer], outputs=[with_watermark_detection_result,session_args,session_tokenizer,html_with_watermark])
         # 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,session_tokenizer], outputs=[without_watermark_detection_result,session_args,session_tokenizer])
-        output_with_watermark.change(fn=detect_partial, inputs=[output_with_watermark,session_args,session_tokenizer], outputs=[with_watermark_detection_result,session_args,session_tokenizer])
         # Register main detection tab click
-        # detect_btn.click(fn=detect_partial, inputs=[detection_input,session_args,session_tokenizer], outputs=[detection_result, session_args,session_tokenizer])
         detect_btn.click(fn=detect_partial, inputs=[detection_input,session_args,session_tokenizer], outputs=[detection_result, session_args,session_tokenizer], api_name="detection")
 
         # State management logic
@@ -727,7 +764,11 @@ def run_gradio(args, model=None, device=None, tokenizer=None):
         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
-        def update_tokenizer(model_name_or_path): return AutoTokenizer.from_pretrained(model_name_or_path)
+        def update_tokenizer(model_name_or_path): 
+            # if model_name_or_path == ALPACA_MODEL_NAME:
+            #     return ALPACA_MODEL_TOKENIZER.from_pretrained(ALPACA_TOKENIZER_PATH)
+            # else:
+            return AutoTokenizer.from_pretrained(model_name_or_path)
         # registering callbacks for toggling the visibilty of certain parameters based on the values of others
         decoding.change(toggle_sampling_vis,inputs=[decoding], outputs=[sampling_temp])
         decoding.change(toggle_sampling_vis,inputs=[decoding], outputs=[generation_seed])
@@ -829,19 +870,6 @@ def main(args):
         "on their body and head. The diamondback terrapin has large webbed "
         "feet.[9] The species is"
     )
-    
-    # teaser example
-    # 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"
-    # )
 
     args.default_prompt = input_text
     
@@ -854,19 +882,28 @@ def main(args):
         print("Prompt:")
         print(input_text)
 
-        _, _, decoded_output_without_watermark, decoded_output_with_watermark, _, _ = generate(input_text, 
-                                                                                            args, 
-                                                                                            model=model, 
-                                                                                            device=device, 
-                                                                                            tokenizer=tokenizer)
+        # a generator that yields (without_watermark, with_watermark) pairs
+        generator_outputs = generate(input_text, 
+                                    args, 
+                                    model=model, 
+                                    device=device, 
+                                    tokenizer=tokenizer)
+        # we need to iterate over it, 
+        # but we only want the last output in this case
+        for out in generator_outputs:
+            decoded_output_without_watermark = out[0]
+            decoded_output_with_watermark = out[1]
+
         without_watermark_detection_result = detect(decoded_output_without_watermark, 
                                                     args, 
                                                     device=device, 
-                                                    tokenizer=tokenizer)
+                                                    tokenizer=tokenizer,
+                                                    return_green_token_mask=False)
         with_watermark_detection_result = detect(decoded_output_with_watermark, 
                                                  args, 
                                                  device=device, 
-                                                 tokenizer=tokenizer)
+                                                 tokenizer=tokenizer,
+                                                 return_green_token_mask=False)
 
         print("#"*term_width)
         print("Output without watermark:")

requirements.txt

@@ -1,9 +1,9 @@
 spacy
-gradio
+gradio>=3.21.0
 nltk
 scipy
 torch
 transformers
 tokenizers
 accelerate
-text-generation>=0.3.0
+text-generation>=0.3.1