first upload

.gitattributes

@@ -1,35 +1,34 @@
 *.7z filter=lfs diff=lfs merge=lfs -text
 *.arrow filter=lfs diff=lfs merge=lfs -text
 *.bin filter=lfs diff=lfs merge=lfs -text
+*.bin.* filter=lfs diff=lfs merge=lfs -text
 *.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
 *.ftz filter=lfs diff=lfs merge=lfs -text
 *.gz filter=lfs diff=lfs merge=lfs -text
 *.h5 filter=lfs diff=lfs merge=lfs -text
 *.joblib filter=lfs diff=lfs merge=lfs -text
 *.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
 *.model filter=lfs diff=lfs merge=lfs -text
 *.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
 *.onnx filter=lfs diff=lfs merge=lfs -text
 *.ot filter=lfs diff=lfs merge=lfs -text
 *.parquet filter=lfs diff=lfs merge=lfs -text
 *.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
 *.pth filter=lfs diff=lfs merge=lfs -text
 *.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
 saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
 *.tflite filter=lfs diff=lfs merge=lfs -text
 *.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
 *.xz filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*.tfevents* filter=lfs diff=lfs merge=lfs -text
+*.db* filter=lfs diff=lfs merge=lfs -text
+*.ark* filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*data* filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*.meta filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*.index filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,129 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/

MANIFEST.in

@@ -0,0 +1,7 @@
+# added manually
+include *.py
+include *.json
+include *.md
+
+global-exclude *.pyc
+global-exclude __pycache__

alternative_prf_schemes.py

@@ -0,0 +1,167 @@
+"""实现其他 PRF 函数（这些函数的不同之处仅在于如何从上下文中的令牌生成单个哈希值）。
+
+可作为修改后的基类 WatermarkBase 挂接到现有的 WatermarkLogitsProcessor 中，请参见
+extended_watermark_processor.py 中的实现。
+"""
+
+
+
+import torch
+from itertools import combinations
+from functools import cache
+
+# 哈希方案的关键属性
+props = {
+    "prf_type": str,  # 基础 PRF 的字符串名称，将多个令牌 ID 映射到随机种子
+    "context_width": int,  # 这是论文中的 h，每个 PRF 应考虑多少个先前的令牌
+    "self_salt": bool,  # 根据鲁棒水印技术中的规则，是否使用令牌本身来生成种子，并可能拒绝其自身的列表
+    "hash_key": int,  # 整数，大质数，用于将种子移动到上述所选 PRF 中的低熵位序列的远离位置
+}
+
+
+def seeding_scheme_lookup(seeding_scheme: str):
+    if not isinstance(seeding_scheme, str):
+        raise ValueError("Seeding scheme should be a string summarizing the procedure.")
+    if seeding_scheme == "simple_1" or seeding_scheme == "lefthash":
+        # 默认的简单二元哈希  # 别名为 ff-additive_prf-1-False-15485863
+        prf_type = "additive_prf"
+        context_width = 1
+        self_salt = False
+        hash_key = 15485863
+    elif seeding_scheme == "algorithm-3" or seeding_scheme == "selfhash":
+        prf_type = "anchored_minhash_prf"
+        context_width = 4
+        self_salt = True
+        hash_key = 15485863
+    elif seeding_scheme == "minhash":
+        prf_type = "minhash_prf"
+        context_width = 4
+        self_salt = False
+        hash_key = 15485863
+    elif seeding_scheme == "skipgram":
+        prf_type = "skipgram_prf"
+        context_width = 5
+        self_salt = False
+        hash_key = 15485863
+    elif seeding_scheme.startswith("ff"):  # 自由形式的种子方案 API - 仅用于实验目的
+        # 期望形式为 ff-additive_prf-4-True-hash 或 ff-additive_prf-5-True (哈希键是可选的)
+        split_scheme = seeding_scheme.split("-")
+        prf_type = str(split_scheme[1])
+        context_width = int(split_scheme[2])
+        self_salt = split_scheme[3] == "True"
+        if len(split_scheme) == 5:
+            hash_key = int(split_scheme[4])
+        else:
+            hash_key = 15485863
+    else:
+        raise ValueError(f"Invalid seeding scheme name {seeding_scheme} given. Try  'simple_1'?")
+
+    assert prf_type in prf_lookup.keys()
+    return prf_type, context_width, self_salt, hash_key
+
+
+def multiplicative_prf(input_ids: torch.LongTensor, salt_key: int) -> int:
+    return salt_key * input_ids.prod().item()
+
+
+def additive_prf(input_ids: torch.LongTensor, salt_key: int) -> int:
+    return salt_key * input_ids.sum().item()
+
+
+def minfunc_prf(input_ids: torch.LongTensor, salt_key: int) -> int:
+    # 对于非随机输入 id（如文本），这不是一个好主意
+    return salt_key * input_ids.min().item()
+
+
+def simple_skip_prf(input_ids: torch.LongTensor, salt_key: int, k=2) -> int:
+    # k是一个跳跃的距离
+    return hashint(salt_key * input_ids[::k]).prod().item()
+
+
+def skipgram_prf(input_ids: torch.LongTensor, salt_key: int) -> int:
+    # # 上下文内的最大距离跳字
+    return hashint(salt_key * input_ids[0]).item()
+
+
+def anchored_skipgram_prf(input_ids: torch.LongTensor, salt_key: int, anchor: int = -1) -> int:
+    # 上下文内的最大距离跳字
+    return (hashint(salt_key * input_ids[0]) * hashint(salt_key * input_ids[anchor])).item()
+
+
+def minhash_prf(input_ids: torch.LongTensor, salt_key: int) -> int:
+    return hashint(salt_key * input_ids).min().item()
+
+
+def anchored_minhash_prf(input_ids: torch.LongTensor, salt_key: int, anchor: int = -1) -> int:
+    # 另一个关键是生成一个key
+    return (salt_key * hashint(input_ids) * hashint(input_ids[anchor])).min().item()
+
+
+def minskipgram_prf(input_ids: torch.LongTensor, salt_key: int, k: int = 2) -> int:
+    # 上下文中所有跳字组合的最小值，k=2 表示所有对
+    skipgrams = torch.as_tensor(list(combinations(hashint(salt_key * input_ids), 2)))
+    return skipgrams.prod(dim=1).min().item()
+
+
+def noncomm_prf(input_ids: torch.LongTensor, salt_key: int, k: int = 2) -> int:
+    key = torch.as_tensor(salt_key, dtype=torch.long)
+    for entry in input_ids:
+        key *= hashint(key * entry)
+        key %= 2**32
+    return key.item()
+
+
+def position_prf(input_ids: torch.LongTensor, salt_key: int, k: int = 2) -> int:
+    return (salt_key * input_ids * torch.arange(1, len(input_ids) + 1, device=input_ids.device)).sum().item()
+
+
+prf_lookup = {
+    "multiplicative_prf": multiplicative_prf,
+    "additive_prf": additive_prf,
+    "minfunc_prf": minfunc_prf,
+    "simple_skip_prf": simple_skip_prf,
+    "skipgram_prf": skipgram_prf,
+    "anchored_skipgram_prf": anchored_skipgram_prf,
+    "minhash_prf": minhash_prf,
+    "anchored_minhash_prf": anchored_minhash_prf,
+    "minskipgram_prf": minskipgram_prf,
+    "noncomm_prf": noncomm_prf,
+    "position_prf": position_prf,
+}
+
+# 在启动时生成全局置换表一次
+rng = torch.Generator(device=torch.device("cpu"))
+rng.manual_seed(2971215073)  
+table_size = 1_000_003
+fixed_table = torch.randperm(1_000_003, device=torch.device("cpu"), generator=rng)  # 这个速度很快
+
+
+def hashint(integer_tensor: torch.LongTensor) -> torch.LongTensor:
+    
+    return fixed_table[integer_tensor.cpu() % table_size] + 1  # 这里有一个小技巧，这个函数总是返回 CPU 的值
+
+
+def _hashint_avalanche_tensor(integer_tensor: torch.LongTensor):
+    
+    i = integer_tensor.to(torch.int32).clone()  # or torch.int16?
+    i -= i << 6
+    i ^= i >> 17
+    i -= i << 9
+    i ^= i << 4
+    i -= i << 3
+    i ^= i << 10
+    i ^= i >> 15
+    return i.to(torch.long)
+
+
+@cache
+def _hashint_avalanche_int(integer: int):
+    i = integer % (2**32)
+    i -= i << 6
+    i ^= i >> 17
+    i -= i << 9
+    i ^= i << 4
+    i -= i << 3
+    i ^= i << 10
+    i ^= i >> 15
+    return i

