adding files

README.md

@@ -1,13 +1,14 @@
 ---
-title: Lm Watermarking
-emoji: 🌍
-colorFrom: green
-colorTo: yellow
+title: A Watermark for LLMs
+emoji: 💧
+colorFrom: blue
+colorTo: purple
 sdk: gradio
 sdk_version: 3.18.0
 app_file: app.py
 pinned: false
 license: apache-2.0
+python_version: 3.10.6
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,41 @@
+# 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.
+
+from argparse import Namespace
+args = Namespace()
+
+arg_dict = {  
+    "run_gradio": True,
+    "model_name_or_path": "facebook/opt-125m",
+    # "model_name_or_path": "facebook/opt-1.3b",
+    # "model_name_or_path": "facebook/opt-2.7b",
+    "max_new_tokens": 200,
+    "use_sampling": True,
+    "sampling_temp": 0.7,
+    "use_gpu": True,
+    "seeding_scheme": "markov_1",
+    "gamma": 0.25,
+    "delta": 2.0, 
+    "normalizers": "",
+    "ignore_repeated_bigrams": False,
+}
+
+args.__dict__.update(arg_dict)
+print(args)
+
+from demo_watermark import main
+
+main(args)

demo_watermark.py

@@ -0,0 +1,379 @@
+# 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)

homoglyphs.py

