demo_watermark.py

# coding=utf-8
# Copyright 2023 Authors of "A Watermark for Large Language Models" 
# available at https://arxiv.org/abs/2301.10226
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import argparse
from pprint import pprint
from functools import partial

import torch

from transformers import (AutoTokenizer,
                          AutoModelForSeq2SeqLM,
                          AutoModelForCausalLM,
                          LogitsProcessorList)

from watermark_processor import WatermarkLogitsProcessor, WatermarkDetector

def str2bool(v):
    if isinstance(v, bool):
        return v
    if v.lower() in ('yes', 'true', 't', 'y', '1'):
        return True
    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
        return False
    else:
        raise argparse.ArgumentTypeError('Boolean value expected.')

def parse_args():

    parser = argparse.ArgumentParser(description="A minimum working example of applying the watermark to any LLM that supports the huggingface 🤗 `generate` API")

    parser.add_argument(
        "--run_gradio",
        type=str2bool,
        default=False,
        help="Whether to launch as a gradio demo.",
    )
    parser.add_argument(
        "--demo_public",
        type=str2bool,
        default=False,
        help="Whether to expose the gradio demo to the internet.",
    )
    parser.add_argument(
        "--model_name_or_path",
        type=str,
        default="facebook/opt-6.7b",
        help="Main model, path to pretrained model or model identifier from huggingface.co/models.",
    )
    parser.add_argument(
        "--prompt_max_length",
        type=int,
        default=None,
        help="Truncation length for prompt, overrides model config's max length field.",
    )
    parser.add_argument(
        "--max_new_tokens",
        type=int,
        default=200,
        help="Maximmum number of new tokens to generate.",
    )
    parser.add_argument(
        "--generation_seed",
        type=int,
        default=123,
        help="Seed for setting the torch global rng prior to generation.",
    )
    parser.add_argument(
        "--use_sampling",
        type=str2bool,
        default=True,
        help="Whether to generate using multinomial sampling.",
    )
    parser.add_argument(
        "--sampling_temp",
        type=float,
        default=0.7,
        help="Sampling temperature to use when generating using multinomial sampling.",
    )
    parser.add_argument(
        "--use_gpu",
        type=str2bool,
        default=True,
        help="Whether to run inference and watermark hashing/seeding/permutation on gpu.",
    )
    parser.add_argument(
        "--seeding_scheme",
        type=str,
        default="markov_1",
        help="Seeding scheme to use to generate the greenlists at each generation and verification step.",
    )
    parser.add_argument(
        "--gamma",
        type=float,
        default=0.25,
        help="The fraction of the vocabulary to partition into the greenlist at each generation and verification step.",
    )
    parser.add_argument(
        "--delta",
        type=float,
        default=2.0,
        help="The amount/bias to add to each of the greenlist token logits before each token sampling step.",
    )
    parser.add_argument(
        "--normalizers",
        type=str,
        default="",
        help="Single or comma separated list of the preprocessors/normalizer names to use when performing watermark detection.",
    )
    parser.add_argument(
        "--ignore_repeated_bigrams",
        type=str2bool,
        default=False,
        help="Whether to use the detection method that only counts each unqiue bigram once as either a green or red hit.",
    )
    parser.add_argument(
        "--detection_z_threshold",
        type=float,
        default=4.0,
        help="The test statistic threshold for the detection hypothesis test.",
    )
    parser.add_argument(
        "--select_green_tokens",
        type=str2bool,
        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.",
    )
    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:
        model = AutoModelForSeq2SeqLM.from_pretrained(args.model_name_or_path)
    elif 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}")

    if args.use_gpu:
        device = "cuda" if torch.cuda.is_available() else "cpu"
        model = model.to(device)
    else:
        device = "cpu"
    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

    # 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)
    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)

    # Launch the app to generate and detect interactively (implements the hf space demo)

    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
                
                generate_btn.click(fn=generate, inputs=[prompt], outputs=[redecoded_input, truncation_warning, output_without_watermark, output_with_watermark])
                
                # 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()

    return

if __name__ == "__main__":

    args = parse_args()
    print(args)

    main(args)