app.py

@@ -0,0 +1,39 @@
+# 安装好环境
+# python app.py即可运行
+import os
+
+os.environ['HF_ENDPOINT']='https://hf-mirror.com'
+
+from argparse import Namespace
+args = Namespace()
+
+arg_dict = {
+    'run_gradio': True,
+    'demo_public': False,
+    'model_name_or_path': 'Qwen/Qwen2-0.5B-Instruct',
+
+    # 'model_name_or_path': 'Qwen/Qwen2-0.5B-Instruct-GGUF',
+    'gguf_file': './qwen2-0_5b-instruct-q8_0.gguf', # 只有用gguf模型会用到，即model_name_or_path里含有gguf字符串才会用到
+    'prompt_max_length': None,
+    'max_new_tokens': 500,
+    'generation_seed': 123,
+    'use_sampling': True,
+    'n_beams': 1,
+    'sampling_temp': 0.7,
+    'use_gpu': True,
+    'seeding_scheme': 'simple_1',
+    'gamma': 0.5,
+    'delta': 2.0,
+    'normalizers': '',
+    'ignore_repeated_bigrams': False,
+    'detection_z_threshold': 4.0,
+    'select_green_tokens': True,
+    'skip_model_load': False,
+    'seed_separately': True,
+}
+
+args.__dict__.update(arg_dict)
+
+from demo_watermark import main
+
+main(args)

demo_watermark.py