@@ -0,0 +1,268 @@
+"""Updated version of core.py from
+https://github.com/yamatt/homoglyphs/tree/main/homoglyphs_fork
+for modern python3
+"""
+
+from collections import defaultdict
+import json
+from itertools import product
+import os
+import unicodedata
+
+import homoglyphs_fork as hg
+
+CURRENT_DIR = hg.core.CURRENT_DIR
+
+# Actions if char not in alphabet
+STRATEGY_LOAD = 1  # load category for this char
+STRATEGY_IGNORE = 2  # add char to result
+STRATEGY_REMOVE = 3  # remove char from result
+
+ASCII_RANGE = range(128)
+
+
+class Categories:
+    """
+    Work with aliases from ISO 15924.
+    https://en.wikipedia.org/wiki/ISO_15924#List_of_codes
+    """
+
+    fpath = os.path.join(CURRENT_DIR, "categories.json")
+
+    @classmethod
+    def _get_ranges(cls, categories):
+        """
+        :return: iter: (start code, end code)
+        :rtype: list
+        """
+        with open(cls.fpath, encoding="utf-8") as f:
+            data = json.load(f)
+
+        for category in categories:
+            if category not in data["aliases"]:
+                raise ValueError("Invalid category: {}".format(category))
+
+        for point in data["points"]:
+            if point[2] in categories:
+                yield point[:2]
+
+    @classmethod
+    def get_alphabet(cls, categories):
+        """
+        :return: set of chars in alphabet by categories list
+        :rtype: set
+        """
+        alphabet = set()
+        for start, end in cls._get_ranges(categories):
+            chars = (chr(code) for code in range(start, end + 1))
+            alphabet.update(chars)
+        return alphabet
+
+    @classmethod
+    def detect(cls, char):
+        """
+        :return: category
+        :rtype: str
+        """
+        with open(cls.fpath, encoding="utf-8") as f:
+            data = json.load(f)
+
+        # try detect category by unicodedata
+        try:
+            category = unicodedata.name(char).split()[0]
+        except TypeError:
+            # In Python2 unicodedata.name raise error for non-unicode chars
+            pass
+        else:
+            if category in data["aliases"]:
+                return category
+
+        # try detect category by ranges from JSON file.
+        code = ord(char)
+        for point in data["points"]:
+            if point[0] <= code <= point[1]:
+                return point[2]
+
+    @classmethod
+    def get_all(cls):
+        with open(cls.fpath, encoding="utf-8") as f:
+            data = json.load(f)
+        return set(data["aliases"])
+
+
+class Languages:
+    fpath = os.path.join(CURRENT_DIR, "languages.json")
+
+    @classmethod
+    def get_alphabet(cls, languages):
+        """
+        :return: set of chars in alphabet by languages list
+        :rtype: set
+        """
+        with open(cls.fpath, encoding="utf-8") as f:
+            data = json.load(f)
+        alphabet = set()
+        for lang in languages:
+            if lang not in data:
+                raise ValueError("Invalid language code: {}".format(lang))
+            alphabet.update(data[lang])
+        return alphabet
+
+    @classmethod
+    def detect(cls, char):
+        """
+        :return: set of languages which alphabet contains passed char.
+        :rtype: set
+        """
+        with open(cls.fpath, encoding="utf-8") as f:
+            data = json.load(f)
+        languages = set()
+        for lang, alphabet in data.items():
+            if char in alphabet:
+                languages.add(lang)
+        return languages
+
+    @classmethod
+    def get_all(cls):
+        with open(cls.fpath, encoding="utf-8") as f:
+            data = json.load(f)
+        return set(data.keys())
+
+
+class Homoglyphs:
+    def __init__(
+        self,
+        categories=None,
+        languages=None,
+        alphabet=None,
+        strategy=STRATEGY_IGNORE,
+        ascii_strategy=STRATEGY_IGNORE,
+        ascii_range=ASCII_RANGE,
+    ):
+        # strategies
+        if strategy not in (STRATEGY_LOAD, STRATEGY_IGNORE, STRATEGY_REMOVE):
+            raise ValueError("Invalid strategy")
+        self.strategy = strategy
+        self.ascii_strategy = ascii_strategy
+        self.ascii_range = ascii_range
+
+        # Homoglyphs must be initialized by any alphabet for correct work
+        if not categories and not languages and not alphabet:
+            categories = ("LATIN", "COMMON")
+
+        # cats and langs
+        self.categories = set(categories or [])
+        self.languages = set(languages or [])
+
+        # alphabet
+        self.alphabet = set(alphabet or [])
+        if self.categories:
+            alphabet = Categories.get_alphabet(self.categories)
+            self.alphabet.update(alphabet)
+        if self.languages:
+            alphabet = Languages.get_alphabet(self.languages)
+            self.alphabet.update(alphabet)
+        self.table = self.get_table(self.alphabet)
+
+    @staticmethod
+    def get_table(alphabet):
+        table = defaultdict(set)
+        # removed CURRENT_DIR here:
+        with open(os.path.join("confusables_sept2022.json")) as f:
+            data = json.load(f)
+        for char in alphabet:
+            if char in data:
+                for homoglyph in data[char]:
+                    if homoglyph in alphabet:
+                        table[char].add(homoglyph)
+        return table
+
+    @staticmethod
+    def get_restricted_table(source_alphabet, target_alphabet):
+        table = defaultdict(set)
+        # removed CURRENT_DIR here:
+        with open(os.path.join("confusables_sept2022.json")) as f:
+            data = json.load(f)
+        for char in source_alphabet:
+            if char in data:
+                for homoglyph in data[char]:
+                    if homoglyph in target_alphabet:
+                        table[char].add(homoglyph)
+        return table
+
+    @staticmethod
+    def uniq_and_sort(data):
+        result = list(set(data))
+        result.sort(key=lambda x: (-len(x), x))
+        return result
+
+    def _update_alphabet(self, char):
+        # try detect languages
+        langs = Languages.detect(char)
+        if langs:
+            self.languages.update(langs)
+            alphabet = Languages.get_alphabet(langs)
+            self.alphabet.update(alphabet)
+        else:
+            # try detect categories
+            category = Categories.detect(char)
+            if category is None:
+                return False
+            self.categories.add(category)
+            alphabet = Categories.get_alphabet([category])
+            self.alphabet.update(alphabet)
+        # update table for new alphabet
+        self.table = self.get_table(self.alphabet)
+        return True
+
+    def _get_char_variants(self, char):
+        if char not in self.alphabet:
+            if self.strategy == STRATEGY_LOAD:
+                if not self._update_alphabet(char):
+                    return []
+            elif self.strategy == STRATEGY_IGNORE:
+                return [char]
+            elif self.strategy == STRATEGY_REMOVE:
+                return []
+
+        # find alternative chars for current char
+        alt_chars = self.table.get(char, set())
+        if alt_chars:
+            # find alternative chars for alternative chars for current char
+            alt_chars2 = [self.table.get(alt_char, set()) for alt_char in alt_chars]
+            # combine all alternatives
+            alt_chars.update(*alt_chars2)
+        # add current char to alternatives
+        alt_chars.add(char)
+
+        # uniq, sort and return
+        return self.uniq_and_sort(alt_chars)
+
+    def _get_combinations(self, text, ascii=False):
+        variations = []
+        for char in text:
+            alt_chars = self._get_char_variants(char)
+
+            if ascii:
+                alt_chars = [
+                    char for char in alt_chars if ord(char) in self.ascii_range
+                ]
+                if not alt_chars and self.ascii_strategy == STRATEGY_IGNORE:
+                    return
+
+            if alt_chars:
+                variations.append(alt_chars)
+        if variations:
+            for variant in product(*variations):
+                yield "".join(variant)
+
+    def get_combinations(self, text):
+        return list(self._get_combinations(text))
+
+    def _to_ascii(self, text):
+        for variant in self._get_combinations(text, ascii=True):
+            if max(map(ord, variant)) in self.ascii_range:
+                yield variant
+
+    def to_ascii(self, text):
+        return self.uniq_and_sort(self._to_ascii(text))

