lm-watermarking / watermark_processor.py
jwkirchenbauer's picture
rm fixme
6f4883b
raw
history blame
12.3 kB
# 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