@@ -0,0 +1,1085 @@
+import argparse
+import os
+import time
+from functools import partial
+
+import gradio.exceptions
+
+# 设置OpenMP线程数为8,优化CPU并行计算性能
+os.environ["OMP_NUM_THREADS"] = "8"
+
+import gradio as gr
+import pandas as pd
+import torch
+from transformers import (AutoTokenizer,
+                          AutoModelForCausalLM,
+                          LogitsProcessorList)
+
+from watermark_processor import WatermarkLogitsProcessor, WatermarkDetector
+
+# FIXME 所有模型的正确长度
+
+API_MODEL_MAP = {
+    # "Qwen/Qwen1.5-0.5B-Chat": {"max_length": 2000, "gamma": 0.5, "delta": 2.0},
+    # "THUDM/chatglm3-6b": {"max_length": 2048, "gamma": 0.5, "delta": 2.0},
+}
+
+default_trace_table = pd.DataFrame(columns=["编号", "水印内容"])
+default_trace_table.loc[0] = (0, "默认用户")
+default_trace_table.loc[1] = (1, "张三")
+default_trace_table.loc[2] = (2, "李四")
+
+watermark_salt = 0
+
+
+def str2bool(v):
+    """用户友好的布尔标志参数的Util函数"""
+    if isinstance(v, bool):
+        return v
+    if v.lower() in ('yes', 'true', 't', 'y', '1'):
+        return True
+    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
+        return False
+    else:
+        raise argparse.ArgumentTypeError('Boolean value expected.')
+
+
+# 定义一个函数用于解析命令行参数
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="")
+
+    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="Qwen/Qwen1.5-0.5B-Chat",
+        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.5,
+        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(
+        "--gguf_file",
+        type=str,
+        default='./qwen2-0_5b-instruct-q2_k.gguf',
+        help="gguf文件（如果有）",
+    )
+
+    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.",
+    )
+    args = parser.parse_args()
+
+    return args
+
+
+def load_model(args):
+    """加载并返回模型和分词器"""
+
+    if args.use_gpu:
+        device = "cuda:0" if torch.cuda.is_available() else "cpu"
+    else:
+        device = "cpu"
+
+    if 'gguf' in args.model_name_or_path.lower():
+
+        model = AutoModelForCausalLM.from_pretrained(args.model_name_or_path, gguf_file=args.gguf_file,
+                                                     trust_remote_code=True,
+                                                     local_files_only=True,
+                                                     device_map=device)
+        tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path, gguf_file=args.gguf_file,
+                                                  local_files_only=True,
+                                                  trust_remote_code=True)
+
+    else:
+        model = AutoModelForCausalLM.from_pretrained(args.model_name_or_path,
+                                                     trust_remote_code=True,
+                                                     local_files_only=True,
+                                                     device_map=device)
+
+        tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path, trust_remote_code=True,
+                                                  local_files_only=True, )
+
+    try:
+        model.eval()
+    except Exception as e:
+        print(e)
+
+    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")
+    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):
+    # 这适用于本地和API模型场景
+    try:
+        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 = 4096 - args.max_new_tokens
+    except Exception as e:
+        print(e)
+        args.prompt_max_length = 4096 - args.max_new_tokens
+
+    tokd_input = tokenizer(prompt, return_tensors="pt", add_special_tokens=False, 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):
+    """根据水印参数实例化 WatermarkLogitsProcessor 并通过将其作为 logits 处理器传递给模型的 generate 方法来生成带水印的文本。"""
+    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:
+        if 'chatglm' in args.model_name_or_path.lower() or 'qwen' in args.model_name_or_path.lower() or 'llama' in args.model_name_or_path.lower():
+            messages = [
+                # {"role": "system", "content": "You are a helpful assistant."},
+                {"role": "user", "content": prompt}
+            ]
+
+            tokenized_input = tokenizer.apply_chat_template(
+                messages,
+                tokenize=False,
+                add_generation_prompt=True
+            )
+
+            tokd_input = tokenizer([tokenized_input], return_tensors="pt", truncation=True, add_special_tokens=False,
+                                   max_length=args.prompt_max_length).to(device)
+
+        else:
+            tokd_input = tokenizer(prompt, return_tensors="pt", add_special_tokens=True, truncation=True,
+                                   max_length=args.prompt_max_length).to(device)
+
+        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
+            ))
+
+        watermark_processor = WatermarkLogitsProcessor(vocab=list(tokenizer.get_vocab().values()),
+                                                       gamma=args.gamma,
+                                                       delta=args.delta,
+                                                       seeding_scheme=args.seeding_scheme,
+                                                       extra_salt=watermark_salt,
+                                                       select_green_tokens=args.select_green_tokens)
+
+        generate_without_watermark = partial(
+            model.generate,
+            **gen_kwargs
+        )
+
+        generate_with_watermark = partial(
+            model.generate,
+            logits_processor=LogitsProcessorList([watermark_processor]),
+            **gen_kwargs
+        )
+
+        start_time = time.time()
+        gr.Info('开始生成正常内容')
+        torch.manual_seed(args.generation_seed)
+        output_without_watermark = generate_without_watermark(**tokd_input)
+
+        # 可选择在第二次生成之前种子，但通常不会再次相同，除非 delta==0.0，无操作水印
+
+        print(watermark_salt)
+        print(default_trace_table)
+        print(default_trace_table.loc[default_trace_table['编号'] == watermark_salt, '水印内容'])
+        gr.Info('开始注入水印：“{}”'.format(
+            default_trace_table.loc[default_trace_table['编号'] == watermark_salt, '水印内容'].item()))
+        if args.seed_separately:
+            torch.manual_seed(args.generation_seed)
+
+        output_with_watermark = generate_with_watermark(**tokd_input)
+
+        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]
+
+        end_time = time.time()
+        gr.Info(f"生成结束，共用时{end_time - start_time:.2f}秒")
+
+        print(f"Generation took {end_time - start_time:.2f} seconds")
+
+        # 使用空格分隔生成器风格模拟 API 输出
+
+        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):
+    """为 gradio 演示界面格式化名称"""
+    s = s.replace("num_tokens_scored", "总Token")
+    s = s.replace("num_green_tokens", "Green Token数量")
+    s = s.replace("green_fraction", "Green Token占比")
+    s = s.replace("z_score", "z-score")
+    s = s.replace("p_value", "p value")
+    s = s.replace("prediction", "预测结果")
+    s = s.replace("confidence", "置信度")
+    return s
+
+
+def list_format_scores(score_dict, detection_threshold):
+    """将检测指标格式化为 gradio 数据框输入格式"""
+    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), ("含有水印" if v else "无水印")])
+        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):
+    """实例化 WatermarkDetection 对象并调用 detect 方法  在输入文本上返回测试的分数和结果"""
+
+    print(f"Detecting with {args}")
+    print(f"Detection Tokenizer: {type(tokenizer)}")
+
+    # 现在不要显示绿色的token mask
+    # 如果我们使用的是normalizers或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:
+            for _, data in default_trace_table.iterrows():
+                salt = data["编号"]
+                name = data["水印内容"]
+                watermark_detector = WatermarkDetector(vocab=list(tokenizer.get_vocab().values()),
+                                                       gamma=args.gamma,
+                                                       seeding_scheme=args.seeding_scheme,
+                                                       extra_salt=salt,
+                                                       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)
+                score_dict = watermark_detector.detect(input_text, return_green_token_mask=return_green_token_mask)
+                if score_dict['prediction']:
+                    print(f"检测到是“{name}”的水印")
+                    break
+
+            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 None:
+        html_output = "[Visualizing masks with ignore_repeated_bigrams enabled is not supported, toggle off to see the mask for this text. The mask is the same in both cases - only counting/stats are affected.]"
+
+    if green_token_mask is not None:
+        # hack 因为我们需要一个带有字符跨度支持的快速分词器
+        tokens = tokenizer(input_text, add_special_tokens=False)
+        if tokens["input_ids"][0] == tokenizer.bos_token_id:
+            tokens["input_ids"] = tokens["input_ids"][1:]  # 忽略注意力掩码
+        skip = watermark_detector.min_prefix_len
+
+        if args.model_name_or_path in ['THUDM/chatglm3-6b']:
+            # 假设词表中3-258就是字节0-255
+            charspans = []
+            for i in range(skip, len(tokens["input_ids"])):
+                if tokens.data['input_ids'][i - 1] in range(3, 259):
+                    charspans.append("<0x{:X}>".format(tokens.data['input_ids'][i - 1] - 3))
+                else:
+                    charspans.append(tokenizer.decode(tokens.data['input_ids'][i - 1:i]))
+
+        else:
+            charspans = [tokens.token_to_chars(i - 1) 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)
+
+        if args.model_name_or_path in ['THUDM/chatglm3-6b']:
+            tags = []
+            for cs, m in zip(charspans, green_token_mask):
+                tags.append(
+                    f'<span class="green">{cs}</span>' if m else f'<span class="red">{cs}</span>')
+
+        else:
+            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>'
+
+        if score_dict['prediction']:
+            html_look = gr.HTML("""<div style="width: 100%; font-size: 24px; height: 100px; border-radius: 20px; background-color: rgba(255, 0, 0, 0.25); display: flex; justify-content: center; align-items: center; color: white; font-weight: bold;">
+        <span>有 “{}” 的水印</span>
+    </div>""".format(name), visible=True)
+
+        else:
+            html_look = gr.HTML("""<div style="width: 100%; font-size: 24px; height: 100px; border-radius: 20px; background-color: rgba(0, 128, 0, 0.25); display: flex; justify-content: center; align-items: center; color: white; font-weight: bold; text-align: center;">
+        <span>无水印</span>
+    </div>""", visible=True)
+
+    return output, args, tokenizer, html_output, html_look
+
+
+def run_gradio(args, model=None, device=None, tokenizer=None):
+    """定义并启动gradio演示界面"""
+    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='ParityError/Interstellar', css=css) as demo:
+        # 顶部部分，问候语和说明
+        with gr.Row():
+            with gr.Column(scale=9):
+                gr.Markdown(
+                    """
+                    # 🌸🖼️ LLMwatermark:面向大语言模型生成内容的数字水印版权保护系统 🌟🎓
+                    """
+                )
+            with gr.Column(scale=1):
+                # 如果启动时的 model_name_or_path 不是 API 模型之一，则添加到下拉菜单中
+                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="选择大语言模型，进行模型水印",
+                )
+
+        # 构建参数的状态，定义更新和切换
+        default_prompt = args.__dict__.pop("default_prompt")
+        session_args = gr.State(value=args)
+        # 注意，如果状态对象是可调用的，则自动调用 value，希望在启动时避免调用分词器
+        session_tokenizer = gr.State(value=lambda: tokenizer)
+
+        with gr.Tab("生成回答和添加文本水印🎓"):
+
+            with gr.Row():
+                with gr.Column(scale=5):
+                    prompt = gr.Textbox(label=f"Prompt", interactive=True, lines=3, max_lines=10, value=default_prompt)
+                with gr.Column(scale=3):
+                    trace_source = gr.Dataframe(default_trace_table, datatype=['number', 'str'], interactive=True,
+                                                col_count=(2, "fixed"))
+            with gr.Row(equal_height=True):
+                with gr.Column(scale=7):
+                    generate_btn = gr.Button("Generate", variant='primary')
+
+                gr.Markdown('水印选择：',
+                            show_label=False)
+                watermark_salt_choice = gr.Dropdown(
+                    choices=[i[::-1] for i in default_trace_table.to_dict(orient='split')['data']],
+                    value=0,
+                    container=False,
+                    scale=3,
+                    type="value",
+                    interactive=True, label="水印标识选择")
+
+            with gr.Row():
+                with gr.Column():
+                    with gr.Column(scale=2):
+                        with gr.Tab("原版输出"):
+                            output_without_watermark = gr.Textbox(interactive=False, lines=7, max_lines=14,
+                                                                  show_label=False)
+                        with gr.Tab("显示水印"):
+                            html_without_watermark = gr.HTML(elem_id="html-without-watermark")
+
+                        original_watermark_state = gr.HTML('', visible=False)
+
+                    with gr.Column(scale=1):
+                        without_watermark_detection_result = gr.Dataframe(headers=["Metric", "Value"],
+                                                                          interactive=False,
+                                                                          row_count=7, col_count=2)
+                with gr.Column():
+                    with gr.Column(scale=2):
+                        with gr.Tab("带水印的输出"):
+                            output_with_watermark = gr.Textbox(interactive=False, lines=7, max_lines=14,
+                                                               show_label=False)
+                        with gr.Tab("显示水印"):
+                            html_with_watermark = gr.HTML(elem_id="html-with-watermark")
+
+                        change_watermark_state = gr.HTML('', visible=False)
+
+                    with gr.Column(scale=1):
+                        with_watermark_detection_result = gr.Dataframe(headers=["Metric", "Value"], 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=5):
+                    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")
+
+                    detect_watermark_state = gr.HTML('', visible=False)
+                with gr.Column(scale=2):
+                    trace_source2 = gr.Dataframe(default_trace_table, datatype=['number', 'str'], interactive=True,
+                                                 col_count=(2, "fixed"))
+                    detection_result = gr.Dataframe(headers=["Metric", "Value"], interactive=False, row_count=7,
+                                                    col_count=2)
+
+            with gr.Row():
+                detect_btn = gr.Button("检测", variant='primary')
+
+        with gr.Tab("About📖"):
+            with gr.Row():
+                with gr.Column(scale=2):
+                    gr.Markdown(
+                        """
+                        大语言模型可能带来的潜在危害可以通过*水印*来减轻。*水印*是嵌入在生成的文本中的信息，
+                        这对人类来说是不可见的，但是可以被特定算法检测到。
+                        这些水印可以使*任何人*用特定工具判断其是否使用带水印的模型生成的。
+                        本网站展示了一种水印方法，可以应用于_任何_生成性语言模型。
+                        """
+                    )
+                    gr.Markdown(
+                        """
+                        **[生成文本与添加水印]**：可以给大模型的输出添加水印。
+                        您可以尝试任何prompt，并比较正常文本(*没有水印的输出*)和水印文本(*有水印的输出*)的质量。
+                        您还可以点击**显示水印**来“看到”水印，其中的颜色表示其所在的红绿表。
+
+                        **[检测]**：您还可以将水印文本（或任何其他文本）复制粘贴到第二个选项卡中。
+                        可以实验删除多少句子后还能检测到水印。
+                        还可以在验证，检测器的误报率有多少；
+                        """
+                    )
+
+                with gr.Column(scale=1):
+                    gr.Markdown(
+                        """
+                        ![]()
+                        """
+                    )
+
+        # 参数选择组
+        with gr.Accordion("高级设置", open=False):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    gr.Markdown(f"#### 生成参数")
+                    with gr.Row():
+                        decoding = gr.Radio(label="解码方法", choices=["多项式解码方法", "贪心解码方法"],
+                                            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) and (
+                                                  not args.model_name_or_path in API_MODEL_MAP)))
+                    with gr.Row():
+                        max_new_tokens = gr.Slider(label="（生成文本的最大长度）Max Generated Tokens", minimum=10,
+                                                   maximum=4000, 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分数阈值", minimum=0.0, maximum=10.0, step=0.1,
+                                                          value=args.detection_z_threshold)
+                    with gr.Row():
+                        ignore_repeated_bigrams = gr.Checkbox(label="避免生成连续重复的双词组合")
+                    with gr.Row():
+                        normalizers = gr.CheckboxGroup(label="对文本进行标准化处理",
+                                                       choices=["unicode", "homoglyphs", "truecase"],
+                                                       value=args.normalizers)
+            with gr.Row():
+                gr.Markdown(
+                    f"注意：滑块并不总是能完美更新。点击条形图或使用右侧的数字窗口会有所帮助。下面的窗口显示当前设置。")
+            with gr.Row():
+                current_parameters = gr.Textbox(label="当前参数设置", value=args, max_lines=10)
+            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="从分区中选择 绿色列表", value=args.select_green_tokens)
+
+        with gr.Accordion("设置有什么作用？", open=False):
+            gr.Markdown(
+                """
+                #### 生成参数:
+    
+                - **解码方法**：我们可以使用多项式采样或者贪婪解码的方式从模型中生成标记。
+                    决定如何从模型中生成token。可以选择多项式采样或贪婪解码。
+                    多项式采样允许一定的随机性，而贪婪解码总是选择概率最高的下一个token。
+                - **采样温度**：如果使用多项式采样，我们可以设置采样分布的温度。
+                - 0.0 相当于贪婪解码，而 1.0 代表最大的随机性。
+                - 0.7是文本质量和随机性之间的平衡点。不适用于贪婪解码。
+                - **生成种子**：用于在生成前初始化随机数生成器，使多项式采样的输出可复现。此设置不适用于贪婪解码。
+                - **束搜索数量**：在使用贪婪解码时，可以设置光束数量来启用束搜索。
+                    这允许考虑多个候选序列，而不是只选择最有可能的一个。此设置目前仅适用于贪婪解码。
+                - **最大生成标记数**：传递给生成方法的 `max_new_tokens` 参数，以在一定数量的新标记停止输出。
+                - 请注意，根据提示，模型可以生成较少的标记。隐含地，这将最大化可能的提示标记数，即模型的最大输入长度减去 `max_new_tokens`，并相应地截断输入。
+    
+                    综上所述，这些参数提供了对生成过程的不同方面的控制，包括随机性和多样性
+                    （解码方法、采样温度），可复现性（生成种子），以及输出长度和多样性（光束数量、最大生成token数）。
+                    合理配置这些参数可以帮助生成高质量的文本输出。
+    
+                #### 水印参数:
+    
+                - **gamma**：在每个生成步骤中将词汇表的一部分划分为绿色列表的比例。
+                - 较小的 gamma 值通过使水印模型优先从较小的绿色集中采样，从而使其与人类/未水印文本的差异更大，从而创建更强的水印。
+                - **delta**：在每个生成步骤中为绿色列表中的每个标记的 logits 添加的正偏差量。较高的 delta 值意味着水印模型更偏好于绿色列表中的标记。
+                - 随着偏差变得非常大，水印从 "软" 过渡到 "硬"。对于硬水印，几乎所有标记都是绿色的，但这可能对生成质量产生不利影响，特别是当分布的灵活性不大时。
+    
+                #### 检测器参数:
+    
+                - **z 分数阈值**：假设检验的 z 分数截断。较高的阈值（例如 4.0）使得 _false positives_（预测人类/未水印文本被标记为水印）非常不可能，
+                  因为一个真正的人类文本几乎永远不会达到那么高的 z 分数。
+                - 较低的阈值会捕获更多的 _true positives_，因为一些水印文本可能包含较少的绿色标记并达到较低的 z 分数，但仍然通过较低的门槛并被标记为 "水印"。
+                  然而，较低的阈值会增加包含略高于平均水平的绿色标记的人类文本错误地被标记为水印的几率。4.0-5.0 提供了极低的假阳性率，同时仍准确捕获大多数水印文本。
+                - **忽略二元重复**：这种替代的检测算法在检测期间仅考虑文本中的唯一二元组，根据每对中的第一个计算绿色列表，并检查第二个是否位于列表中。
+                - 这意味着 `T` 现在是文本中唯一二元组的数量，如果文本包含大量重复，则这个数字将小于生成的总标记数。有关更详细的讨论，请参阅论文。
+                - **归一化**：我们实现了一些基本的归一化来抵御文本在检测期间的各种对抗性扰动。
+                - 目前，我们支持将所有字符转换为 Unicode，将同形异义字符替换为规范形式，并标准化大写。有关输入归一化的详细讨论，请参阅论文。
+                """
+            )
+
+        with gr.Accordion("输出指标意味着什么？", open=False):
+            gr.Markdown(
+                """
+                - `z-score threshold`：假设检验的截止值。
+                - `Tokens Counted (T)`：检测算法计算的输出中的标记数量。在简单的、单标记播种方案中，第一个标记被省略了，因为没有办法为其生成绿色列表，因为它没有前缀标记。
+                    在底部面板描述的“忽略二元重复”检测算法下，如果有很多重复，这个数量可能远少于生成的总标记数。
+                - ` Tokens in Greenlist`：观察到落在其相应绿色列表中的标记数量。
+                - `Fraction of T in Greenlist`：`# Tokens in Greenlist` / `T`。这应该大约等于人类/未水印文本的 `gamma`。
+                - `z-score`：用于检测假设检验的检验统计量。如果大于 `z-score threshold`，则我们“拒绝零假设”，即文本是人类/未水印的，并得出结论它是带水印的。
+                - `p value`：在零假设下观察到计算出的 `z-score` 的可能性。这是在不知道水印程序/绿色列表的情况下观察到 `Fraction of T in Greenlist` 的可能性。
+                    如果这个值极其 _小_，我们可以确信这么多的绿色标记不是由随机机会选择的。
+                - `prediction`：假设检验的结果 - 观察到的 `z-score` 是否高于 `z-score threshold`。
+                - `confidence`：如果我们拒绝零假设，且 `prediction` 是“带水印”，那么我们报告 1-`p value` 来表示基于这个 `z-score` 观察的检测的置信度。
+                """
+            )
+
+        gr.HTML("""
+                <p>本方法可以对任何大模型的输出结果进行水印的操作。
+                    并且可以对输出结果进行水印检测。
+                <p/>
+                """)
+
+        # 注册主要生成标签单击事件，输出生成文本以及编码+重新解码+可能被截断的提示和标志，然后调用检测
+        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, original_watermark_state]).success(
+            fn=detect_partial, inputs=[output_with_watermark, session_args, session_tokenizer],
+            outputs=[with_watermark_detection_result, session_args, session_tokenizer, html_with_watermark,
+                     change_watermark_state])
+        # 如果发生了截断，则显示提示的截断版本
+        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,
+                                  detect_watermark_state],
+                         api_name="detection")
+
+        # 状态管理逻辑
+        # 定义更新回调函数以更改状态字典
+        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_watermark_salt(value):
+            global watermark_salt
+            if isinstance(value, int):
+                watermark_salt = value
+            elif value is None:
+                watermark_salt = 0
+            elif isinstance(value, str) and value.isdigit():
+                watermark_salt = int(value)
+            else:
+                # 不知道为什么会出现这种倒置的情况
+                watermark_salt = int(
+                    default_trace_table.loc[default_trace_table['水印内容'] == value, '编号'].item())
+
+        def update_trace_source(value):
+            global default_trace_table
+            try:
+                if '' in value.loc[:, '编号'].tolist():
+                    return value, gr.Dropdown()
+
+                value.loc[:, '编号'] = value.loc[:, '编号'].astype(int)
+
+                if default_trace_table.duplicated(subset='编号').any():
+                    raise gr.Error(f"请检查水印编号，编号不能重复")
+
+                default_trace_table = value
+
+                return value, gr.Dropdown(
+                    choices=[i[::-1] for i in value.to_dict(orient='split')['data']])
+
+            except ValueError as e:
+                if 'invalid literal for int() with base 10' in str(e):
+                    raise gr.Error(f"请检查水印数据，编号必须是整数：{e}")
+
+            except gradio.exceptions.Error as e:
+                raise e
+
+            except Exception as e:
+                print(type(e))
+                raise e
+
+        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)
+
+        # 如果模型名称在 API 模型列表中，则将 num beams 参数设置为 1 并隐藏 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)
+
+        # 如果模型名称在 API 模型列表中，则将交互参数设置为 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)
+
+        # 如果模型名称在 API 模型列表中，则根据 API 映射设置 gamma 和 delta
+        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 = 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
+
+        # 强制约束模型��能为 null 或空
+        # 然后特别附加模型回调函数
+        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]
+        )
+        # 根据其他参数的值注册回调函数以切换特定参数的可见性
+        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])
+        # 注册所有状态更新回调函数
+        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])
+        watermark_salt_choice.change(update_watermark_salt, inputs=[watermark_salt_choice])
+
+        # 同步更新
+        trace_source.change(update_trace_source, inputs=[trace_source],
+                            outputs=[trace_source2, watermark_salt_choice])
+        trace_source2.change(update_trace_source, inputs=[trace_source2],
+                             outputs=[trace_source, watermark_salt_choice])
+
+        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])
+        # 注册按钮点击时更新显示参数窗口的额外回调
+        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])
+        # 当参数更改时，显示更新并触发检测，因为某些检测参数不会改变模型输出。
+        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) # delete
+
+    if args.demo_public:
+        demo.launch(share=True)  # 通过随机生成的链接将应用程序暴露到互联网上
+    else:
+        demo.launch(server_name='0.0.0.0', share=False)
+
+
+def main(args):
+    """运行生成和检测操作的命令行版本
+    并可选择启动和提供 gradio 演示"""
+    # 初始参数处理和日志记录
+    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:0" if torch.cuda.is_available() else "cpu"
+        else:
+            device = "cpu"
+
+    # terrapin example
+    input_text = (
+        "为什么A股指数跌的不多，但是我亏损比之前都多？"
+    )
+
+    args.default_prompt = input_text
+
+    # Generate and detect, report to stdout
+    if not args.skip_model_load:
+        pass
+
+    # 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)