normalizers.py

@@ -0,0 +1,195 @@
+""" Text-based normalizers, used to mitigate simple attacks against watermarking.
+
+This implementation is unlikely to be a complete list of all possible exploits within the unicode standard,
+it represents our best effort at the time of writing.
+
+These normalizers can be used as stand-alone normalizers. They could be made to conform to HF tokenizers standard, but that would
+require messing with the limited rust interface of tokenizers.NormalizedString
+"""
+from collections import defaultdict
+from functools import cache
+
+import re
+import unicodedata
+import homoglyphs as hg
+
+
+def normalization_strategy_lookup(strategy_name: str) -> object:
+    if strategy_name == "unicode":
+        return UnicodeSanitizer()
+    elif strategy_name == "homoglyphs":
+        return HomoglyphCanonizer()
+    elif strategy_name == "truecase":
+        return TrueCaser()
+
+
+class HomoglyphCanonizer:
+    """Attempts to detect homoglyph attacks and find a consistent canon.
+
+    This function does so on a per-ISO-category level. Language-level would also be possible (see commented code).
+    """
+
+    def __init__(self):
+        self.homoglyphs = None
+
+    def __call__(self, homoglyphed_str: str) -> str:
+        # find canon:
+        target_category, all_categories = self._categorize_text(homoglyphed_str)
+        homoglyph_table = self._select_canon_category_and_load(target_category, all_categories)
+        return self._sanitize_text(target_category, homoglyph_table, homoglyphed_str)
+
+    def _categorize_text(self, text: str) -> dict:
+        iso_categories = defaultdict(int)
+        # self.iso_languages = defaultdict(int)
+
+        for char in text:
+            iso_categories[hg.Categories.detect(char)] += 1
+            # for lang in hg.Languages.detect(char):
+            #     self.iso_languages[lang] += 1
+        target_category = max(iso_categories, key=iso_categories.get)
+        all_categories = tuple(iso_categories)
+        return target_category, all_categories
+
+    @cache
+    def _select_canon_category_and_load(self, target_category: str, all_categories: tuple[str]) -> dict:
+        homoglyph_table = hg.Homoglyphs(categories=(target_category, "COMMON"))  # alphabet loaded here from file
+
+        source_alphabet = hg.Categories.get_alphabet(all_categories)
+        restricted_table = homoglyph_table.get_restricted_table(source_alphabet, homoglyph_table.alphabet)  # table loaded here from file
+        return restricted_table
+
+    def _sanitize_text(self, target_category: str, homoglyph_table: dict, homoglyphed_str: str) -> str:
+        sanitized_text = ""
+        for char in homoglyphed_str:
+            # langs = hg.Languages.detect(char)
+            cat = hg.Categories.detect(char)
+            if target_category in cat or "COMMON" in cat or len(cat) == 0:
+                sanitized_text += char
+            else:
+                sanitized_text += list(homoglyph_table[char])[0]
+        return sanitized_text
+
+
+class UnicodeSanitizer:
+    """Regex-based unicode sanitzer. Has different levels of granularity.
+
+    * ruleset="whitespaces"    - attempts to remove only whitespace unicode characters
+    * ruleset="IDN.blacklist"  - does its best to remove unusual unicode based on  Network.IDN.blacklist characters
+    * ruleset="ascii"          - brute-forces all text into ascii
+
+    This is unlikely to be a comprehensive list.
+
+    You can find a more comprehensive discussion at https://www.unicode.org/reports/tr36/
+    and https://www.unicode.org/faq/security.html
+    """
+
+    def __init__(self, ruleset="whitespaces"):
+        if ruleset == "whitespaces":
+
+            """Documentation:
+            \u00A0: Non-breaking space
+            \u1680: Ogham space mark
+            \u180E: Mongolian vowel separator
+            \u2000-\u200B: Various space characters, including en space, em space, thin space, hair space, zero-width space, and zero-width non-joiner
+            \u200C\u200D: Zero-width non-joiner and zero-width joiner
+            \u200E,\u200F: Left-to-right-mark, Right-to-left-mark
+            \u2060: Word joiner
+            \u2063: Invisible separator
+            \u202F: Narrow non-breaking space
+            \u205F: Medium mathematical space
+            \u3000: Ideographic space
+            \uFEFF: Zero-width non-breaking space
+            \uFFA0: Halfwidth hangul filler
+            \uFFF9\uFFFA\uFFFB: Interlinear annotation characters
+            \uFE00-\uFE0F: Variation selectors
+            \u202A-\u202F: Embedding characters
+            \u3164: Korean hangul filler.
+
+            Note that these characters are not always superfluous whitespace characters!
+            """
+
+            self.pattern = re.compile(
+                r"[\u00A0\u1680\u180E\u2000-\u200B\u200C\u200D\u200E\u200F\u2060\u2063\u202F\u205F\u3000\uFEFF\uFFA0\uFFF9\uFFFA\uFFFB"
+                r"\uFE00\uFE01\uFE02\uFE03\uFE04\uFE05\uFE06\uFE07\uFE08\uFE09\uFE0A\uFE0B\uFE0C\uFE0D\uFE0E\uFE0F\u3164\u202A\u202B\u202C\u202D"
+                r"\u202E\u202F]"
+            )
+        elif ruleset == "IDN.blacklist":
+
+            """Documentation:
+            [\u00A0\u1680\u180E\u2000-\u200B\u202F\u205F\u2060\u2063\uFEFF]: Matches any whitespace characters in the Unicode character
+                        set that are included in the IDN blacklist.
+            \uFFF9-\uFFFB: Matches characters that are not defined in Unicode but are used as language tags in various legacy encodings.
+                        These characters are not allowed in domain names.
+            \uD800-\uDB7F: Matches the first part of a surrogate pair. Surrogate pairs are used to represent characters in the Unicode character
+                        set that cannot be represented by a single 16-bit value. The first part of a surrogate pair is in the range U+D800 to U+DBFF,
+                        and the second part is in the range U+DC00 to U+DFFF.
+            \uDB80-\uDBFF][\uDC00-\uDFFF]?: Matches the second part of a surrogate pair. The second part of a surrogate pair is in the range U+DC00
+                        to U+DFFF, and is optional.
+            [\uDB40\uDC20-\uDB40\uDC7F][\uDC00-\uDFFF]: Matches certain invalid UTF-16 sequences which should not appear in IDNs.
+            """
+
+            self.pattern = re.compile(
+                r"[\u00A0\u1680\u180E\u2000-\u200B\u202F\u205F\u2060\u2063\uFEFF\uFFF9-\uFFFB\uD800-\uDB7F\uDB80-\uDBFF]"
+                r"[\uDC00-\uDFFF]?|[\uDB40\uDC20-\uDB40\uDC7F][\uDC00-\uDFFF]"
+            )
+        else:
+            """Documentation:
+            This is a simple restriction to "no-unicode", using only ascii characters. Control characters are included.
+            """
+            self.pattern = re.compile(r"[^\x00-\x7F]+")
+
+    def __call__(self, text: str) -> str:
+        text = unicodedata.normalize("NFC", text)  # canon forms
+        text = self.pattern.sub(" ", text)  # pattern match
+        text = re.sub(" +", " ", text)  # collapse whitespaces
+        text = "".join(c for c in text if unicodedata.category(c) != "Cc")  # Remove any remaining non-printable characters
+        return text
+
+
+class TrueCaser:
+    """True-casing, is a capitalization normalization that returns text to its original capitalization.
+
+    This defends against attacks that wRIte TeXt lIkE spOngBoB.
+
+    Here, a simple POS-tagger is used.
+    """
+
+    uppercase_pos = ["PROPN"]  # Name POS tags that should be upper-cased
+
+    def __init__(self, backend="spacy"):
+        if backend == "spacy":
+            import spacy
+
+            self.nlp = spacy.load("en_core_web_sm")
+            self.normalize_fn = self._spacy_truecasing
+        else:
+            from nltk import pos_tag, word_tokenize  # noqa
+            import nltk
+
+            nltk.download("punkt")
+            nltk.download("averaged_perceptron_tagger")
+            nltk.download("universal_tagset")
+            self.normalize_fn = self._nltk_truecasing
+
+    def __call__(self, random_capitalized_string: str) -> str:
+        truecased_str = self.normalize_fn(random_capitalized_string)
+        return truecased_str
+
+    def _spacy_truecasing(self, random_capitalized_string: str):
+        doc = self.nlp(random_capitalized_string.lower())
+        POS = self.uppercase_pos
+        truecased_str = "".join([w.text_with_ws.capitalize() if w.pos_ in POS or w.is_sent_start else w.text_with_ws for w in doc])
+        return truecased_str
+
+    def _nltk_truecasing(self, random_capitalized_string: str):
+        from nltk import pos_tag, word_tokenize
+        import nltk
+
+        nltk.download("punkt")
+        nltk.download("averaged_perceptron_tagger")
+        nltk.download("universal_tagset")
+        POS = ["NNP", "NNPS"]
+
+        tagged_text = pos_tag(word_tokenize(random_capitalized_string.lower()))
+        truecased_str = " ".join([w.capitalize() if p in POS else w for (w, p) in tagged_text])
+        return truecased_str

