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