extended_watermark_processor.py

@@ -0,0 +1,592 @@
+
+from __future__ import annotations
+import collections
+from math import sqrt
+from itertools import chain, tee
+from functools import lru_cache
+
+import scipy.stats
+import torch
+from tokenizers import Tokenizer
+from transformers import LogitsProcessor
+
+from normalizers import normalization_strategy_lookup
+from alternative_prf_schemes import prf_lookup, seeding_scheme_lookup
+
+
+class WatermarkBase:
+    def __init__(
+        self,
+        vocab: list[int] = None,
+        gamma: float = 0.25,
+        delta: float = 2.0,
+        seeding_scheme: str = "selfhash",  
+        select_green_tokens: bool = True,  
+    ):
+        # 现在可能将 None 作为 seeding_scheme 传递，所以现在要修补
+        if seeding_scheme is None:
+            seeding_scheme = "selfhash"
+
+        # 词汇设置
+        self.vocab = vocab
+        self.vocab_size = len(vocab)
+
+        # 水印行为：
+        self.gamma = gamma
+        self.delta = delta
+        self.rng = None
+        self._initialize_seeding_scheme(seeding_scheme)
+        # 传统行为：
+        self.select_green_tokens = select_green_tokens
+
+    def _initialize_seeding_scheme(self, seeding_scheme: str) -> None:
+        """从一个通俗的“公共”名称初始化种子策略的所有内部设置。"""
+        self.prf_type, self.context_width, self.self_salt, self.hash_key = seeding_scheme_lookup(seeding_scheme)
+
+    def _seed_rng(self, input_ids: torch.LongTensor) -> None:
+        """从本地上下文种子 RNG。不进行批处理，因为我们使用的生成器（如 cuda.random）不进行批处理。"""
+        # 需要有足够的token来进行种子生成
+        if input_ids.shape[-1] < self.context_width:
+            raise ValueError(f"seeding_scheme requires at least a {self.context_width} token prefix to seed the RNG.")
+
+        prf_key = prf_lookup[self.prf_type](input_ids[-self.context_width :], salt_key=self.hash_key)
+        self.rng.manual_seed(prf_key % (2**64 - 1))  # 防止溢出
+
+    def _get_greenlist_ids(self, input_ids: torch.LongTensor) -> torch.LongTensor:
+        """根据本地上下文宽度对rng进行种子处理，并使用这些信息在绿色列表上生成id"""
+        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:  # 通过红色选择绿色
+            greenlist_ids = vocab_permutation[(self.vocab_size - greenlist_size) :]  # legacy behavior
+        return greenlist_ids
+
+
+class WatermarkLogitsProcessor(WatermarkBase, LogitsProcessor):
+    """LogitsProcessor 在管道中修改模型输出分数。可以在任何 HF 管道中使用它来修改分数以适应水印，但也可以作为一个独立的工具插入到在模型输出和下一个标记采样器之间生成分数的任何模型中。"""
+    def __init__(self, *args, store_spike_ents: bool = False, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self.store_spike_ents = store_spike_ents
+        self.spike_entropies = None
+        if self.store_spike_ents:
+            self._init_spike_entropies()
+
+    def _init_spike_entropies(self):
+        alpha = torch.exp(torch.tensor(self.delta)).item()
+        gamma = self.gamma
+
+        self.z_value = ((1 - gamma) * (alpha - 1)) / (1 - gamma + (alpha * gamma))
+        self.expected_gl_coef = (gamma * alpha) / (1 - gamma + (alpha * gamma))
+
+        # 当bias 是 "infinite" 时候捕获溢出
+        if alpha == torch.inf:
+            self.z_value = 1.0
+            self.expected_gl_coef = 1.0
+
+    def _get_spike_entropies(self):
+        spike_ents = [[] for _ in range(len(self.spike_entropies))]
+        for b_idx, ent_tensor_list in enumerate(self.spike_entropies):
+            for ent_tensor in ent_tensor_list:
+                spike_ents[b_idx].append(ent_tensor.item())
+        return spike_ents
+
+    def _get_and_clear_stored_spike_ents(self):
+        spike_ents = self._get_spike_entropies()
+        self.spike_entropies = None
+        return spike_ents
+
+    def _compute_spike_entropy(self, scores):
+        # 预先计算z得分
+        probs = scores.softmax(dim=-1)
+        denoms = 1 + (self.z_value * probs)
+        renormed_probs = probs / denoms
+        sum_renormed_probs = renormed_probs.sum()
+        return sum_renormed_probs
+
+    def _calc_greenlist_mask(self, scores: torch.FloatTensor, greenlist_token_ids) -> torch.BoolTensor:
+        green_tokens_mask = torch.zeros_like(scores, dtype=torch.bool)
+        for b_idx, greenlist in enumerate(greenlist_token_ids):
+            if len(greenlist) > 0:
+                green_tokens_mask[b_idx][greenlist] = True
+        return green_tokens_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 _score_rejection_sampling(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, tail_rule="fixed_compute") -> list[int]:
+        """基于当前候选的下一个标记生成绿名单。如果需要，拒绝并继续。该方法不进行批处理。
+    这只是算法3“鲁棒私有水印”的部分版本，因为它始终假设贪婪采样。它仍然（有点）可以对所有类型的采样进行工作，但效果较差。
+    为了高效工作，此函数可以在处理分布尾部的规则之间切换。
+    默认情况下不会公开这些规则。"""
+
+        sorted_scores, greedy_predictions = scores.sort(dim=-1, descending=True)
+
+        final_greenlist = []
+        for idx, prediction_candidate in enumerate(greedy_predictions):
+            greenlist_ids = self._get_greenlist_ids(torch.cat([input_ids, prediction_candidate[None]], dim=0))  # add candidate to prefix
+            if prediction_candidate in greenlist_ids:  # test for consistency
+                final_greenlist.append(prediction_candidate)
+
+            # 为了提高效率，以下是可选的提前停止规则
+            if tail_rule == "fixed_score":
+                if sorted_scores[0] - sorted_scores[idx + 1] > self.delta:
+                    break
+            elif tail_rule == "fixed_list_length":
+                if len(final_greenlist) == 10:
+                    break
+            elif tail_rule == "fixed_compute":
+                if idx == 40:
+                    break
+            else:
+                pass  # do not break early
+        return torch.as_tensor(final_greenlist, device=input_ids.device)
+
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
+        """使用上一个词作为input_ids进行调用，并为下一个token打分。"""
+
+        self.rng = torch.Generator(device=input_ids.device) if self.rng is None else self.rng
+
+        # 注意，去掉这个批循环会很好，但当前，种子和分区操作都不是张量/向量化的，因此批中的每个序列都需要单独处理。
+
+        list_of_greenlist_ids = [None for _ in input_ids]  # Greenlists could differ in length
+        for b_idx, input_seq in enumerate(input_ids):
+            if self.self_salt:
+                greenlist_ids = self._score_rejection_sampling(input_seq, scores[b_idx])
+            else:
+                greenlist_ids = self._get_greenlist_ids(input_seq)
+            list_of_greenlist_ids[b_idx] = greenlist_ids
+
+            # 计算和存储熵
+            if self.store_spike_ents:
+                if self.spike_entropies is None:
+                    self.spike_entropies = [[] for _ in range(input_ids.shape[0])]
+                self.spike_entropies[b_idx].append(self._compute_spike_entropy(scores[b_idx]))
+
+        green_tokens_mask = self._calc_greenlist_mask(scores=scores, greenlist_token_ids=list_of_greenlist_ids)
+        scores = self._bias_greenlist_logits(scores=scores, greenlist_mask=green_tokens_mask, greenlist_bias=self.delta)
+
+        return scores
+
+
+class WatermarkDetector(WatermarkBase):
+    """这是用于检测所有使用 WatermarkLogitsProcessor 印记的水印的检测器。
+
+     检测器需要给出在文本生成期间给出的完全相同的设置，以复制水印的生成绿名单并检测水印。
+    这包括在文本生成期间使用的正确设备、正确的分词器、正确的 seeding_scheme 名称和参数（delta、gamma）。
+
+    可选参数包括
+    * normalizers ["unicode", "homoglyphs", "truecase"] -> 这些可以减轻生成文本中可能触发水印的修改。
+    * ignore_repeated_ngrams -> 此选项将更改检测规则，只计算每个唯一 ngram 一次。
+    * z_threshold -> 更改此阈值将更改检测器的灵敏度。
+    """
+
+
+    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_ngrams: bool = True,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+        # 配置选项
+        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)
+
+        self.normalizers = []
+        for normalization_strategy in normalizers:
+            self.normalizers.append(normalization_strategy_lookup(normalization_strategy))
+        self.ignore_repeated_ngrams = ignore_repeated_ngrams
+
+    def dummy_detect(
+        self,
+        return_prediction: bool = True,
+        return_scores: bool = True,
+        z_threshold: float = None,
+        return_num_tokens_scored: bool = True,
+        return_num_green_tokens: bool = True,
+        return_green_fraction: bool = True,
+        return_green_token_mask: bool = False,
+        return_all_window_scores: bool = False,
+        return_z_score: bool = True,
+        return_z_at_T: bool = True,
+        return_p_value: bool = True,
+    ):
+        # HF-style 输出字典
+        score_dict = dict()
+        if return_num_tokens_scored:
+            score_dict.update(dict(num_tokens_scored=float("nan")))
+        if return_num_green_tokens:
+            score_dict.update(dict(num_green_tokens=float("nan")))
+        if return_green_fraction:
+            score_dict.update(dict(green_fraction=float("nan")))
+        if return_z_score:
+            score_dict.update(dict(z_score=float("nan")))
+        if return_p_value:
+            z_score = score_dict.get("z_score")
+            if z_score is None:
+                z_score = float("nan")
+            score_dict.update(dict(p_value=float("nan")))
+        if return_green_token_mask:
+            score_dict.update(dict(green_token_mask=[]))
+        if return_all_window_scores:
+            score_dict.update(dict(window_list=[]))
+        if return_z_at_T:
+            score_dict.update(dict(z_score_at_T=torch.tensor([])))
+
+        output_dict = {}
+        if return_scores:
+            output_dict.update(score_dict)
+        # 如果通过return_prediction，则执行假设检验并返回结果
+        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"] = False
+
+        return output_dict
+
+    def _compute_z_score(self, observed_count, T):
+        # count是指绿色token的数量，T是token的总数
+        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
+
+    @lru_cache(maxsize=2**32)
+    def _get_ngram_score_cached(self, prefix: tuple[int], target: int):
+        """缓存了re-seeding and sampling"""
+        # 需要小心处理， 理想情况下应在__getattribute__访问self.prof_type、self.text_width、self.self_salt、self.hash_key时重置
+        greenlist_ids = self._get_greenlist_ids(torch.as_tensor(prefix, device=self.device))
+        return True if target in greenlist_ids else False
+
+    def _score_ngrams_in_passage(self, input_ids: torch.Tensor):
+        """核心功能是收集输入中的所有ngram并计算其水印"""
+        if len(input_ids) - self.context_width < 1:
+            raise ValueError(
+                f"Must have at least {1} token to score after "
+                f"the first min_prefix_len={self.context_width} tokens required by the seeding scheme."
+            )
+
+        # 计算文章中所有ngrams上下文的分数:
+        token_ngram_generator = ngrams(input_ids.cpu().tolist(), self.context_width + 1 - self.self_salt)
+        frequencies_table = collections.Counter(token_ngram_generator)
+        ngram_to_watermark_lookup = {}
+        for idx, ngram_example in enumerate(frequencies_table.keys()):
+            prefix = ngram_example if self.self_salt else ngram_example[:-1]
+            target = ngram_example[-1]
+            ngram_to_watermark_lookup[ngram_example] = self._get_ngram_score_cached(prefix, target)
+
+        return ngram_to_watermark_lookup, frequencies_table
+
+    def _get_green_at_T_booleans(self, input_ids, ngram_to_watermark_lookup) -> tuple[torch.Tensor]:
+        """生成基于每个标记的绿色与红色二值列表，一个忽略重复n-gram的独立列表，以及一个用于在两种表示法之间转换的偏移量列表：
+    green_token_mask = green_token_mask_unique[offsets]，除了在所有会被计算为重复的位置之外
+    """
+
+        green_token_mask, green_token_mask_unique, offsets = [], [], []
+        used_ngrams = {}
+        unique_ngram_idx = 0
+        ngram_examples = ngrams(input_ids.cpu().tolist(), self.context_width + 1 - self.self_salt)
+
+        for idx, ngram_example in enumerate(ngram_examples):
+            green_token_mask.append(ngram_to_watermark_lookup[ngram_example])
+            if self.ignore_repeated_ngrams:
+                if ngram_example in used_ngrams:
+                    pass
+                else:
+                    used_ngrams[ngram_example] = True
+                    unique_ngram_idx += 1
+                    green_token_mask_unique.append(ngram_to_watermark_lookup[ngram_example])
+            else:
+                green_token_mask_unique.append(ngram_to_watermark_lookup[ngram_example])
+                unique_ngram_idx += 1
+            offsets.append(unique_ngram_idx - 1)
+        return (
+            torch.tensor(green_token_mask),
+            torch.tensor(green_token_mask_unique),
+            torch.tensor(offsets),
+        )
+
+    def _score_sequence(
+        self,
+        input_ids: torch.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_z_at_T: bool = True,
+        return_p_value: bool = True,
+    ):
+        ngram_to_watermark_lookup, frequencies_table = self._score_ngrams_in_passage(input_ids)
+        green_token_mask, green_unique, offsets = self._get_green_at_T_booleans(input_ids, ngram_to_watermark_lookup)
+
+        # 把所有ngrams的分数加起来
+        if self.ignore_repeated_ngrams:
+            # 一个方法，只对每个唯一的n-gram计算一次绿色/红色命中。
+            # 新的总标记评分数（T）变为唯一n-gram的数量。
+            # 我们遍历输入中的所有唯一的标记n-gram，计算每个上下文诱导的绿名单，
+            # 然后检查最后一个标记是否落在该绿名单中。
+
+            num_tokens_scored = len(frequencies_table.keys())
+            green_token_count = sum(ngram_to_watermark_lookup.values())
+        else:
+            num_tokens_scored = sum(frequencies_table.values())
+            assert num_tokens_scored == len(input_ids) - self.context_width + self.self_salt
+            green_token_count = sum(freq * outcome for freq, outcome in zip(frequencies_table.values(), ngram_to_watermark_lookup.values()))
+        assert green_token_count == green_unique.sum()
+
+        # HF-style 字典
+        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.tolist()))
+        if return_z_at_T:
+            # Score z_at_T:
+            sizes = torch.arange(1, len(green_unique) + 1)
+            seq_z_score_enum = torch.cumsum(green_unique, dim=0) - self.gamma * sizes
+            seq_z_score_denom = torch.sqrt(sizes * self.gamma * (1 - self.gamma))
+            z_score_at_effective_T = seq_z_score_enum / seq_z_score_denom
+            z_score_at_T = z_score_at_effective_T[offsets]
+            assert torch.isclose(z_score_at_T[-1], torch.tensor(z_score))
+
+            score_dict.update(dict(z_score_at_T=z_score_at_T))
+
+        return score_dict
+
+    def _score_windows_impl_batched(
+        self,
+        input_ids: torch.Tensor,
+        window_size: str,
+        window_stride: int = 1,
+    ):
+        # 实现细节：
+            # 1) --ignore_repeated_ngrams 选项被全局应用，然后在减少的二值向量上应用窗口化处理。
+            #    这只是实现的一种方式，另一种方法是在每个窗口内忽略bigram（这可能更难并行化处理）。
+            # 2) 这些窗口在绿色/红色命中的二值向量上，独立于 context_width，与 Kezhi 的第一个实现不同。
+            # 3) 由于窗口化的处理，这个实现得到的 z-分数不能直接转换为 p-值，并且应该只用作对选定 FPR 进行校准的 ROC 图的标签。
+            #    由于多次假设测试，整体得分将被提高。
+            #    naive_count_correction=True 是对这个问题的一种部分解决方法。
+
+
+        ngram_to_watermark_lookup, frequencies_table = self._score_ngrams_in_passage(input_ids)
+        green_mask, green_ids, offsets = self._get_green_at_T_booleans(input_ids, ngram_to_watermark_lookup)
+        len_full_context = len(green_ids)
+
+        partial_sum_id_table = torch.cumsum(green_ids, dim=0)
+
+        if window_size == "max":
+            # 可以稍后启动，小窗口无法产生足够的能量
+            # solve (T * Spike_Entropy - g * T) / sqrt(T * g * (1 - g)) = z_thresh for T
+            sizes = range(1, len_full_context)
+        else:
+            sizes = [int(x) for x in window_size.split(",") if len(x) > 0]
+
+        z_score_max_per_window = torch.zeros(len(sizes))
+        cumulative_eff_z_score = torch.zeros(len_full_context)
+        s = window_stride
+
+        window_fits = False
+        for idx, size in enumerate(sizes):
+            if size <= len_full_context:
+                # 并行计算窗口内所有位置的hit：
+                window_score = torch.zeros(len_full_context - size + 1, dtype=torch.long)
+                # 包括第0个窗口
+                window_score[0] = partial_sum_id_table[size - 1]
+                # 从1号开始的所有其他窗口：
+                window_score[1:] = partial_sum_id_table[size::s] - partial_sum_id_table[:-size:s]
+
+                # 现在计算批处理的z_scores
+                batched_z_score_enum = window_score - self.gamma * size
+                z_score_denom = sqrt(size * self.gamma * (1 - self.gamma))
+                batched_z_score = batched_z_score_enum / z_score_denom
+
+                # 找到最大的hit
+                maximal_z_score = batched_z_score.max()
+                z_score_max_per_window[idx] = maximal_z_score
+
+                z_score_at_effective_T = torch.cummax(batched_z_score, dim=0)[0]
+                cumulative_eff_z_score[size::s] = torch.maximum(cumulative_eff_z_score[size::s], z_score_at_effective_T[:-1])
+                window_fits = True  # 成功计算所有大小的窗口
+
+        if not window_fits:
+            raise ValueError(
+                f"Could not find a fitting window with window sizes {window_size} for (effective) context length {len_full_context}."
+            )
+
+        # 计算最佳窗口大小和z得分
+        cumulative_z_score = cumulative_eff_z_score[offsets]
+        optimal_z, optimal_window_size_idx = z_score_max_per_window.max(dim=0)
+        optimal_window_size = sizes[optimal_window_size_idx]
+        return (
+            optimal_z,
+            optimal_window_size,
+            z_score_max_per_window,
+            cumulative_z_score,
+            green_mask,
+        )
+
+    def _score_sequence_window(
+        self,
+        input_ids: torch.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_z_at_T: bool = True,
+        return_p_value: bool = True,
+        window_size: str = None,
+        window_stride: int = 1,
+    ):
+        (
+            optimal_z,
+            optimal_window_size,
+            _,
+            z_score_at_T,
+            green_mask,
+        ) = self._score_windows_impl_batched(input_ids, window_size, window_stride)
+
+        # HF-style 字典
+        score_dict = dict()
+        if return_num_tokens_scored:
+            score_dict.update(dict(num_tokens_scored=optimal_window_size))
+
+        denom = sqrt(optimal_window_size * self.gamma * (1 - self.gamma))
+        green_token_count = int(optimal_z * denom + self.gamma * optimal_window_size)
+        green_fraction = green_token_count / optimal_window_size
+        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_fraction))
+        if return_z_score:
+            score_dict.update(dict(z_score=optimal_z))
+        if return_z_at_T:
+            score_dict.update(dict(z_score_at_T=z_score_at_T))
+        if return_p_value:
+            z_score = score_dict.get("z_score", optimal_z)
+            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_mask.tolist()))
+
+        return score_dict
+
+    def detect(
+        self,
+        text: str = None,
+        tokenized_text: list[int] = None,
+        window_size: str = None,
+        window_stride: int = None,
+        return_prediction: bool = True,
+        return_scores: bool = True,
+        z_threshold: float = None,
+        convert_to_float: bool = False,
+        **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  # 返回阳性检测的 "confidence":=1-p
+
+        # 对文本运行可选的normalizers
+        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:
+            # 尝试从一开始就删除bos_tok（如果它在那里的话）
+            if (self.tokenizer is not None) and (tokenized_text[0] == self.tokenizer.bos_token_id):
+                tokenized_text = tokenized_text[1:]
+
+        # 调用score方法
+        output_dict = {}
+
+        if window_size is not None:
+            score_dict = self._score_sequence_window(
+                tokenized_text,
+                window_size=window_size,
+                window_stride=window_stride,
+                **kwargs,
+            )
+            output_dict.update(score_dict)
+        else:
+            score_dict = self._score_sequence(tokenized_text, **kwargs)
+        if return_scores:
+            output_dict.update(score_dict)
+        # 如果通过return_prediction，则执行假设检验并返回结果
+        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"]
+
+        # 如果需要的话，将任何数值转换为浮点值
+        if convert_to_float:
+            for key, value in output_dict.items():
+                if isinstance(value, int):
+                    output_dict[key] = float(value)
+
+        return output_dict
+
+
+
+
+
+def ngrams(sequence, n, pad_left=False, pad_right=False, pad_symbol=None):
+    sequence = iter(sequence)
+    if pad_left:
+        sequence = chain((pad_symbol,) * (n - 1), sequence)
+    if pad_right:
+        sequence = chain(sequence, (pad_symbol,) * (n - 1))
+    iterables = tee(sequence, n)
+
+    for i, sub_iterable in enumerate(iterables):  # For each window,
+        for _ in range(i):  # iterate through every order of ngrams
+            next(sub_iterable, None)  # generate the ngrams within the window.
+    return zip(*iterables)  # Unpack and flattens the iterables.

homoglyphs.py

@@ -0,0 +1,249 @@
+
+from collections import defaultdict
+import json
+from itertools import product
+import os
+import unicodedata
+
+STRATEGY_LOAD = 1  # 加载类别
+STRATEGY_IGNORE = 2  # 对结果添加字符
+STRATEGY_REMOVE = 3  # 对结果移除字符
+
+ASCII_RANGE = range(128)
+
+
+CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
+DATA_LOCATION = os.path.join(CURRENT_DIR, "homoglyph_data")
+
+
+class Categories:
+
+
+    fpath = os.path.join(DATA_LOCATION, "categories.json")
+
+    @classmethod
+    def _get_ranges(cls, categories):
+
+        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):
+      
+        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)
+
+        # 尝试用unicodedata检测类别
+        try:
+            category = unicodedata.name(char).split()[0]
+        except (TypeError, ValueError):
+            pass
+        else:
+            if category in data["aliases"]:
+                return category
+
+        # 尝试从JSON文件中按范围检测类别
+        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必须由任何字母表初始化才能正确工作
+        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):
+        # 尝试检测语言
+        langs = Languages.detect(char)
+        if langs:
+            self.languages.update(langs)
+            alphabet = Languages.get_alphabet(langs)
+            self.alphabet.update(alphabet)
+        else:
+            # 尝试检测类别
+            category = Categories.detect(char)
+            if category is None:
+                return False
+            self.categories.add(category)
+            alphabet = Categories.get_alphabet([category])
+            self.alphabet.update(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 []
+
+        # 查找当前字符的替代字符
+        alt_chars = self.table.get(char, set())
+        if alt_chars:
+            # 为当前字符查找可选字符
+            alt_chars2 = [self.table.get(alt_char, set()) for alt_char in alt_chars]
+            # 合并所有备选方案
+            alt_chars.update(*alt_chars2)
+        # 将当前字符添加到备选项
+        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,199 @@
+"""文本基础规范化器，用于减轻针对水印的简单攻击。
+
+这个实现不太可能是所有可能的Unicode标准中的所有漏洞的完整列表，
+它代表了我们在撰写时的最佳努力。
+
+这些规范化器可以作为独立的规范化器使用。它们可以被制作成符合HF分词器标准的规范化器，
+但这将需要涉及tokenizers.NormalizedString的有限Rust接口。
+"""
+
+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:
+    """尝试检测同形字攻击并找到一致的标准形式。
+
+    这个函数是在ISO分类级别上进行的。也可以在语言级别上进行（参见注释掉的代码）。
+    """
+
+
+    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")
+        )  # 从文件中加载到此处的字母表
+
+        source_alphabet = hg.Categories.get_alphabet(all_categories)
+        restricted_table = homoglyph_table.get_restricted_table(
+            source_alphabet, homoglyph_table.alphabet
+        )  # 从文件中加载到此处的表
+        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:
+
+    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.
+
+            """
+
+            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"
+        )  # 删除所有剩余的不可打印字符
+        return text
+
+
+class TrueCaser:
+    """真大小写还原，是一种将文本还原为其原始大小写形式的大小写规范化处理。
+
+    这可以防御那些像 spOngBoB 那样随机大小写的攻击。
+
+    这里使用了简单的词性标注器。
+    """
+
+
+    uppercase_pos = ["PROPN"]  # 应使用大写字母命名POS
+
+    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,12 @@
+nltk
+scipy
+torch
+transformers
+tokenizers
+accelerate
+text-generation==0.3.1
+optimum
+auto-gptq
+cpm_kernels
+bitsandbytes
+gradio==3.50.2

watermark_processor.py

@@ -0,0 +1,336 @@
+from __future__ import annotations
+
+import collections
+from math import sqrt
+
+import scipy.stats
+import torch
+from nltk.util import ngrams
+from tokenizers import Tokenizer
+from torch import Tensor
+from transformers import LogitsProcessor
+
+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",
+            hash_key: int = 15485863,  # 只需要一个大素数就可以创建一个具有足够位宽的rng种子
+            extra_salt: int = 0,
+            select_green_tokens: bool = True,
+    ):
+
+        # 水印参数
+        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.extra_salt = extra_salt
+        self.select_green_tokens = select_green_tokens
+
+    def _seed_rng(self, input_ids: torch.LongTensor, seeding_scheme: str = None) -> None:
+        # 可以选择覆盖种子设定方案，但默认情况下使用实例属性
+        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 + self.extra_salt)
+        else:
+            raise NotImplementedError(f"Unexpected seeding_scheme: {seeding_scheme}")
+        return
+
+    def _get_greenlist_ids(self, input_ids: torch.LongTensor) -> list[int]:
+
+        self._seed_rng(input_ids)
+
+        greenlist_size = int(self.vocab_size * self.gamma)
+
+        if input_ids.device != 'cpu':
+            # 为了确保能在不同设备上复现，这里的随机数生成都用cpu
+            vocab_permutation = torch.randperm(self.vocab_size, device='cpu', generator=self.rng)
+            vocab_permutation = vocab_permutation.to(input_ids.device)
+
+        else:
+            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:  # 从红色中挑选绿色
+            greenlist_ids = vocab_permutation[(self.vocab_size - greenlist_size):]
+        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:
+        if self.rng is None:
+            # self.rng = torch.Generator(device=input_ids.device)
+            self.rng = torch.Generator(device='cpu')
+
+
+        # 注意：理想情况下应该去掉这个批处理循环，但目前，
+        # 种子和分区操作还没有实现向量化，因此
+        # 批处理中的每个序列需要单独处理。
+
+        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)
+
+        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)
+        self.rng = torch.Generator(device='cpu')
+
+        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是指绿色token的数量，T是token的总数
+        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:
+            # 一个方法，只对每个唯一的bigram计算一次绿色/红色命中。
+            # 新的总标记评分数（T）变为唯一bigram的数量。
+            # 我们遍历输入中的所有唯一的标记bigram，计算每个bigram的第一个标记诱导的绿名单，
+            # 然后检查第二个标记是否落在该绿名单中。
+
+            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."))
+            # 标准方法
+            # 由于我们通常至少需要1个token（对于最简单的方案）
+            # 我们从最小数量的token开始迭代token序列，作为种子方案的第一个前缀，
+            # 在每一步中，计算当前前缀诱导的绿名单，
+            # 并检查当前token是否落在绿名单中。
+
+            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  # 返回阳性检测的"confidence":=1-p
+
+        # 运行可选的normalizers
+        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:
+            # 尝试一开始就删除bos_tok（如果它在那里的话）
+            if (self.tokenizer is not None) and (tokenized_text[0] == self.tokenizer.bos_token_id):
+                tokenized_text = tokenized_text[1:]
+
+        # 调用score方法
+        output_dict = {}
+        score_dict = self._score_sequence(tokenized_text, **kwargs)
+        if return_scores:
+            output_dict.update(score_dict)
+        # 如果通过return_prediction，则执行假设检验并返回结果
+        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
+
+
+if __name__ == "__main__":
+    from transformers import AutoTokenizer, AutoModelForCausalLM, LogitsProcessorList
+
+    tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen1.5-0.5B-Chat")
+    model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen1.5-0.5B-Chat")
+
+    watermark_processor = WatermarkLogitsProcessor(vocab=list(tokenizer.get_vocab().values()),
+                                                   gamma=0.5,
+                                                   delta=2,
+                                                   seeding_scheme="simple_1",
+                                                   extra_salt=0,
+                                                   select_green_tokens=True)
+
+    messages = [
+        # {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "content": "讲一段明代的历史"}
+    ]
+
+    tokenized_input = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True
+    )
+
+    tokd_input = tokenizer([tokenized_input], return_tensors="pt", truncation=True, add_special_tokens=False,
+                           max_length=2000).to(model.device)
+
+    logits_processor = LogitsProcessorList([watermark_processor])
+
+    output = model.generate(
+        tokd_input.input_ids,
+        max_new_tokens=500,
+        logits_processor=logits_processor,
+        do_sample=True,
+        temperature=0.7,
+    )
+
+    print(tokenizer.decode(output[0]))
+
+    watermark_detector = WatermarkDetector(vocab=list(tokenizer.get_vocab().values()),
+                                           gamma=0.5,
+                                           seeding_scheme="simple_1",
+                                           extra_salt=0,
+                                           device=torch.device("cpu"),
+                                           tokenizer=tokenizer,
+                                           z_threshold=4,
+                                           # normalizers='',
+                                           ignore_repeated_bigrams=False,
+                                           select_green_tokens=True)
+
+    print(watermark_detector.detect(tokenizer.decode(output[0]), return_prediction=True, return_scores=True,
+                                    z_threshold=4))
+
+    # print(watermark_detector.detect("抱歉，作为人工智能语言模型，我无法提供有关希腊历史的信息。我的目的是为用户提供有用的和有用的回答，而不仅仅是提供错误的信息。请告诉我您想要了解的是什么内容。 ", return_prediction=True, return_scores=True, z_threshold=4))
+    #
+    # print(watermark_detector.detect("明朝是中国历史上一个重要的朝代，从1368年至1542年，中国经历了从宋朝的衰败到明朝的兴盛，这个时期的朝代特征鲜明，政治、经济和社会都取得了显著的进步。 政治方面：明朝的君主制度比较完善，以皇权为核心，实行中央集权。此外，明朝还实行了科举考试制度，选拔了许多优秀人才，使得社会���气更加开放。 经济上：明朝的经济实力非常强大，尤其是手工业和农业发展迅速。明朝的海上贸易也非常发达，被誉为“海路不穷”。另外，明朝还通过派遣郑和等船队出使西洋，传播中国文化，扩大对外交流。 社会上：明朝的社会结构相对稳定，人们生活水平不断提高。然而，这个时期也存在一些问题，如农民起义、土地兼并等。 文化方面：明朝的文化非常丰富多样，有诗词歌赋、绘画、建筑等众多艺术形式。此外，明人的文学创作也非常出色，如《西游记》、《红楼梦》等经典作品。 总之，明朝是中国历史上的一个重要阶段，它的繁荣与衰败、创新与发展深深地影响了后世的人民和社会的发展。 ", return_prediction=True, return_scores=True, z_threshold=4))