requirements.txt

@@ -0,0 +1,6 @@
+homoglyphs_fork
+nltk
+scipy
+torch
+transformers
+tokenizers

watermark_processor.py

@@ -0,0 +1,281 @@
+# 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.
+
+from __future__ import annotations
+import collections
+from math import sqrt
+
+import scipy.stats
+
+import torch
+from torch import Tensor
+from tokenizers import Tokenizer
+from transformers import LogitsProcessor
+
+from nltk.util import ngrams
+
+from normalizers import normalization_strategy_lookup
+
+class WatermarkBase:
+    def __init__(
+        self,
+        vocab: list[int] = None,
+        gamma: float = 0.5,
+        delta: float = 2.0,
+        seeding_scheme: str = "markov_1",  # mostly unused/always default
+        hash_key: int = 15485863,  # just a large prime number to create a rng seed with sufficient bit width
+        select_green_tokens: bool = True,
+    ):
+
+        # watermarking parameters
+        self.vocab = vocab
+        self.vocab_size = len(vocab)
+        self.gamma = gamma
+        self.delta = delta
+        self.seeding_scheme = seeding_scheme
+        self.rng = None
+        self.hash_key = hash_key
+        self.select_green_tokens = select_green_tokens
+
+    def _seed_rng(self, input_ids: torch.LongTensor, seeding_scheme: str = None) -> None:
+        # can optionally override the seeding scheme,
+        # but uses the instance attr by default
+        if seeding_scheme is None:
+            seeding_scheme = self.seeding_scheme
+
+        if seeding_scheme == "markov_1":
+            assert input_ids.shape[-1] >= 1, f"seeding_scheme={seeding_scheme} requires at least a 1 token prefix sequence to seed rng"
+            prev_token = input_ids[-1].item()
+            self.rng.manual_seed(self.hash_key * prev_token)
+        else:
+            raise NotImplementedError(f"Unexpected seeding_scheme: {seeding_scheme}")
+        return
+
+    def _get_greenlist_ids(self, input_ids: torch.LongTensor) -> list[int]:
+        # seed the rng using the previous tokens/prefix
+        # according to the seeding_scheme
+        self._seed_rng(input_ids)
+
+        greenlist_size = int(self.vocab_size * self.gamma)
+        vocab_permutation = torch.randperm(self.vocab_size, device=input_ids.device, generator=self.rng)
+        if self.select_green_tokens: # directly
+            greenlist_ids = vocab_permutation[:greenlist_size] # new
+        else: # select green via red
+            greenlist_ids = vocab_permutation[(self.vocab_size - greenlist_size) :]  # legacy behavior
+        return greenlist_ids
+
+
+class WatermarkLogitsProcessor(WatermarkBase, LogitsProcessor):
+
+    # FIXME maybe make this explict instead of args/kwargs
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def _calc_greenlist_mask(self, scores: torch.FloatTensor, greenlist_token_ids) -> torch.BoolTensor:
+        # TODO lets see if we can lose this loop
+        green_tokens_mask = torch.zeros_like(scores)
+        for b_idx in range(len(greenlist_token_ids)):
+            green_tokens_mask[b_idx][greenlist_token_ids[b_idx]] = 1
+        final_mask = green_tokens_mask.bool()
+        return final_mask
+
+    def _bias_greenlist_logits(self, scores: torch.Tensor, greenlist_mask: torch.Tensor, greenlist_bias: float) -> torch.Tensor:
+        scores[greenlist_mask] = scores[greenlist_mask] + greenlist_bias
+        return scores
+
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
+
+        # this is lazy to allow us to colocate on the watermarked model's device
+        if self.rng is None:
+            self.rng = torch.Generator(device=input_ids.device)
+
+        # NOTE, it would be nice to get rid of this batch loop, but currently,
+        # the seed and partition operations are not tensor/vectorized, thus
+        # each sequence in the batch needs to be treated separately.
+        batched_greenlist_ids = [None for _ in range(input_ids.shape[0])]
+
+        for b_idx in range(input_ids.shape[0]):
+            greenlist_ids = self._get_greenlist_ids(input_ids[b_idx])
+            batched_greenlist_ids[b_idx] = greenlist_ids
+
+        green_tokens_mask = self._calc_greenlist_mask(scores=scores, greenlist_token_ids=batched_greenlist_ids)
+
+        scores = self._bias_greenlist_logits(scores=scores, greenlist_mask=green_tokens_mask, greenlist_bias=self.delta)
+        return scores
+
+
+class WatermarkDetector(WatermarkBase):
+    def __init__(
+        self,
+        *args,
+        device: torch.device = None,
+        tokenizer: Tokenizer = None,
+        z_threshold: float = 4.0,
+        normalizers: list[str] = ["unicode"],  # or also: ["unicode", "homoglyphs", "truecase"]
+        ignore_repeated_bigrams: bool = False,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+        # also configure the metrics returned/preprocessing options
+        assert device, "Must pass device"
+        assert tokenizer, "Need an instance of the generating tokenizer to perform detection"
+
+        self.tokenizer = tokenizer
+        self.device = device
+        self.z_threshold = z_threshold
+        self.rng = torch.Generator(device=self.device)
+
+        if self.seeding_scheme == "markov_1":
+            self.min_prefix_len = 1
+        else:
+            raise NotImplementedError(f"Unexpected seeding_scheme: {self.seeding_scheme}")
+
+        self.normalizers = []
+        for normalization_strategy in normalizers:
+            self.normalizers.append(normalization_strategy_lookup(normalization_strategy))
+        
+        self.ignore_repeated_bigrams = ignore_repeated_bigrams
+        if self.ignore_repeated_bigrams: 
+            assert self.seeding_scheme == "markov_1", "No repeated bigram credit variant assumes the single token seeding scheme."
+
+
+    def _compute_z_score(self, observed_count, T):
+        # count refers to number of green tokens, T is total number of tokens
+        expected_count = self.gamma
+        numer = observed_count - expected_count * T
+        denom = sqrt(T * expected_count * (1 - expected_count))
+        z = numer / denom
+        return z
+
+    def _compute_p_value(self, z):
+        p_value = scipy.stats.norm.sf(z)
+        return p_value
+
+    def _score_sequence(
+        self,
+        input_ids: Tensor,
+        return_num_tokens_scored: bool = True,
+        return_num_green_tokens: bool = True,
+        return_green_fraction: bool = True,
+        return_green_token_mask: bool = False,
+        return_z_score: bool = True,
+        return_p_value: bool = True,
+    ):
+        if self.ignore_repeated_bigrams:
+            # Method that only counts a green/red hit once per unique bigram.
+            # New num total tokens scored (T) becomes the number unique bigrams.
+            # We iterate over all unqiue token bigrams in the input, computing the greenlist
+            # induced by the first token in each, and then checking whether the second
+            # token falls in that greenlist.
+            assert return_green_token_mask == False, "Can't return the green/red mask when ignoring repeats."
+            bigram_table = {}
+            token_bigram_generator = ngrams(input_ids.cpu().tolist(), 2)
+            freq = collections.Counter(token_bigram_generator)
+            num_tokens_scored = len(freq.keys())
+            for idx, bigram in enumerate(freq.keys()):
+                prefix = torch.tensor([bigram[0]], device=self.device) # expects a 1-d prefix tensor on the randperm device
+                greenlist_ids = self._get_greenlist_ids(prefix)
+                bigram_table[bigram] = True if bigram[1] in greenlist_ids else False
+            green_token_count = sum(bigram_table.values())
+        else:
+            num_tokens_scored = len(input_ids) - self.min_prefix_len
+            if num_tokens_scored < 1:
+                raise ValueError((f"Must have at least {1} token to score after "
+                                f"the first min_prefix_len={self.min_prefix_len} tokens required by the seeding scheme."))
+            # Standard method.
+            # Since we generally need at least 1 token (for the simplest scheme)
+            # we start the iteration over the token sequence with a minimum 
+            # num tokens as the first prefix for the seeding scheme,
+            # and at each step, compute the greenlist induced by the
+            # current prefix and check if the current token falls in the greenlist.
+            green_token_count, green_token_mask = 0, []
+            for idx in range(self.min_prefix_len, len(input_ids)):
+                curr_token = input_ids[idx]
+                greenlist_ids = self._get_greenlist_ids(input_ids[:idx])
+                if curr_token in greenlist_ids:
+                    green_token_count += 1
+                    green_token_mask.append(True)
+                else:
+                    green_token_mask.append(False)
+
+        score_dict = dict()
+        if return_num_tokens_scored:
+            score_dict.update(dict(num_tokens_scored=num_tokens_scored))
+        if return_num_green_tokens:
+            score_dict.update(dict(num_green_tokens=green_token_count))
+        if return_z_score:
+            score_dict.update(dict(z_score=self._compute_z_score(green_token_count, num_tokens_scored)))
+        if return_p_value:
+            z_score = score_dict.get("z_score")
+            if z_score is None:
+                z_score = self._compute_z_score(green_token_count, num_tokens_scored)
+            score_dict.update(dict(p_value=self._compute_p_value(z_score)))
+        if return_green_fraction:
+            score_dict.update(dict(green_fraction=(green_token_count / num_tokens_scored)))
+        if return_green_token_mask:
+            score_dict.update(dict(green_token_mask=green_token_mask))
+
+        return score_dict
+
+    def detect(
+        self,
+        text: str = None,
+        tokenized_text: list[int] = None,
+        return_prediction: bool = True,
+        return_scores: bool = True,
+        z_threshold: float = None,
+        **kwargs,
+    ) -> dict:
+
+        assert (text is not None) ^ (tokenized_text is not None), "Must pass either the raw or tokenized string"
+        if return_prediction:
+            kwargs["return_p_value"] = True  # to return the "confidence":=1-p of positive detections
+
+        # run optional normalizers on text
+        for normalizer in self.normalizers:
+            text = normalizer(text)
+        if len(self.normalizers) > 0: 
+            print(f"Text after normalization:\n\n{text}\n")
+
+        if tokenized_text is None:
+            assert self.tokenizer is not None, (
+                "Watermark detection on raw string ",
+                "requires an instance of the tokenizer ",
+                "that was used at generation time.",
+            )
+            tokenized_text = self.tokenizer(text, return_tensors="pt", add_special_tokens=False)["input_ids"][0].to(self.device)
+            if tokenized_text[0] == self.tokenizer.bos_token_id:
+                tokenized_text = tokenized_text[1:]
+        else:
+            # try to remove the bos_tok at beginning if it's there
+            if (self.tokenizer is not None) and (tokenized_text[0] == self.tokenizer.bos_token_id):
+                tokenized_text = tokenized_text[1:]
+
+        # call score method
+        output_dict = {}
+        score_dict = self._score_sequence(tokenized_text, **kwargs)
+        if return_scores:
+            output_dict.update(score_dict)
+        # if passed return_prediction then perform the hypothesis test and return the outcome
+        if return_prediction:
+            z_threshold = z_threshold if z_threshold else self.z_threshold
+            assert z_threshold is not None, "Need a threshold in order to decide outcome of detection test"
+            output_dict["prediction"] = score_dict["z_score"] > z_threshold
+            if output_dict["prediction"]:
+                output_dict["confidence"] = 1 - score_dict["p_value"]
+
+        return output_dict
+