Upload 6 files

app.py

@@ -0,0 +1,52 @@
+# 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, 
+    'demo_public': False,
+    # 'model_name_or_path': 'facebook/opt-125m', 
+    # 'model_name_or_path': 'facebook/opt-1.3b', 
+    # 'model_name_or_path': 'facebook/opt-2.7b', 
+    'model_name_or_path': 'facebook/opt-6.7b',
+    # 'model_name_or_path': 'facebook/opt-13b',
+    # 'load_fp16' : True,
+    'load_fp16' : False,
+    'prompt_max_length': None, 
+    'max_new_tokens': 200, 
+    'generation_seed': 123, 
+    'use_sampling': True, 
+    'n_beams': 1, 
+    'sampling_temp': 0.7, 
+    'use_gpu': True, 
+    'seeding_scheme': 'simple_1', 
+    'gamma': 0.25, 
+    'delta': 2.0, 
+    'normalizers': '', 
+    'ignore_repeated_bigrams': False, 
+    'detection_z_threshold': 4.0, 
+    'select_green_tokens': True,
+    'skip_model_load': True,
+    'seed_separately': True,
+}
+
+args.__dict__.update(arg_dict)
+
+from demo_watermark import main
+
+main(args)

demo_watermark.py

@@ -0,0 +1,1772 @@
+# # 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 os
+# import argparse
+# from argparse import Namespace
+# from pprint import pprint
+# from functools import partial
+#
+# import numpy # for gradio hot reload
+# import gradio as gr
+#
+# import torch
+#
+# from transformers import (AutoTokenizer,
+#                           AutoModelForSeq2SeqLM,
+#                           AutoModelForCausalLM,
+#                           LogitsProcessorList)
+#
+# from watermark_processor import WatermarkLogitsProcessor, WatermarkDetector
+#
+# from typing import Iterable
+# from gradio.themes.base import Base
+# from gradio.themes.utils import colors, fonts, sizes
+# import time
+#
+# def str2bool(v):
+#     """Util function for user friendly boolean flag args"""
+#     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():
+#     """Command line argument specification"""
+#
+#     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=True,
+#         help="Whether to launch as a gradio demo. Set to False if not installed and want to just run the stdout version.",
+#     )
+#     parser.add_argument(
+#         "--demo_public",
+#         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(
+#         "--n_beams",
+#         type=int,
+#         default=1,
+#         help="Number of beams to use for beam search. 1 is normal greedy decoding",
+#     )
+#     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="simple_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.",
+#     )
+#     parser.add_argument(
+#         "--skip_model_load",
+#         type=str2bool,
+#         default=False,
+#         help="Skip the model loading to debug the interface.",
+#     )
+#     parser.add_argument(
+#         "--seed_separately",
+#         type=str2bool,
+#         default=True,
+#         help="Whether to call the torch seed function before both the unwatermarked and watermarked generate calls.",
+#     )
+#     parser.add_argument(
+#         "--load_fp16",
+#         type=str2bool,
+#         default=False,
+#         help="Whether to run model in float16 precsion.",
+#     )
+#     args = parser.parse_args()
+#     return args
+#
+# def load_model(args):
+#     """Load and return the model and tokenizer"""
+#
+#     args.is_seq2seq_model = any([(model_type in args.model_name_or_path) for model_type in ["t5","T0"]])
+#     args.is_decoder_only_model = any([(model_type in args.model_name_or_path) for model_type in ["gpt","opt","bloom"]])
+#     if args.is_seq2seq_model:
+#         model = AutoModelForSeq2SeqLM.from_pretrained(args.model_name_or_path)
+#     elif args.is_decoder_only_model:
+#         if args.load_fp16:
+#             model = AutoModelForCausalLM.from_pretrained(args.model_name_or_path,torch_dtype=torch.float16, device_map='auto')
+#         else:
+#             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"
+#         if args.load_fp16:
+#             pass
+#         else:
+#             model = model.to(device)
+#     else:
+#         device = "cpu"
+#     model.eval()
+#
+#     tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path)
+#
+#     return model, tokenizer, device
+#
+# def generate(prompt, args, model=None, device=None, tokenizer=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}")
+#
+#     watermark_processor = WatermarkLogitsProcessor(vocab=list(tokenizer.get_vocab().values()),
+#                                                     gamma=args.gamma,
+#                                                     delta=args.delta,
+#                                                     seeding_scheme=args.seeding_scheme,
+#                                                     select_green_tokens=args.select_green_tokens)
+#
+#     gen_kwargs = dict(max_new_tokens=args.max_new_tokens)
+#
+#     if args.use_sampling:
+#         gen_kwargs.update(dict(
+#             do_sample=True,
+#             top_k=0,
+#             temperature=args.sampling_temp
+#         ))
+#     else:
+#         gen_kwargs.update(dict(
+#             num_beams=args.n_beams
+#         ))
+#
+#     generate_without_watermark = partial(
+#         model.generate,
+#         **gen_kwargs
+#     )
+#     generate_with_watermark = partial(
+#         model.generate,
+#         logits_processor=LogitsProcessorList([watermark_processor]),
+#         **gen_kwargs
+#     )
+#     if args.prompt_max_length:
+#         pass
+#     elif hasattr(model.config,"max_position_embedding"):
+#         args.prompt_max_length = model.config.max_position_embeddings-args.max_new_tokens
+#     else:
+#         args.prompt_max_length = 2048-args.max_new_tokens
+#
+#     tokd_input = tokenizer(prompt, return_tensors="pt", add_special_tokens=True, truncation=True, max_length=args.prompt_max_length).to(device)
+#     truncation_warning = True if tokd_input["input_ids"].shape[-1] == args.prompt_max_length else False
+#     redecoded_input = tokenizer.batch_decode(tokd_input["input_ids"], skip_special_tokens=True)[0]
+#
+#     torch.manual_seed(args.generation_seed)
+#     output_without_watermark = generate_without_watermark(**tokd_input)
+#
+#     # optional to seed before second generation, but will not be the same again generally, unless delta==0.0, no-op watermark
+#     if args.seed_separately:
+#         torch.manual_seed(args.generation_seed)
+#     output_with_watermark = generate_with_watermark(**tokd_input)
+#
+#     if args.is_decoder_only_model:
+#         # need to isolate the newly generated tokens
+#         output_without_watermark = output_without_watermark[:,tokd_input["input_ids"].shape[-1]:]
+#         output_with_watermark = output_with_watermark[:,tokd_input["input_ids"].shape[-1]:]
+#
+#     decoded_output_without_watermark = tokenizer.batch_decode(output_without_watermark, skip_special_tokens=True)[0]
+#     decoded_output_with_watermark = tokenizer.batch_decode(output_with_watermark, skip_special_tokens=True)[0]
+#
+#     return (redecoded_input,
+#             int(truncation_warning),
+#             decoded_output_without_watermark,
+#             decoded_output_with_watermark,
+#             args)
+#             # decoded_output_with_watermark)
+#
+# def format_names(s):
+#     """Format names for the gradio demo interface"""
+#     s=s.replace("num_tokens_scored","Tokens Counted (T)")
+#     s=s.replace("num_green_tokens","# Tokens in Greenlist")
+#     s=s.replace("green_fraction","Fraction of T in Greenlist")
+#     s=s.replace("z_score","z-score")
+#     s=s.replace("p_value","p value")
+#     s=s.replace("prediction","Prediction")
+#     s=s.replace("confidence","Confidence")
+#     return 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%}"])
+#         elif k=='confidence':
+#             lst_2d.append([format_names(k), f"{v:.3%}"])
+#         elif isinstance(v, float):
+#             lst_2d.append([format_names(k), f"{v:.3g}"])
+#         elif isinstance(v, bool):
+#             lst_2d.append([format_names(k), ("Watermarked" if v else "Human/Unwatermarked")])
+#         else:
+#             lst_2d.append([format_names(k), f"{v}"])
+#     if "confidence" in score_dict:
+#         lst_2d.insert(-2,["z-score Threshold", f"{detection_threshold}"])
+#     else:
+#         lst_2d.insert(-1,["z-score Threshold", f"{detection_threshold}"])
+#     return lst_2d
+#
+# def detect(input_text, args, device=None, tokenizer=None):
+#     """Instantiate the WatermarkDetection object and call detect on
+#         the input text returning the scores and outcome of the test"""
+#     watermark_detector = WatermarkDetector(vocab=list(tokenizer.get_vocab().values()),
+#                                         gamma=args.gamma,
+#                                         seeding_scheme=args.seeding_scheme,
+#                                         device=device,
+#                                         tokenizer=tokenizer,
+#                                         z_threshold=args.detection_z_threshold,
+#                                         normalizers=args.normalizers,
+#                                         ignore_repeated_bigrams=args.ignore_repeated_bigrams,
+#                                         select_green_tokens=args.select_green_tokens)
+#     if len(input_text)-1 > watermark_detector.min_prefix_len:
+#         score_dict = watermark_detector.detect(input_text)
+#         # output = str_format_scores(score_dict, watermark_detector.z_threshold)
+#         output = list_format_scores(score_dict, watermark_detector.z_threshold)
+#     else:
+#         # output = (f"Error: string not long enough to compute watermark presence.")
+#         output = [["Error","string too short to compute metrics"]]
+#         output += [["",""] for _ in range(6)]
+#     return output, args
+#
+# class Seafoam(Base):
+#     def __init__(
+#         self,
+#         *,
+#         primary_hue: colors.Color | str = colors.emerald,
+#         secondary_hue: colors.Color | str = colors.blue,
+#         neutral_hue: colors.Color | str = colors.blue,
+#         spacing_size: sizes.Size | str = sizes.spacing_md,
+#         radius_size: sizes.Size | str = sizes.radius_md,
+#         text_size: sizes.Size | str = sizes.text_lg,
+#         font: fonts.Font
+#         | str
+#         | Iterable[fonts.Font | str] = (
+#             fonts.GoogleFont("Quicksand"),
+#             "ui-sans-serif",
+#             "sans-serif",
+#         ),
+#         font_mono: fonts.Font
+#         | str
+#         | Iterable[fonts.Font | str] = (
+#             fonts.GoogleFont("IBM Plex Mono"),
+#             "ui-monospace",
+#             "monospace",
+#         ),
+#     ):
+#         super().__init__(
+#             primary_hue=primary_hue,
+#             secondary_hue=secondary_hue,
+#             neutral_hue=neutral_hue,
+#             spacing_size=spacing_size,
+#             radius_size=radius_size,
+#             text_size=text_size,
+#             font=font,
+#             font_mono=font_mono,
+#         )
+#         super().set(
+#             body_background_fill="repeating-linear-gradient(45deg, *primary_200, *primary_200 10px, *primary_50 10px, *primary_50 20px)",
+#             body_background_fill_dark="repeating-linear-gradient(45deg, *primary_800, *primary_800 10px, *primary_900 10px, *primary_900 20px)",
+#             button_primary_background_fill="linear-gradient(90deg, *primary_300, *secondary_400)",
+#             button_primary_background_fill_hover="linear-gradient(90deg, *primary_200, *secondary_300)",
+#             button_primary_text_color="white",
+#             button_primary_background_fill_dark="linear-gradient(90deg, *primary_600, *secondary_800)",
+#             slider_color="*secondary_300",
+#             slider_color_dark="*secondary_600",
+#             block_title_text_weight="600",
+#             block_border_width="3px",
+#             block_shadow="*shadow_drop_lg",
+#             button_shadow="*shadow_drop_lg",
+#             button_large_padding="32px",
+#         )
+#
+# seafoam = Seafoam()
+#
+# 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)
+#
+#     # with gr.Blocks(theme="shivi/calm_seafoam") as demo:
+#     # with gr.Blocks(theme="finlaymacklon/smooth_slate") as demo:
+#     # with gr.Blocks(theme="freddyaboulton/test-blue") as demo:
+#     with gr.Blocks(theme="xiaobaiyuan/theme_brief") as demo:
+#         gr.Markdown(
+#             """
+#             # 💧 大语言模型水印 🔍
+#             """
+#         )
+#
+#         with gr.Accordion("参数说明", open=False):
+#             gr.Markdown(
+#                 """
+#                 - `z分数阈值` : 假设检验的截断值。
+#                 - `标记个数 (T)` : 检测算法计算的输出中计数的标记数。
+#                     在简单的单个标记种子方案中，第一个标记被省略，因为它没有前缀标记，无法为其生成绿色列表。
+#                     在底部面板中描述的“忽略重复二元组”检测算法下，如果存在大量重复，这个数量可能远小于生成的总标记数。
+#                 - `绿色列表中的标记数目` : 观察到的落在各自绿色列表中的标记数。
+#                 - `T中含有绿色列表标记的比例` :  `绿色列表中的标记数目` / `T`。预期对于人类/非水印文本，这个比例大约等于 gamma。
+#                 - `z分数` : 检测假设检验的检验统计量。如果大于 `z分数阈值`，则“拒绝零假设”，即文本是人类/非水印的，推断它是带有水印的。
+#                 - `p值` : 在零假设下观察到计算的 `z-分数` 的概率。
+#                     这是在不知道水印程序/绿色列表的情况下观察到 'T中含有绿色列表标记的比例' 的概率。
+#                     如果这个值非常小，我们有信心认为这么多绿色标记不是随机选择的。
+#                 -  `预测` : 假设检验的结果，即观察到的 `z分数` 是否高于 `z分数阈值`。
+#                 - `置信度` : 如果我们拒绝零假设，并且 `预测` 是“Watermarked”，那么我们报告 1-`p 值` 来表示基于这个 `z分数` 观察的检测置信度的不可能性。
+#                 """
+#             )
+#
+#         with gr.Accordion("关于模型能力的说明", open=True):
+#             gr.Markdown(
+#                 """
+#                 本演示使用适用于单个 GPU 的开源语言模型。这些模型比专有商业工具（如 ChatGPT、Claude 或 Bard）的能力更弱。
+#
+#                 还有一件事，我们使用语言模型旨在“完成”您的提示，而不是经过微调以遵循指令的模型。
+#                 为了获得最佳结果，请使用一些组成段落开头的句子提示模型，然后让它“继续”您的段落。
+#                 一些示例包括维基百科文章的开头段落或故事的前几句话。
+#                 结尾处中断的较长提示将产生更流畅的生成。
+#                 """
+#             )
+#
+#         gr.Markdown(f"语言模型: {args.model_name_or_path} {'(float16 mode)' if args.load_fp16 else ''}")
+#
+#         # Construct state for parameters, define updates and toggles
+#         default_prompt = args.__dict__.pop("default_prompt")
+#         session_args = gr.State(value=args)
+#
+#         with gr.Tab("生成检测"):
+#             with gr.Row():
+#                 prompt = gr.Textbox(label=f"提示词", interactive=True,lines=10,max_lines=10, value=default_prompt)
+#             with gr.Row():
+#                 generate_btn = gr.Button("生成")
+#             with gr.Row():
+#                 with gr.Column(scale=2):
+#                     with gr.Tab("未嵌入水印输出的文本"):
+#                         output_without_watermark = gr.Textbox(label=None, interactive=False, lines=14,
+#                                                                   max_lines=14, show_label=False)
+#                     with gr.Tab("高亮"):
+#                         highlight_output_without_watermark = gr.Textbox(label=None, interactive=False, lines=14,
+#                                                                   max_lines=14, show_label=False)
+#                 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=["参数", "值"], interactive=False,
+#                                                                       row_count=7, col_count=2)
+#
+#
+#             with gr.Row():
+#                 with gr.Column(scale=2):
+#                     with gr.Tab("嵌入了水印输出的文本"):
+#                         output_with_watermark = gr.Textbox(label=None, interactive=False, lines=14,
+#                                                            max_lines=14, show_label=False)
+#                     with gr.Tab("高亮"):
+#                         highlight_output_with_watermark = gr.Textbox(label=None, interactive=False, lines=14,
+#                                                            max_lines=14, show_label=False)
+#                 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=["参数", "值"], interactive=False,
+#                                                                    row_count=7, col_count=2)
+#
+#
+#             redecoded_input = gr.Textbox(visible=False)
+#             truncation_warning = gr.Number(visible=False)
+#             def truncate_prompt(redecoded_input, truncation_warning, orig_prompt, args):
+#                 if truncation_warning:
+#                     return redecoded_input + f"\n\n[Prompt was truncated before generation due to length...]", args
+#                 else:
+#                     return orig_prompt, args
+#
+#         with gr.Tab("仅检测"):
+#             with gr.Row():
+#                 with gr.Column(scale=2):
+#                     detection_input = gr.Textbox(label="待分析文本", 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=["参数", "值"], interactive=False, row_count=7, col_count=2)
+#             with gr.Row():
+#                 detect_btn = gr.Button("检测")
+#
+#         # Parameter selection group
+#         with gr.Accordion("高级设置", open=False):
+#             with gr.Row():
+#                 with gr.Column(scale=1):
+#                     gr.Markdown(f"#### 生成参数")
+#                     with gr.Row():
+#                         decoding = gr.Radio(label="解码方式", choices=["multinomial", "greedy"],
+#                                             value=("multinomial" if args.use_sampling else "greedy"))
+#
+#                     with gr.Row():
+#                         sampling_temp = gr.Slider(label="采样随机性多样性权重", minimum=0.1, maximum=1.0, step=0.1,
+#                                                   value=args.sampling_temp, visible=True)
+#                     with gr.Row():
+#                         generation_seed = gr.Number(label="生成种子", value=args.generation_seed, interactive=True)
+#                     with gr.Row():
+#                         n_beams = gr.Dropdown(label="束搜索路数", choices=list(range(1, 11, 1)), value=args.n_beams,
+#                                               visible=(not args.use_sampling))
+#                     with gr.Row():
+#                         max_new_tokens = gr.Slider(label="生成最大标记数", minimum=10, maximum=1000, step=10,
+#                                                    value=args.max_new_tokens)
+#
+#                 with gr.Column(scale=1):
+#                     gr.Markdown(f"#### 水印参数")
+#                     with gr.Row():
+#                         gamma = gr.Slider(label="gamma", minimum=0.1, maximum=0.9, step=0.05, value=args.gamma)
+#                     with gr.Row():
+#                         delta = gr.Slider(label="delta", minimum=0.0, maximum=10.0, step=0.1, value=args.delta)
+#                     gr.Markdown(f"#### 检测参数")
+#                     with gr.Row():
+#                         detection_z_threshold = gr.Slider(label="z-score 阈值", minimum=0.0, maximum=10.0, step=0.1,
+#                                                           value=args.detection_z_threshold)
+#                     with gr.Row():
+#                         ignore_repeated_bigrams = gr.Checkbox(label="忽略重复 Bigram")
+#                     with gr.Row():
+#                         normalizers = gr.CheckboxGroup(label="正则化器",
+#                                                        choices=["unicode", "homoglyphs", "truecase"],
+#                                                        value=args.normalizers)
+#             # with gr.Accordion("Actual submitted parameters:",open=False):
+#             with gr.Row():
+#                 gr.Markdown(
+#                     f"_提示: 滑块更新有延迟。点击滑动条或使用右侧的数字窗口可以帮助更新。下方窗口显示当前的设置。_")
+#             with gr.Row():
+#                 current_parameters = gr.Textbox(label="当前参数", value=args, interactive=False, lines=6)
+#             with gr.Accordion("保留设置", open=False):
+#                 with gr.Row():
+#                     with gr.Column(scale=1):
+#                         seed_separately = gr.Checkbox(label="红绿分别生成", value=args.seed_separately)
+#                     with gr.Column(scale=1):
+#                         select_green_tokens = gr.Checkbox(label="从分区中选择'greenlist'",
+#                                                           value=args.select_green_tokens)
+#
+#         with gr.Accordion("关于设置", open=False):
+#             gr.Markdown(
+#                 """
+#                 #### 生成参数:
+#
+#                 - 解码方法：我们可以使用多项式采样或贪婪解码来从模型中生成标记。
+#                 - 采样温度：如果使用多项式采样，可以设置采样分布的温度。
+#                               0.0 相当于贪婪解码，而 1.0 是下一个标记分布中的最大变异性/熵。
+#                               0.7 在保持对模型对前几个候选者的估计准确性的同时增加了多样性。对于贪婪解码无效。
+#                 - 生成种子：在运行生成之前传递给 torch 随机数生成器的整数。使多项式采样策略输出可复现。对于贪婪解码无效。
+#                 - 并行数：当使用贪婪解码时，还可以将并行数设置为 > 1 以启用波束搜索。
+#                             这在多项式采样中未实现/排除在论文中，但可能会在未来添加。
+#                 - 最大生成标记数：传递给生成方法的 `max_new_tokens` 参数，以在特定数量的新标记处停止输出。
+#                                     请注意，根据提示，模型可以生成较少的标记。
+#                                     这将隐含地将可能的提示标记数量设置为模型的最大输入长度减去 `max_new_tokens`，
+#                                     并且输入将相应地被截断。
+#
+#                 #### 水印参数:
+#
+#                 - gamma：每次生成步骤将词汇表分成绿色列表的部分。较小的 gamma 值通过使得有水印的模型能够更好地与人类/无水印文本区分，
+#                           从而创建了更强的水印，因为它会更倾向于从较小的绿色集合中进行采样，使得这些标记不太可能是偶然发生的。
+#                 - delta：在每个生成步骤中，在采样/选择下一个标记之前，为绿色列表中的每个标记的对数概率添加正偏差。
+#                           较高的 delta 值意味着绿色列表标记更受有水印的模型青睐，并且随着偏差的增大，水印从“软性”过渡到“硬性”。
+#                           对于硬性水印，几乎所有的标记都是绿色的，但这可能对生成质量产生不利影响，特别是当分布的灵活性有限时。
+#
+#                 #### 检测器参数:
+#
+#                 - z-score 阈值：假设检验的 z-score 截断值。较高的阈值（例如 4.0）使得预测人类/无水印文本是有水印的
+#                                   （_false positives_）的可能性非常低，因为一个真正的包含大量标记的人类文本几乎不可能达到那么高的 z-score。
+#                                   较低的阈值将捕捉更多的真正有水印的文本，因为一些有水印的文本可能包含较少的绿色标记并获得较低的 z-score，
+#                                   但仍然通过较低的门槛被标记为“有水印”。然而，较低的阈值会增加被错误地标记为有水印的具有略高于平均绿色标记数的人类文本的几率。
+#                                   4.0-5.0 提供了极低的误报率，同时仍然准确地捕捉到大多数有水印的文本。
+#                 - 忽略重复的双字母组合：此备用检测算法在检测期间只考虑文本中的唯一双字母组合，
+#                                           根据每对中的第一个计算绿色列表，并检查第二个是否在列表内。
+#                                           这意味着 `T` 现在是文本中唯一的双字母组合的数量，
+#                                           如果文本包含大量重复，那么它将少于生成的总标记数。
+#                                           有关更详细的讨论，请参阅论文。
+#                 - 标准化：我们实现了一些基本的标准化，以防止文本在检测过程中受到各种对抗性扰动。
+#                           目前，我们支持将所有字符转换为 Unicode，使用规范形式替换同形字符，并标准化大小写。
+#                 """
+#             )
+#
+#         # gr.HTML("""
+#         #         <p>For faster inference without waiting in queue, you may duplicate the space and upgrade to GPU in settings.
+#         #             Follow the github link at the top and host the demo on your own GPU hardware to test out larger models.
+#         #         <br/>
+#         #         <a href="https://huggingface.co/spaces/tomg-group-umd/lm-watermarking?duplicate=true">
+#         #         <img style="margin-top: 0em; margin-bottom: 0em" src="https://bit.ly/3gLdBN6" alt="Duplicate Space"></a>
+#         #         <p/>
+#         #         """)
+#
+#         # 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], outputs=[redecoded_input, truncation_warning, output_without_watermark, output_with_watermark,session_args])
+#         # Show truncated version of prompt if truncation occurred
+#         redecoded_input.change(fn=truncate_prompt, inputs=[redecoded_input,truncation_warning,prompt,session_args], outputs=[prompt,session_args])
+#         # Call detection when the outputs (of the generate function) are updated
+#         output_without_watermark.change(fn=detect_partial, inputs=[output_without_watermark,session_args], outputs=[without_watermark_detection_result,session_args])
+#         output_with_watermark.change(fn=detect_partial, inputs=[output_with_watermark,session_args], outputs=[with_watermark_detection_result,session_args])
+#         # Register main detection tab click
+#         detect_btn.click(fn=detect_partial, inputs=[detection_input,session_args], outputs=[detection_result, session_args])
+#
+#         # State management logic
+#         # update callbacks that change the state dict
+#         def update_sampling_temp(session_state, value): session_state.sampling_temp = float(value); return session_state
+#         def update_generation_seed(session_state, value): session_state.generation_seed = int(value); return session_state
+#         def update_gamma(session_state, value): session_state.gamma = float(value); return session_state
+#         def update_delta(session_state, value): session_state.delta = float(value); return session_state
+#         def update_detection_z_threshold(session_state, value): session_state.detection_z_threshold = float(value); return session_state
+#         def update_decoding(session_state, value):
+#             if value == "multinomial":
+#                 session_state.use_sampling = True
+#             elif value == "greedy":
+#                 session_state.use_sampling = False
+#             return session_state
+#         def toggle_sampling_vis(value):
+#             if value == "multinomial":
+#                 return gr.update(visible=True)
+#             elif value == "greedy":
+#                 return gr.update(visible=False)
+#         def toggle_sampling_vis_inv(value):
+#             if value == "multinomial":
+#                 return gr.update(visible=False)
+#             elif value == "greedy":
+#                 return gr.update(visible=True)
+#         def update_n_beams(session_state, value): session_state.n_beams = int(value); return session_state
+#         def update_max_new_tokens(session_state, value): session_state.max_new_tokens = int(value); return session_state
+#         def update_ignore_repeated_bigrams(session_state, value): session_state.ignore_repeated_bigrams = value; return session_state
+#         def update_normalizers(session_state, value): session_state.normalizers = value; return session_state
+#         def update_seed_separately(session_state, value): session_state.seed_separately = value; return session_state
+#         def update_select_green_tokens(session_state, value): session_state.select_green_tokens = value; return session_state
+#         # registering callbacks for toggling the visibilty of certain parameters
+#         decoding.change(toggle_sampling_vis,inputs=[decoding], outputs=[sampling_temp])
+#         decoding.change(toggle_sampling_vis,inputs=[decoding], outputs=[generation_seed])
+#         decoding.change(toggle_sampling_vis_inv,inputs=[decoding], outputs=[n_beams])
+#         # registering all state update callbacks
+#         decoding.change(update_decoding,inputs=[session_args, decoding], outputs=[session_args])
+#         sampling_temp.change(update_sampling_temp,inputs=[session_args, sampling_temp], outputs=[session_args])
+#         generation_seed.change(update_generation_seed,inputs=[session_args, generation_seed], outputs=[session_args])
+#         n_beams.change(update_n_beams,inputs=[session_args, n_beams], outputs=[session_args])
+#         max_new_tokens.change(update_max_new_tokens,inputs=[session_args, max_new_tokens], outputs=[session_args])
+#         gamma.change(update_gamma,inputs=[session_args, gamma], outputs=[session_args])
+#         delta.change(update_delta,inputs=[session_args, delta], outputs=[session_args])
+#         detection_z_threshold.change(update_detection_z_threshold,inputs=[session_args, detection_z_threshold], outputs=[session_args])
+#         ignore_repeated_bigrams.change(update_ignore_repeated_bigrams,inputs=[session_args, ignore_repeated_bigrams], outputs=[session_args])
+#         normalizers.change(update_normalizers,inputs=[session_args, normalizers], outputs=[session_args])
+#         seed_separately.change(update_seed_separately,inputs=[session_args, seed_separately], outputs=[session_args])
+#         select_green_tokens.change(update_select_green_tokens,inputs=[session_args, select_green_tokens], outputs=[session_args])
+#         # register additional callback on button clicks that updates the shown parameters window
+#         generate_btn.click(lambda value: str(value), inputs=[session_args], outputs=[current_parameters])
+#         detect_btn.click(lambda value: str(value), inputs=[session_args], outputs=[current_parameters])
+#         # When the parameters change, display the update and fire detection, since some detection params dont change the model output.
+#         gamma.change(lambda value: str(value), inputs=[session_args], outputs=[current_parameters])
+#         gamma.change(fn=detect_partial, inputs=[output_without_watermark,session_args], outputs=[without_watermark_detection_result,session_args])
+#         gamma.change(fn=detect_partial, inputs=[output_with_watermark,session_args], outputs=[with_watermark_detection_result,session_args])
+#         gamma.change(fn=detect_partial, inputs=[detection_input,session_args], outputs=[detection_result,session_args])
+#         detection_z_threshold.change(lambda value: str(value), inputs=[session_args], outputs=[current_parameters])
+#         detection_z_threshold.change(fn=detect_partial, inputs=[output_without_watermark,session_args], outputs=[without_watermark_detection_result,session_args])
+#         detection_z_threshold.change(fn=detect_partial, inputs=[output_with_watermark,session_args], outputs=[with_watermark_detection_result,session_args])
+#         detection_z_threshold.change(fn=detect_partial, inputs=[detection_input,session_args], outputs=[detection_result,session_args])
+#         ignore_repeated_bigrams.change(lambda value: str(value), inputs=[session_args], outputs=[current_parameters])
+#         ignore_repeated_bigrams.change(fn=detect_partial, inputs=[output_without_watermark,session_args], outputs=[without_watermark_detection_result,session_args])
+#         ignore_repeated_bigrams.change(fn=detect_partial, inputs=[output_with_watermark,session_args], outputs=[with_watermark_detection_result,session_args])
+#         ignore_repeated_bigrams.change(fn=detect_partial, inputs=[detection_input,session_args], outputs=[detection_result,session_args])
+#         normalizers.change(lambda value: str(value), inputs=[session_args], outputs=[current_parameters])
+#         normalizers.change(fn=detect_partial, inputs=[output_without_watermark,session_args], outputs=[without_watermark_detection_result,session_args])
+#         normalizers.change(fn=detect_partial, inputs=[output_with_watermark,session_args], outputs=[with_watermark_detection_result,session_args])
+#         normalizers.change(fn=detect_partial, inputs=[detection_input,session_args], outputs=[detection_result,session_args])
+#         select_green_tokens.change(lambda value: str(value), inputs=[session_args], outputs=[current_parameters])
+#         select_green_tokens.change(fn=detect_partial, inputs=[output_without_watermark,session_args], outputs=[without_watermark_detection_result,session_args])
+#         select_green_tokens.change(fn=detect_partial, inputs=[output_with_watermark,session_args], outputs=[with_watermark_detection_result,session_args])
+#         select_green_tokens.change(fn=detect_partial, inputs=[detection_input,session_args], outputs=[detection_result,session_args])
+#
+#
+#     demo.queue(concurrency_count=3)
+#
+#     if args.demo_public:
+#         demo.launch(share=True) # exposes app to the internet via randomly generated link
+#     else:
+#         demo.launch()
+#
+# def main(args):
+#     """Run a command line version of the generation and detection operations
+#         and optionally launch and serve the gradio demo"""
+#     # Initial arg processing and log
+#     args.normalizers = (args.normalizers.split(",") if args.normalizers else [])
+#     print(args)
+#
+#     if not args.skip_model_load:
+#         model, tokenizer, device = load_model(args)
+#     else:
+#         model, tokenizer, device = None, None, None
+#
+#     # Generate and detect, report to stdout
+#     if not args.skip_model_load:
+#         input_text = (
+#         "The diamondback terrapin or simply terrapin (Malaclemys terrapin) is a "
+#         "species of turtle native to the brackish coastal tidal marshes of the "
+#         "Northeastern and southern United States, and in Bermuda.[6] It belongs "
+#         "to the monotypic genus Malaclemys. It has one of the largest ranges of "
+#         "all turtles in North America, stretching as far south as the Florida Keys "
+#         "and as far north as Cape Cod.[7] The name 'terrapin' is derived from the "
+#         "Algonquian word torope.[8] It applies to Malaclemys terrapin in both "
+#         "British English and American English. The name originally was used by "
+#         "early European settlers in North America to describe these brackish-water "
+#         "turtles that inhabited neither freshwater habitats nor the sea. It retains "
+#         "this primary meaning in American English.[8] In British English, however, "
+#         "other semi-aquatic turtle species, such as the red-eared slider, might "
+#         "also be called terrapins. The common name refers to the diamond pattern "
+#         "on top of its shell (carapace), but the overall pattern and coloration "
+#         "vary greatly. The shell is usually wider at the back than in the front, "
+#         "and from above it appears wedge-shaped. The shell coloring can vary "
+#         "from brown to grey, and its body color can be grey, brown, yellow, "
+#         "or white. All have a unique pattern of wiggly, black markings or spots "
+#         "on their body and head. The diamondback terrapin has large webbed "
+#         "feet.[9] The species is"
+#         )
+#
+#         args.default_prompt = input_text
+#
+#         term_width = 80
+#         print("#"*term_width)
+#         print("Prompt:")
+#         print(input_text)
+#
+#         _, _, decoded_output_without_watermark, decoded_output_with_watermark, _ = generate(input_text,
+#                                                                                             args,
+#                                                                                             model=model,
+#                                                                                             device=device,
+#                                                                                             tokenizer=tokenizer)
+#         without_watermark_detection_result = detect(decoded_output_without_watermark,
+#                                                     args,
+#                                                     device=device,
+#                                                     tokenizer=tokenizer)
+#         with_watermark_detection_result = detect(decoded_output_with_watermark,
+#                                                  args,
+#                                                  device=device,
+#                                                  tokenizer=tokenizer)
+#
+#         print("#"*term_width)
+#         print("Output without watermark:")
+#         print(decoded_output_without_watermark)
+#         print("-"*term_width)
+#         print(f"Detection result @ {args.detection_z_threshold}:")
+#         pprint(without_watermark_detection_result)
+#         print("-"*term_width)
+#
+#         print("#"*term_width)
+#         print("Output with watermark:")
+#         print(decoded_output_with_watermark)
+#         print("-"*term_width)
+#         print(f"Detection result @ {args.detection_z_threshold}:")
+#         pprint(with_watermark_detection_result)
+#         print("-"*term_width)
+#
+#
+#     # Launch the app to generate and detect interactively (implements the hf space demo)
+#     if args.run_gradio:
+#         run_gradio(args, model=model, tokenizer=tokenizer, device=device)
+#
+#     return
+#
+# if __name__ == "__main__":
+#
+#     args = parse_args()
+#     print(args)
+#
+#     main(args)
+
+# 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.
+
+# 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 numpy  # for gradio hot reload
+import gradio as gr
+
+import torch
+
+from transformers import (AutoTokenizer,
+                          AutoModelForSeq2SeqLM,
+                          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 = {
+    "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},
+    # "bigscience/bloom"        : {"max_length": 1000, "gamma": 0.5, "delta": 2.0},
+    # "bigscience/bloomz"       : {"max_length": 1000, "gamma": 0.5, "delta": 2.0},
+}
+
+
+def str2bool(v):
+    """Util function for user friendly boolean flag args"""
+    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():
+    """Command line argument specification"""
+
+    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=True,
+        help="Whether to launch as a gradio demo. Set to False if not installed and want to just run the stdout version.",
+    )
+    parser.add_argument(
+        "--demo_public",
+        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(
+        "--n_beams",
+        type=int,
+        default=1,
+        help="Number of beams to use for beam search. 1 is normal greedy decoding",
+    )
+    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="simple_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.",
+    )
+    parser.add_argument(
+        "--skip_model_load",
+        type=str2bool,
+        default=False,
+        help="Skip the model loading to debug the interface.",
+    )
+    parser.add_argument(
+        "--seed_separately",
+        type=str2bool,
+        default=True,
+        help="Whether to call the torch seed function before both the unwatermarked and watermarked generate calls.",
+    )
+    parser.add_argument(
+        "--load_fp16",
+        type=str2bool,
+        default=False,
+        help="Whether to run model in float16 precsion.",
+    )
+    args = parser.parse_args()
+    return args
+
+
+def load_model(args):
+    """Load and return the model and tokenizer"""
+
+    args.is_seq2seq_model = any([(model_type in args.model_name_or_path) for model_type in ["t5", "T0"]])
+    args.is_decoder_only_model = any(
+        [(model_type in args.model_name_or_path) for model_type in ["gpt", "opt", "bloom"]])
+    if args.is_seq2seq_model:
+        model = AutoModelForSeq2SeqLM.from_pretrained(args.model_name_or_path)
+    elif args.is_decoder_only_model:
+        if args.load_fp16:
+            model = AutoModelForCausalLM.from_pretrained(args.model_name_or_path, torch_dtype=torch.float16,
+                                                         device_map='auto')
+        else:
+            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"
+        if args.load_fp16:
+            pass
+        else:
+            model = model.to(device)
+    else:
+        device = "cpu"
+    model.eval()
+
+    tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path)
+
+    return model, tokenizer, device
+
+
+from text_generation import InferenceAPIClient
+from requests.exceptions import ReadTimeout
+
+
+def generate_with_api(prompt, args):
+    # hf_api_key = os.environ.get("HF_API_KEY")
+    hf_api_key = "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"
+    if hf_api_key is None:
+        raise ValueError("HF_API_KEY environment variable not set, cannot use HF API to generate text.")
+
+    client = InferenceAPIClient(args.model_name_or_path, token=hf_api_key, timeout=60)
+
+    assert args.n_beams == 1, "HF API models do not support beam search."
+    generation_params = {
+        "max_new_tokens": args.max_new_tokens,
+        "do_sample": args.use_sampling,
+    }
+    if args.use_sampling:
+        generation_params["temperature"] = args.sampling_temp
+        generation_params["seed"] = args.generation_seed
+
+    timeout_msg = "[Model API timeout error. Try reducing the max_new_tokens parameter or the prompt length.]"
+    try:
+        generation_params["watermark"] = False
+        without_watermark_iterator = client.generate_stream(prompt, **generation_params)
+    except ReadTimeout as e:
+        print(e)
+        without_watermark_iterator = (char for char in timeout_msg)
+    try:
+        generation_params["watermark"] = True
+        with_watermark_iterator = client.generate_stream(prompt, **generation_params)
+    except ReadTimeout as e:
+        print(e)
+        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 check_prompt(prompt, args, tokenizer, model=None, device=None):
+    # This applies to both the local and API model scenarios
+    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:
+        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]
+
+    return (redecoded_input,
+            int(truncation_warning),
+            args)
+
+
+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}")
+    print(f"Prompt: {prompt}")
+
+    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
+        )
+
+        torch.manual_seed(args.generation_seed)
+        output_without_watermark = generate_without_watermark(**tokd_input)
+
+        # optional to seed before second generation, but will not be the same again generally, unless delta==0.0, no-op watermark
+        if args.seed_separately:
+            torch.manual_seed(args.generation_seed)
+        output_with_watermark = generate_with_watermark(**tokd_input)
+
+        if args.is_decoder_only_model:
+            # 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]
+
+        # 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):
+    """Format names for the gradio demo interface"""
+    s = s.replace("num_tokens_scored", "Tokens Counted (T)")
+    s = s.replace("num_green_tokens", "# Tokens in Greenlist")
+    s = s.replace("green_fraction", "Fraction of T in Greenlist")
+    s = s.replace("z_score", "z-score")
+    s = s.replace("p_value", "p value")
+    s = s.replace("prediction", "Prediction")
+    s = s.replace("confidence", "Confidence")
+    return s
+
+
+def list_format_scores(score_dict, detection_threshold):
+    """Format the detection metrics into a gradio dataframe input format"""
+    lst_2d = []
+    for k, v in score_dict.items():
+        if k == 'green_fraction':
+            lst_2d.append([format_names(k), f"{v:.1%}"])
+        elif k == 'confidence':
+            lst_2d.append([format_names(k), f"{v:.3%}"])
+        elif isinstance(v, float):
+            lst_2d.append([format_names(k), f"{v:.3g}"])
+        elif isinstance(v, bool):
+            lst_2d.append([format_names(k), ("Watermarked" if v else "Human/Unwatermarked")])
+        else:
+            lst_2d.append([format_names(k), f"{v}"])
+    if "confidence" in score_dict:
+        lst_2d.insert(-2, ["z-score Threshold", f"{detection_threshold}"])
+    else:
+        lst_2d.insert(-1, ["z-score Threshold", f"{detection_threshold}"])
+    return lst_2d
+
+
+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)}")
+
+    watermark_detector = WatermarkDetector(vocab=list(tokenizer.get_vocab().values()),
+                                           gamma=args.gamma,
+                                           seeding_scheme=args.seeding_scheme,
+                                           device=device,
+                                           tokenizer=tokenizer,
+                                           z_threshold=args.detection_z_threshold,
+                                           normalizers=args.normalizers,
+                                           ignore_repeated_bigrams=args.ignore_repeated_bigrams,
+                                           select_green_tokens=args.select_green_tokens)
+    # for now, just don't display the green token mask
+    # if we're using normalizers or ignore_repeated_bigrams
+    if args.normalizers != [] or args.ignore_repeated_bigrams:
+        return_green_token_mask = False
+
+    error = False
+    green_token_mask = None
+    if input_text == "":
+        error = True
+    else:
+        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)
+            error = True
+    if error:
+        output = [["Error", "string too short to compute metrics"]]
+        output += [["", ""] for _ in range(6)]
+
+    html_output = "[No highlight markup generated]"
+    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"""
+    check_prompt_partial = partial(check_prompt, model=model, device=device)
+    generate_partial = partial(generate, model=model, device=device)
+    detect_partial = partial(detect, device=device)
+
+    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(theme="xiaobaiyuan/theme_brief") as demo:
+    with gr.Blocks(css=css, theme="xiaobaiyuan/theme_brief") as demo:
+        # Top section, greeting and instructions
+        with gr.Row():
+            with gr.Column(scale=9):
+                gr.Markdown(
+                    """
+                    # 💧 大语言模型水印 🔍
+                    """
+                )
+            with gr.Column(scale=1):
+                # 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(
+                    all_models,
+                    value=args.model_name_or_path,
+                    label="Language Model",
+                )
+
+        with gr.Accordion("参数说明", open=False):
+            gr.Markdown(
+                """
+                - `z分数阈值` : 假设检验的截断值。
+                - `标记个数 (T)` : 检测算法计算的输出中计数的标记数。
+                    在简单的单个标记种子方案中，第一个标记被省略，因为它没有前缀标记，无法为其生成绿色列表。
+                    在底部面板中描述的“忽略重复二元组”检测算法下，如果存在大量重复，这个数量可能远小于生成的总标记数。
+                - `绿色列表中的标记数目` : 观察到的落在各自绿色列表中的标记数。
+                - `T中含有绿色列表标记的比例` :  `绿色列表中的标记数目` / `T`。预期对于人类/非水印文本，这个比例大约等于 gamma。
+                - `z分数` : 检测假设检验的检验统计量。如果大于 `z分数阈值`，则“拒绝零假设”，即文本是人类/非水印的，推断它是带有水印的。
+                - `p值` : 在零假设下观察到计算的 `z-分数` 的概率。
+                    这是在不知道水印程序/绿色列表的情况下观察到 'T中含有绿色列表标记的比例' 的概率。
+                    如果这个值非常小，我们有信心认为这么多绿色标记不是随机选择的。
+                -  `预测` : 假设检验的结果，即观察到的 `z分数` 是否高于 `z分数阈值`。
+                - `置信度` : 如果我们拒绝零假设，并且 `预测` 是“Watermarked”，那么我们报告 1-`p 值` 来表示基于这个 `z分数` 观察的检测置信度的不可能性。
+                """
+            )
+
+        with gr.Accordion("关于模型能力的说明", open=True):
+            gr.Markdown(
+                """
+                本演示使用适用于单个 GPU 的开源语言模型��这些模型比专有商业工具（如 ChatGPT、Claude 或 Bard）的能力更弱。
+
+                还有一件事，我们使用语言模型旨在“完成”您的提示，而不是经过微调以遵循指令的模型。
+                为了获得最佳结果，请使用一些组成段落开头的句子提示模型，然后让它“继续”您的段落。
+                一些示例包括维基百科文章的开头段落或故事的前几句话。
+                结尾处中断的较长提示将产生更流畅的生成。
+                """
+            )
+
+        # Construct state for parameters, define updates and toggles
+        default_prompt = args.__dict__.pop("default_prompt")
+        session_args = gr.State(value=args)
+        # note that state obj automatically calls value if it's a callable, want to avoid calling tokenizer at startup
+        session_tokenizer = gr.State(value=lambda: tokenizer)
+        # with gr.Row():
+        #     gr.Markdown(
+        #         """
+        #         温馨提示：若出现ERROR，可能由于api暂未成功载入，稍等片刻即可
+        #         """
+        #     )
+        with gr.Tab("生成检测"):
+            with gr.Row():
+                prompt = gr.Textbox(label=f"提示词", interactive=True,lines=10,max_lines=10, value=default_prompt)
+            with gr.Row():
+                generate_btn = gr.Button("生成")
+            with gr.Row():
+                with gr.Column(scale=2):
+                    with gr.Tab("未嵌入水印输出的文本"):
+                        output_without_watermark = gr.Textbox(label=None, interactive=False, lines=14,
+                                                                  max_lines=14, show_label=False)
+                    with gr.Tab("高亮"):
+                        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=["参数", "值"], interactive=False,
+                                                                      row_count=7, col_count=2)
+
+            with gr.Row():
+                with gr.Column(scale=2):
+                    with gr.Tab("嵌入了水印输出的文本"):
+                        output_with_watermark = gr.Textbox(label=None, interactive=False, lines=14,
+                                                           max_lines=14, show_label=False)
+                    with gr.Tab("高亮"):
+                        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=["参数", "值"], interactive=False,
+                                                                   row_count=7, col_count=2)
+
+
+            redecoded_input = gr.Textbox(visible=False)
+            truncation_warning = gr.Number(visible=False)
+            def truncate_prompt(redecoded_input, truncation_warning, orig_prompt, args):
+                if truncation_warning:
+                    return redecoded_input + f"\n\n[由于长度原因，提示词被截断...]", args
+                else:
+                    return orig_prompt, args
+
+        with gr.Tab("仅检测"):
+            with gr.Row():
+                with gr.Column(scale=2):
+                    with gr.Tab("待分析文本"):
+                        detection_input = gr.Textbox(interactive=True, lines=14, max_lines=14,show_label=False)
+                    with gr.Tab("高亮"):
+                        html_detection_input = gr.HTML(elem_id="html-detection-input")
+                with gr.Column(scale=1):
+                    detection_result = gr.Dataframe(headers=["参数", "值"], interactive=False, row_count=7,
+                                                    col_count=2)
+            with gr.Row():
+                detect_btn = gr.Button("检测")
+
+        # Parameter selection group
+        with gr.Accordion("高级设置", open=False):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    gr.Markdown(f"#### 生成参数")
+                    with gr.Row():
+                        decoding = gr.Radio(label="解码方式", choices=["multinomial", "greedy"],
+                                            value=("multinomial" if args.use_sampling else "greedy"))
+
+                    with gr.Row():
+                        sampling_temp = gr.Slider(label="采样随机性多样性权重", minimum=0.1, maximum=1.0, step=0.1,
+                                                  value=args.sampling_temp, visible=True)
+                    with gr.Row():
+                        generation_seed = gr.Number(label="生成种子", value=args.generation_seed, interactive=True)
+                    with gr.Row():
+                        n_beams = gr.Dropdown(label="束搜索路数", choices=list(range(1, 11, 1)), value=args.n_beams,
+                                              visible=(not args.use_sampling))
+                    with gr.Row():
+                        max_new_tokens = gr.Slider(label="生成最大标记数", minimum=10, maximum=1000, step=10,
+                                                   value=args.max_new_tokens)
+
+                with gr.Column(scale=1):
+                    gr.Markdown(f"#### 水印参数")
+                    with gr.Row():
+                        gamma = gr.Slider(label="gamma", minimum=0.1, maximum=0.9, step=0.05, value=args.gamma)
+                    with gr.Row():
+                        delta = gr.Slider(label="delta", minimum=0.0, maximum=10.0, step=0.1, value=args.delta)
+                    gr.Markdown(f"#### 检测参数")
+                    with gr.Row():
+                        detection_z_threshold = gr.Slider(label="z-score 阈值", minimum=0.0, maximum=10.0, step=0.1,
+                                                          value=args.detection_z_threshold)
+                    with gr.Row():
+                        ignore_repeated_bigrams = gr.Checkbox(label="忽略重复 Bigram")
+                    with gr.Row():
+                        normalizers = gr.CheckboxGroup(label="正则化器",
+                                                       choices=["unicode", "homoglyphs", "truecase"],
+                                                       value=args.normalizers)
+            # with gr.Accordion("Actual submitted parameters:",open=False):
+            with gr.Row():
+                gr.Markdown(
+                    f"_提示: 滑块更新有延迟。点击滑动条或使用右侧的数字窗口可以帮助更新。下方窗口显示当前的设置。_")
+            with gr.Row():
+                current_parameters = gr.Textbox(label="当前参数", value=args, interactive=False, lines=6)
+            with gr.Accordion("保留设置", open=False):
+                with gr.Row():
+                    with gr.Column(scale=1):
+                        seed_separately = gr.Checkbox(label="红绿分别生成", value=args.seed_separately)
+                    with gr.Column(scale=1):
+                        select_green_tokens = gr.Checkbox(label="从分区中选择'greenlist'",
+                                                          value=args.select_green_tokens)
+
+        with gr.Accordion("关于设置", open=False):
+            gr.Markdown(
+                """
+                #### 生成参数:
+
+                - 解码方法：我们可以使用多项式采样或贪婪解码来从模型中生成标记。
+                - 采样温度：如果使用多项式采样，可以设置采样分布的温度。
+                              0.0 相当于贪婪解码，而 1.0 是下一个标记分布中的最大变异性/熵。
+                              0.7 在保持对模型对前几个候选者的估计准确性的同时增加了多样性。对于贪婪解码无效。
+                - 生成种子：在运行生成之前传递给 torch 随机数生成器的整数。使多项式采样策略输出可复现。对于贪婪解码无效。
+                - 并行数：当使用贪婪解码时，还可以将并行数设置为 > 1 以启用波束搜索。
+                            这在多项式采样中未实现/排除在论文中，但可能会在未来添加。
+                - 最大生成标记数：传递给生成方法的 `max_new_tokens` 参数，以在特定数量的新标记处停止输出。
+                                    请注意，根据提示，模型可以生成较少的标记。
+                                    这将隐含地将可能的提示标记数量设置为模型的最大输入长度减去 `max_new_tokens`，
+                                    并且输入将相应地被截断。
+
+                #### 水印参数:
+
+                - gamma：每次生成步骤将词汇表分成绿色列表的部分。较小的 gamma 值通过使得有水印的模型能够更好地与人类/无水印文本区分，
+                          从而创建了更强的水印，因为它会更倾向于从较小的绿色集合中进行采样，使得这些标记不太可能是偶然发生的。
+                - delta：在每个生成步骤中，在采样/选择下一个标记之前，为绿色列表中的每个标记的对数概率添加正偏差。
+                          较高的 delta 值意味着绿色列表标记更受有水印的模型青睐，并且随着偏差的增大，水印从“软性”过渡到“硬性”。
+                          对于硬性水印，几乎所有的标记都是绿色的，但这可能对生成质量产生不利影响，特别是当分布的灵活性有限时。
+
+                #### 检测器参数:
+
+                - z-score 阈值：假设检验的 z-score 截断值。较高的阈值（例如 4.0）使得预测人类/无水印文本是有水印的
+                                  （_false positives_）的可能性非常低，因为一个真正的包含大量标记的人类文本几乎不可能达到那么高的 z-score。
+                                  较低的阈值将捕捉更多的真正有水印的文本，因为一些有水印的文本可能包含较少的绿色标记并获得较低的 z-score，
+                                  但仍然通过较低的门槛被标记为“有水印”。然而，较低的阈值会增加被错误地标记为有水印的具有略高于平均绿色标记数的人类文本的几率。
+                                  4.0-5.0 提供了极低的误报率，同时仍然准确地捕捉到大多数有水印的文本。
+                - 忽略重复的双字母组合：此备用检测算法在检测期间只考虑文本中的唯一双字母组合，
+                                          根据每对中的第一个计算绿色列表，并检查第二个是否在列表内。
+                                          这意味着 `T` 现在是文本中唯一的双字母组合的数量，
+                                          如果文本包含大量重复，那么它将少于生成的总标记数。
+                                          有关更详细的讨论，请参阅论文。
+                - 标准化：我们实现了一些基本的标准化，以防止文本在检测过程中受到各种对抗性扰动。
+                          目前，我们支持将所有字符转换为 Unicode，使用规范形式替换同形字符，并标准化大小写。
+                """
+            )
+
+
+        # 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])
+        # 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, html_detection_input],
+                         api_name="detection")
+
+        # State management logic
+        # define update callbacks that change the state dict
+        def update_model(session_state, value):
+            session_state.model_name_or_path = value; return session_state
+
+        def update_sampling_temp(session_state, value):
+            session_state.sampling_temp = float(value); return session_state
+
+        def update_generation_seed(session_state, value):
+            session_state.generation_seed = int(value); return session_state
+
+        def update_gamma(session_state, value):
+            session_state.gamma = float(value); return session_state
+
+        def update_delta(session_state, value):
+            session_state.delta = float(value); return session_state
+
+        def update_detection_z_threshold(session_state, value):
+            session_state.detection_z_threshold = float(value); return session_state
+
+        def update_decoding(session_state, value):
+            if value == "multinomial":
+                session_state.use_sampling = True
+            elif value == "greedy":
+                session_state.use_sampling = False
+            return session_state
+
+        def toggle_sampling_vis(value):
+            if value == "multinomial":
+                return gr.update(visible=True)
+            elif value == "greedy":
+                return gr.update(visible=False)
+
+        def toggle_sampling_vis_inv(value):
+            if value == "multinomial":
+                return gr.update(visible=False)
+            elif value == "greedy":
+                return gr.update(visible=True)
+
+        # if model name is in the list of api models, set the num beams parameter to 1 and hide n_beams
+        def toggle_vis_for_api_model(value):
+            if value in API_MODEL_MAP:
+                return gr.update(visible=False)
+            else:
+                return gr.update(visible=True)
+
+        def toggle_beams_for_api_model(value, orig_n_beams):
+            if value in API_MODEL_MAP:
+                return gr.update(value=1)
+            else:
+                return gr.update(value=orig_n_beams)
+
+        # if model name is in the list of api models, set the interactive parameter to false
+        def toggle_interactive_for_api_model(value):
+            if value in API_MODEL_MAP:
+                return gr.update(interactive=False)
+            else:
+                return gr.update(interactive=True)
+
+        # if model name is in the list of api models, set gamma and delta based on API map
+        def toggle_gamma_for_api_model(value, orig_gamma):
+            if value in API_MODEL_MAP:
+                return gr.update(value=API_MODEL_MAP[value]["gamma"])
+            else:
+                return gr.update(value=orig_gamma)
+
+        def toggle_delta_for_api_model(value, orig_delta):
+            if value in API_MODEL_MAP:
+                return gr.update(value=API_MODEL_MAP[value]["delta"])
+            else:
+                return gr.update(value=orig_delta)
+
+        def update_n_beams(session_state, value):
+            session_state.n_beams = int(value); return session_state
+
+        def update_max_new_tokens(session_state, value):
+            session_state.max_new_tokens = int(value); return session_state
+
+        def update_ignore_repeated_bigrams(session_state, value):
+            session_state.ignore_repeated_bigrams = value; return session_state
+
+        def update_normalizers(session_state, value):
+            session_state.normalizers = value; return session_state
+
+        def update_seed_separately(session_state, value):
+            session_state.seed_separately = value; return session_state
+
+        def update_select_green_tokens(session_state, value):
+            session_state.select_green_tokens = value; return session_state
+
+        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)
+
+        def check_model(value):
+            return value if (value != "" and value is not None) else args.model_name_or_path
+
+        # enforce constraint that model cannot be null or empty
+        # then attach model callbacks in particular
+        model_selector.change(check_model, inputs=[model_selector], outputs=[model_selector]).then(
+            toggle_vis_for_api_model, inputs=[model_selector], outputs=[n_beams]
+        ).then(
+            toggle_beams_for_api_model, inputs=[model_selector, n_beams], outputs=[n_beams]
+        ).then(
+            toggle_interactive_for_api_model, inputs=[model_selector], outputs=[gamma]
+        ).then(
+            toggle_interactive_for_api_model, inputs=[model_selector], outputs=[delta]
+        ).then(
+            toggle_gamma_for_api_model, inputs=[model_selector, gamma], outputs=[gamma]
+        ).then(
+            toggle_delta_for_api_model, inputs=[model_selector, delta], outputs=[delta]
+        ).then(
+            update_tokenizer, inputs=[model_selector], outputs=[session_tokenizer]
+        ).then(
+            update_model, inputs=[session_args, model_selector], outputs=[session_args]
+        ).then(
+            lambda value: str(value), inputs=[session_args], outputs=[current_parameters]
+        )
+        # 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])
+        decoding.change(toggle_sampling_vis_inv, inputs=[decoding], outputs=[n_beams])
+        decoding.change(toggle_vis_for_api_model, inputs=[model_selector], outputs=[n_beams])
+        # registering all state update callbacks
+        decoding.change(update_decoding, inputs=[session_args, decoding], outputs=[session_args])
+        sampling_temp.change(update_sampling_temp, inputs=[session_args, sampling_temp], outputs=[session_args])
+        generation_seed.change(update_generation_seed, inputs=[session_args, generation_seed], outputs=[session_args])
+        n_beams.change(update_n_beams, inputs=[session_args, n_beams], outputs=[session_args])
+        max_new_tokens.change(update_max_new_tokens, inputs=[session_args, max_new_tokens], outputs=[session_args])
+        gamma.change(update_gamma, inputs=[session_args, gamma], outputs=[session_args])
+        delta.change(update_delta, inputs=[session_args, delta], outputs=[session_args])
+        detection_z_threshold.change(update_detection_z_threshold, inputs=[session_args, detection_z_threshold],
+                                     outputs=[session_args])
+        ignore_repeated_bigrams.change(update_ignore_repeated_bigrams, inputs=[session_args, ignore_repeated_bigrams],
+                                       outputs=[session_args])
+        normalizers.change(update_normalizers, inputs=[session_args, normalizers], outputs=[session_args])
+        seed_separately.change(update_seed_separately, inputs=[session_args, seed_separately], outputs=[session_args])
+        select_green_tokens.change(update_select_green_tokens, inputs=[session_args, select_green_tokens],
+                                   outputs=[session_args])
+        # register additional callback on button clicks that updates the shown parameters window
+        generate_btn.click(lambda value: str(value), inputs=[session_args], outputs=[current_parameters])
+        detect_btn.click(lambda value: str(value), inputs=[session_args], outputs=[current_parameters])
+        # When the parameters change, display the update and also fire detection, since some detection params dont change the model output.
+        delta.change(lambda value: str(value), inputs=[session_args], outputs=[current_parameters])
+        gamma.change(lambda value: str(value), inputs=[session_args], outputs=[current_parameters])
+        gamma.change(fn=detect_partial, inputs=[output_without_watermark, session_args, session_tokenizer],
+                     outputs=[without_watermark_detection_result, session_args, session_tokenizer,
+                              html_without_watermark])
+        gamma.change(fn=detect_partial, inputs=[output_with_watermark, session_args, session_tokenizer],
+                     outputs=[with_watermark_detection_result, session_args, session_tokenizer, html_with_watermark])
+        gamma.change(fn=detect_partial, inputs=[detection_input, session_args, session_tokenizer],
+                     outputs=[detection_result, session_args, session_tokenizer, html_detection_input])
+        detection_z_threshold.change(lambda value: str(value), inputs=[session_args], outputs=[current_parameters])
+        detection_z_threshold.change(fn=detect_partial,
+                                     inputs=[output_without_watermark, session_args, session_tokenizer],
+                                     outputs=[without_watermark_detection_result, session_args, session_tokenizer,
+                                              html_without_watermark])
+        detection_z_threshold.change(fn=detect_partial, inputs=[output_with_watermark, session_args, session_tokenizer],
+                                     outputs=[with_watermark_detection_result, session_args, session_tokenizer,
+                                              html_with_watermark])
+        detection_z_threshold.change(fn=detect_partial, inputs=[detection_input, session_args, session_tokenizer],
+                                     outputs=[detection_result, session_args, session_tokenizer, html_detection_input])
+        ignore_repeated_bigrams.change(lambda value: str(value), inputs=[session_args], outputs=[current_parameters])
+        ignore_repeated_bigrams.change(fn=detect_partial,
+                                       inputs=[output_without_watermark, session_args, session_tokenizer],
+                                       outputs=[without_watermark_detection_result, session_args, session_tokenizer,
+                                                html_without_watermark])
+        ignore_repeated_bigrams.change(fn=detect_partial,
+                                       inputs=[output_with_watermark, session_args, session_tokenizer],
+                                       outputs=[with_watermark_detection_result, session_args, session_tokenizer,
+                                                html_with_watermark])
+        ignore_repeated_bigrams.change(fn=detect_partial, inputs=[detection_input, session_args, session_tokenizer],
+                                       outputs=[detection_result, session_args, session_tokenizer,
+                                                html_detection_input])
+        normalizers.change(lambda value: str(value), inputs=[session_args], outputs=[current_parameters])
+        normalizers.change(fn=detect_partial, inputs=[output_without_watermark, session_args, session_tokenizer],
+                           outputs=[without_watermark_detection_result, session_args, session_tokenizer,
+                                    html_without_watermark])
+        normalizers.change(fn=detect_partial, inputs=[output_with_watermark, session_args, session_tokenizer],
+                           outputs=[with_watermark_detection_result, session_args, session_tokenizer,
+                                    html_with_watermark])
+        normalizers.change(fn=detect_partial, inputs=[detection_input, session_args, session_tokenizer],
+                           outputs=[detection_result, session_args, session_tokenizer, html_detection_input])
+        select_green_tokens.change(lambda value: str(value), inputs=[session_args], outputs=[current_parameters])
+        select_green_tokens.change(fn=detect_partial,
+                                   inputs=[output_without_watermark, session_args, session_tokenizer],
+                                   outputs=[without_watermark_detection_result, session_args, session_tokenizer,
+                                            html_without_watermark])
+        select_green_tokens.change(fn=detect_partial, inputs=[output_with_watermark, session_args, session_tokenizer],
+                                   outputs=[with_watermark_detection_result, session_args, session_tokenizer,
+                                            html_with_watermark])
+        select_green_tokens.change(fn=detect_partial, inputs=[detection_input, session_args, session_tokenizer],
+                                   outputs=[detection_result, session_args, session_tokenizer, html_detection_input])
+
+    demo.queue(concurrency_count=3)
+
+    if args.demo_public:
+        demo.launch(share=True)  # exposes app to the internet via randomly generated link
+    else:
+        demo.launch()
+
+
+def main(args):
+    """Run a command line version of the generation and detection operations
+        and optionally launch and serve the gradio demo"""
+    # Initial arg processing and log
+    args.normalizers = (args.normalizers.split(",") if args.normalizers else [])
+    print(args)
+
+    if not args.skip_model_load:
+        model, tokenizer, device = load_model(args)
+    else:
+        model, tokenizer, device = None, None, None
+        tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path)
+        if args.use_gpu:
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+        else:
+            device = "cpu"
+
+    # terrapin example
+    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"
+    )
+
+    args.default_prompt = input_text
+
+    # Generate and detect, report to stdout
+    if not args.skip_model_load:
+
+        term_width = 80
+        print("#" * term_width)
+        print("Prompt:")
+        print(input_text)
+
+        # 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,
+                                                    return_green_token_mask=False)
+        with_watermark_detection_result = detect(decoded_output_with_watermark,
+                                                 args,
+                                                 device=device,
+                                                 tokenizer=tokenizer,
+                                                 return_green_token_mask=False)
+
+        print("#" * term_width)
+        print("Output without watermark:")
+        print(decoded_output_without_watermark)
+        print("-" * term_width)
+        print(f"Detection result @ {args.detection_z_threshold}:")
+        pprint(without_watermark_detection_result)
+        print("-" * term_width)
+
+        print("#" * term_width)
+        print("Output with watermark:")
+        print(decoded_output_with_watermark)
+        print("-" * term_width)
+        print(f"Detection result @ {args.detection_z_threshold}:")
+        pprint(with_watermark_detection_result)
+        print("-" * term_width)
+
+    # Launch the app to generate and detect interactively (implements the hf space demo)
+    if args.run_gradio:
+        run_gradio(args, model=model, tokenizer=tokenizer, device=device)
+
+    return
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    print(args)
+
+    main(args)
+

homoglyphs.py

@@ -0,0 +1,265 @@
+"""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
+
+# 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)
+
+
+CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
+DATA_LOCATION = os.path.join(CURRENT_DIR, "homoglyph_data")
+
+
+class Categories:
+    """
+    Work with aliases from ISO 15924.
+    https://en.wikipedia.org/wiki/ISO_15924#List_of_codes
+    """
+
+    fpath = os.path.join(DATA_LOCATION, "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, ValueError):
+            # In Python2 unicodedata.name raise error for non-unicode chars
+            # Python3 raise ValueError for non-unicode characters
+            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(DATA_LOCATION, "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)
+        with open(os.path.join(DATA_LOCATION, "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)
+        with open(os.path.join(DATA_LOCATION, "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,7 @@
+gradio
+nltk
+scipy
+torch
+transformers
+tokenizers
+text_generation

watermark_processor.py

@@ -0,0 +1,280 @@
+# 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 = "simple_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 == "simple_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):
+
+    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 == "simple_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 == "simple_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_green_fraction:
+            score_dict.update(dict(green_fraction=(green_token_count / num_tokens_scored)))
+        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_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
+