lm-watermarking-main/experiments/watermark.py

# micro lib to implement the watermarking extensions to LM generation
# as well as utils for eval/validaiton

from typing import List, Optional, Callable

import time
import random
import math
import torch
import numpy as np

from torch import Tensor
from tokenizers import Tokenizer
from transformers import LogitsProcessor, LogitsProcessorList, set_seed


def tokenize_and_truncate(example: dict,
                          completion_length: int = None,
                          prompt_length: int = None,
                          hf_model_name: str = None,
                          tokenizer = None,
                          model_max_seq_len: int = 4096):
    """take hf dataset entry and preprocess it for completion by a model"""
    assert hf_model_name is not None, "need model name to know whether to adjust wrt special tokens"
    assert "text" in example, "expects 'text' field to be present"
    # tokenize
    inputs = tokenizer.encode(example["text"], return_tensors="pt", truncation=True, max_length=model_max_seq_len)
    example.update({"untruncated_inputs": inputs})

    if (completion_length is not None) and (prompt_length is None):
        # leave at least one token as prefix # FIXME I think plus 1 since 0 is start tok
        slice_length = min(inputs.shape[1]-1, completion_length)
    elif (prompt_length is not None) and (completion_length is None):
        desired_comp_len = (inputs.shape[1]-1) - prompt_length
        slice_length = desired_comp_len if desired_comp_len > 0 else 0
    else:
        raise ValueError((f"Can only tokenize and truncate based on either the desired prompt length or desired completion length,",
                          f" but got completion_length:{completion_length},prompt_length:{prompt_length}"))

    # truncate
    inputs = inputs[:,:inputs.shape[1]-slice_length]
    # logic depending on special tokens for the model
    if "t5" in hf_model_name or "T0" in hf_model_name: 
        inputs[0,-1] = 1
    # else: pass
    example.update({"inputs": inputs})
    return example


class BlacklistLogitsProcessor(LogitsProcessor):
    """
    [`LogitsProcessor`] that enforces that specified sequences will never be sampled.

    Args:
        bad_words_ids (`List[List[int]]`):
            List of list of token ids that are not allowed to be generated. In order to get the token ids of the words
            that should not appear in the generated text, use `tokenizer(bad_words, add_prefix_space=True,
            add_special_tokens=False).input_ids`.
        eos_token_id (`int`):
            The id of the *end-of-sequence* token.
    """

    def __init__(self, 
                bad_words_ids: List[List[int]], 
                eos_token_id: int,
                vocab: list[int], 
                vocab_size: int,
                bl_proportion: float=0.5,
                bl_logit_bias: float=1.0,
                bl_type: str = "hard", # "soft"
                initial_seed: int=None, 
                dynamic_seed: str=None, # "initial", "markov_1", None
                store_bl_ids: bool=True,
                store_spike_ents: bool = False,
                noop_blacklist: bool = False,
                ):
        
        self.vocab = vocab
        self.vocab_size = vocab_size
        self.bl_proportion = bl_proportion
        self.bl_logit_bias = bl_logit_bias
        self.bl_type = bl_type

        if initial_seed is None: 
            self.initial_seed = None
            assert dynamic_seed != "initial"
        else:
            random.seed(initial_seed)
            self.initial_seed = initial_seed
            
        self.dynamic_seed = dynamic_seed

        self.eos_token_id = eos_token_id
        # self.bad_words_id_length_1 = self._prepare_bad_words(bad_words_ids)

        self.bad_words_mask: Optional[torch.LongTensor] = None

        self.store_bl_ids = store_bl_ids
        self.bl_ids = None

        self.store_spike_ents = store_spike_ents
        self.spike_entropies = None

        # hack to replace this with an approximation of infinite bias
        # so that the expectation coefficient will come out to 1.0
        if self.bl_type == "hard":
            self.bl_logit_bias = 10000 # FIXME to a value that is actually close to the largest soft watermark used

        alpha = torch.exp(torch.tensor(self.bl_logit_bias)).item()
        # gamma = self.bl_proportion
        gamma = 1.0-self.bl_proportion
        self.alpha = alpha
        self.gamma = gamma

        self.z_value = ((1-gamma)*(alpha-1))/(1-gamma+(alpha*gamma))
        self.expected_wl_coef = (gamma*alpha)/(1-gamma+(alpha*gamma))

        # catch for overflow when bias is "infinite"
        if self.alpha == torch.inf:
            self.z_value = 1.0
            self.expected_wl_coef = 1.0

        self.noop_blacklist = noop_blacklist
        if self.noop_blacklist: print(f"Blacklist processor for accounting only, no rescoring of logits")

        self.g_cuda = None
        self.large_prime = 15485863

    @property
    def blacklisted_ids(self):
        assert self.store_bl_ids, "Need to instantiate processor with `store_bl_ids` to be able to retrieve them later"
        # flatten the each indexes blacklist
        l_of_bl_ids = [[] for _ in range(len(self.bl_ids))]
        for b_idx, batch in enumerate(self.bl_ids):
            for l_of_l, seed in batch:
                bl_ids = [l[0] for l in l_of_l] # this was the main line, maybe unnecessary now?
                l_of_bl_ids[b_idx].append((bl_ids,seed))
        return l_of_bl_ids

    def get_and_clear_stored_bl_ids(self):
        old_bl_ids = self.bl_ids
        self.bl_ids = None
        return old_bl_ids

    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):
        # precomputed z value in init
        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 __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:

        self.bad_words_id_length_1 = [None for _ in range(input_ids.shape[0])]

        if self.g_cuda is None:
            self.g_cuda = torch.Generator(device=input_ids.device)

        for b_idx in range(input_ids.shape[0]):
            if self.dynamic_seed == "initial":
                self.g_cuda.manual_seed(self.large_prime*self.initial_seed)
            elif self.dynamic_seed == "markov_1":
                self.g_cuda.manual_seed(self.large_prime*input_ids[b_idx][-1].item())
            elif self.dynamic_seed is None:
                # let the rng evolve naturally - this is not a realistic setting
                pass

            bl_ct = int(self.vocab_size*self.bl_proportion)
            blacklist_ids = torch.randperm(self.vocab_size, device=input_ids.device, generator=self.g_cuda)[:bl_ct] # ty Yuxin :]

            if self.store_bl_ids: 
                if self.bl_ids is None: self.bl_ids = [[] for _ in range(input_ids.shape[0])]
                self.bl_ids[b_idx].append((blacklist_ids,input_ids.tolist()[b_idx][-1]))

            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]))
            
            # self.bad_words_id_length_1[b_idx] = self._prepare_bad_words(blacklist_ids)
            # this logic may not really be necessary for our usecase
            self.bad_words_id_length_1[b_idx] = blacklist_ids
        
        if not self.noop_blacklist:
            self.bad_words_mask = self._calc_curr_bad_word_mask(scores)
            scores = self._set_scores_to_inf_for_banned_tokens(scores)

        return scores

    def _prepare_bad_words(self, bad_words_ids: List[List[int]]) -> list[int]:
        bad_words_ids = list(filter(lambda bad_token_seq: bad_token_seq != [self.eos_token_id], bad_words_ids))
        return bad_words_ids
        # used to have more logic, not used now

    def _calc_curr_bad_word_mask(self, scores: torch.FloatTensor) -> torch.BoolTensor:
        bad_words_mask = torch.zeros_like(scores)
        for b_idx in range(len(self.bad_words_id_length_1)):
            bad_words_mask[b_idx][self.bad_words_id_length_1[b_idx]] = 1
        final_mask = bad_words_mask.bool()
        return final_mask

    def _set_scores_to_inf_for_banned_tokens(
        self, scores: torch.Tensor
    ) -> torch.Tensor:
        """
        Modifies the scores in place by setting the banned token positions to `-inf`. Banned token is expected to be a
        list of list of banned tokens to ban in the format [[batch index, vocabulary position],...

        Args:
            scores: logits distribution of shape (batch size, vocabulary size)
            banned_tokens: list of list of tokens to ban of length (batch_size)
            # NOTE^^ Omitted logic for dynamic mask based on multi-token ban words
        """
        if self.bl_type == "hard":
            scores = scores.masked_fill(self.bad_words_mask, -float("inf"))
        elif self.bl_type == "soft":
            whitelist_mask = torch.logical_not(self.bad_words_mask)
            blacklist_mask = self.bad_words_mask
            scores[whitelist_mask] = scores[whitelist_mask] + self.bl_logit_bias
            # scores[blacklist_mask] = scores[blacklist_mask] - self.bl_logit_bias  # additive only
        else:
            raise NotImplementedError(f"unrecognized bl type {self.bl_type}!")          
        return scores


def score_sequence(inputs: Tensor = None, 
                   outputs: Tensor = None, 
                   tokenizer: Tokenizer = None,
                #    logits_processor: LogitsProcessor = None,
                   initial_seed: int = None,
                   dynamic_seed: str = None,
                   bl_proportion: float = None,
                   use_cuda: bool = True,
                   record_hits: bool = False,
                   debug: bool = True,
                #    trim_tokens: int = 2,
                ):

    assert  (inputs is not None) and \
            (outputs is not None) and \
            (tokenizer is not None),"output tensor, tokenizer, and bl params req'd"
            # (logits_processor is not None), 

    vocabulary = list(tokenizer.get_vocab().values())
    vocab_size = len(vocabulary)

    model_generations = outputs.tolist()[0] # these are tensors unpack once for speed
    # toks_generated = model_generations[num_orig_input_tokens:]
    toks_generated = model_generations
    num_toks_generated = len(toks_generated)

    # num_toks_to_trim = trim_tokens*2
    # if (num_toks_generated-num_toks_to_trim > 0) == False: 
    #     return -1, -1
        
    # assert num_toks_generated > num_toks_to_trim, f"Need more than {num_toks_to_trim} toks total since we trim start and end a bit."

    # toks_generated = toks_generated[trim_tokens:-trim_tokens]

    if initial_seed is not None:
        random.seed(initial_seed)
    
    device = (torch.device("cuda") if use_cuda else torch.device("cpu"))
    g_cuda = torch.Generator(device=device)
    large_prime = 15485863

    bl_hits, hit_list = 0, []

    prev_token = inputs[0][-1].item()
    # prev_token = toks_generated[0] # haven't decided whether this edge effect matters

    # for idx,tok_gend in enumerate(toks_generated[1:]):
    for idx,tok_gend in enumerate(toks_generated):

        # prev_token = model_generations[num_orig_input_tokens+idx-1]
        
        if dynamic_seed == "initial":
            g_cuda.manual_seed(large_prime*initial_seed)
        elif dynamic_seed == "markov_1":
            g_cuda.manual_seed(large_prime*prev_token)
        elif dynamic_seed is None:
            # let the rng evolve naturally - this is not a realistic setting
            pass

        bl_ct = int(vocab_size*bl_proportion)
        posthoc_blacklist = torch.randperm(vocab_size, device=device, generator=g_cuda)[:bl_ct] # ty Yuxin :]

        tok_in_ph_bl = tok_gend in posthoc_blacklist
        if tok_in_ph_bl:
            bl_hits += 1
            hit_list.append(True)
        else:
            hit_list.append(False)

        if debug: 
            decoded_token = tokenizer.decode(tok_gend, skip_special_tokens=True)
            print(f"Token generated: '{decoded_token}' was in the blacklist {tok_in_ph_bl}")
        
        prev_token = tok_gend

    if debug: 
        print(f"wl hits / num tokens : {num_toks_generated-bl_hits}/{num_toks_generated} = {(num_toks_generated-bl_hits)/num_toks_generated:.02f}")
        print(f"bl hits / num tokens : {bl_hits}/{num_toks_generated} = {bl_hits/num_toks_generated:.02f}")

    if record_hits:
        return bl_hits, num_toks_generated, hit_list
    # bl_fraction = bl_hits/num_toks_generated
    return bl_hits, num_toks_generated


def tokenize_for_generation(example: dict, 
                            idx: int,
                            max_new_tokens: int=None,
                            min_prompt_tokens: int=None,
                            hf_model_name : str=None,
                            tokenizer: Tokenizer=None,
                            model: torch.nn.Module=None):

    # preprocessing, generation & scoring
    assert isinstance(example, dict), "Expect no batch dimension currently!"

    # preprocess for model generation/completion
    example = tokenize_and_truncate(example, 
                                    completion_length=max_new_tokens, 
                                    prompt_length=min_prompt_tokens,
                                    hf_model_name=hf_model_name,
                                    tokenizer=tokenizer,
                                    # model_max_seq_len=model.config.max_position_embeddings)
                                    model_max_seq_len=None)
    inputs = example["inputs"]
    # for calculating the baseline violation rate across the "gold" completion
    untruncated_inputs = example["untruncated_inputs"]

    # decode the preprocessed input to store for audit
    re_decoded_input = tokenizer.batch_decode(inputs, skip_special_tokens=True)[0]
    example.update({"truncated_input":re_decoded_input})

    # also decode the original suffix of the input for audit as the baseline
    decoded_untruncated_input = tokenizer.batch_decode(untruncated_inputs, skip_special_tokens=True)[0]
    example.update({"baseline_completion":decoded_untruncated_input.replace(re_decoded_input,"")})

    example.update({
        "orig_sample_length"            : untruncated_inputs.shape[1],
        "prompt_length"                 : inputs.shape[1],
        "real_completion_length"        : untruncated_inputs.shape[1] - inputs.shape[1],
    })
    return example


def generate_completions(example: dict, 
                        idx: int,
                        max_new_tokens: int=None,
                        hf_model_name : str=None,
                        tokenizer: Tokenizer=None,
                        model: torch.nn.Module=None,
                        no_bl_partial: Callable=None,
                        w_bl_partial: Callable=None,
                        # return_logits: bool=False,
                        bl_processor_list: LogitsProcessorList=None):

    # preprocessing, generation & scoring
    assert isinstance(example, dict), "Expect no batch dimension currently!"

    # # preprocess for model generation/completion
    # example = tokenize_and_truncate(example, 
    #                                 completion_length=max_new_tokens, 
    #                                 hf_model_name=hf_model_name,
    #                                 tokenizer=tokenizer,
    #                                 model_max_seq_len=model.config.max_position_embeddings)
    # inputs = example["inputs"]
    # # for calculating the baseline violation rate across the "gold" completion
    # untruncated_inputs = example["untruncated_inputs"]

    # # decode the preprocessed input to store for audit
    # re_decoded_input = tokenizer.batch_decode(inputs, skip_special_tokens=True)[0]
    # example.update({"truncated_input":re_decoded_input})

    # # also decode the original suffix of the input for audit as the baseline
    # decoded_untruncated_input = tokenizer.batch_decode(untruncated_inputs, skip_special_tokens=True)[0]
    # example.update({"baseline_completion":decoded_untruncated_input.replace(re_decoded_input,"")})

    inputs = example["inputs"]
    re_decoded_input = example["truncated_input"]

    # call the vanilla and watermarked generation function wrappers with the preprocessed inputs
    with torch.no_grad():

        samples_taken = 0
        max_retries = 10
        success = False
        while (success is False) and (samples_taken < max_retries):
            samples_taken += 1
            
            # set_seed(1234) # debugging the error when using sampling # leaving this off for now

            start_generation = time.time()
            outputs_no_bl = no_bl_partial(inputs.to(model.device))
            example["no_bl_gen_time"] = time.time() - start_generation

            # set_seed(1234) # debugging the error when using sampling

            start_generation = time.time()
            outputs_w_bl = w_bl_partial(inputs.to(model.device))
            example["w_bl_gen_time"] = time.time() - start_generation

            # if return_logits:
            #     output_no_bl_dict = outputs_no_bl
            #     logits_no_bl = output_no_bl_dict.scores
            #     outputs_no_bl = output_no_bl_dict.sequences
            #     example["logits_no_bl"] = logits_no_bl

            #     output_w_bl_dict = outputs_w_bl
            #     logits_w_bl = output_w_bl_dict.scores
            #     outputs_w_bl = output_w_bl_dict.sequences
            #     example["logits_w_bl"] = logits_w_bl


            if bl_processor_list:
                if bl_processor_list[0].bl_ids is not None:
                    example["bl_ids"] = bl_processor_list[0].get_and_clear_stored_bl_ids()
                if bl_processor_list[0].spike_entropies is not None:
                    example["spike_entropies"] = bl_processor_list[0].get_and_clear_stored_spike_ents()

            try: 
                # decode and store the new generations for auditing
                no_bl_decoded_output = tokenizer.batch_decode(outputs_no_bl, skip_special_tokens=True)[0]
                example.update({"no_bl_output":no_bl_decoded_output.replace(re_decoded_input,"")})

                w_bl_decoded_output = tokenizer.batch_decode(outputs_w_bl, skip_special_tokens=True)[0]
                example.update({"w_bl_output":w_bl_decoded_output.replace(re_decoded_input,"")})
                
                success = True

            except:
                # log what happened
                print(f"Error while trying to decode the outputs of the model...")
                if samples_taken == 1:
                    print(f"truncated_input: {inputs.tolist()}")
                print(f"Result of attempt {samples_taken}")
                print(f"shape outputs_no_bl: {outputs_no_bl.shape}")
                no_bl_toks = outputs_no_bl.tolist()[0]
                print(f"outputs_no_bl: {no_bl_toks}")
                print(f"outputs_no_bl min: {min(no_bl_toks)}")
                print(f"outputs_no_bl max: {max(no_bl_toks)}")

                print(f"shape outputs_w_bl: {outputs_w_bl.shape}")
                w_bl_toks = outputs_w_bl.tolist()[0]
                print(f"outputs_w_bl: {w_bl_toks}")
                print(f"outputs_w_bl min: {min(w_bl_toks)}")
                print(f"outputs_w_bl max: {max(w_bl_toks)}")
        
        if success is False:
            print(f"Unable to get both a no_bl and w_bl output that were decodeable after {samples_taken} tries, returning empty strings.")
            example.update({"no_bl_output":""})
            example.update({"w_bl_output":""})
            if bl_processor_list:
                if bl_processor_list[0].bl_ids is not None:
                    example["bl_ids"] = []
                if bl_processor_list[0].spike_entropies is not None:
                    example["spike_entropies"] = []

    # Able to get lengths in here by checking 
    # truncated input shape versus the output shape

    example.update({
        # "baseline_num_tokens_generated" : untruncated_inputs.shape[1] - inputs.shape[1], # want this earlier now
        "no_bl_num_tokens_generated"    : outputs_no_bl.shape[1] - inputs.shape[1],
        "w_bl_num_tokens_generated"     : outputs_w_bl.shape[1] - inputs.shape[1]
    })
    example.update({
        "no_bl_sec_per_tok"             : example["no_bl_gen_time"]/example["no_bl_num_tokens_generated"],
        "no_bl_tok_per_sec"             : example["no_bl_num_tokens_generated"]/example["no_bl_gen_time"],
        "w_bl_sec_per_tok"             : example["w_bl_gen_time"]/example["w_bl_num_tokens_generated"],
        "w_bl_tok_per_sec"             : example["w_bl_num_tokens_generated"]/example["w_bl_gen_time"],
    })
    
    # now done externally because these persist outside this func
    # # remove any fields we don't need to keep
    # del example["inputs"]
    # del example["untruncated_inputs"]
    
    return example


def compute_bl_metrics(example: dict, 
                        idx: int,
                        hf_model_name: str = None,
                        tokenizer: Tokenizer=None,
                        initial_seed: int = None,
                        dynamic_seed: str = None,
                        bl_proportion: float = None,
                        use_cuda: bool = None,
                        record_hits: bool = False,
                        limit_output_tokens: int = 0):

    # if example["idx"] == 3: breakpoint()
    # okay need to catch an odd bug here and fix things
    baseline_before = example["baseline_completion"]
    example["baseline_completion"] = baseline_before.replace(example["truncated_input"][:-1],"")
    if example["baseline_completion"] != baseline_before:
        print("baseline input replacement bug occurred!")
    
    no_bl_before = example["no_bl_output"]
    example["no_bl_output"] = no_bl_before.replace(example["truncated_input"][:-1],"")
    if example["no_bl_output"] != no_bl_before:
        print("no_bl_output input replacement bug occurred!")
    
    w_bl_before = example["w_bl_output"]
    example["w_bl_output"] = w_bl_before.replace(example["truncated_input"][:-1],"")
    if example["w_bl_output"] != w_bl_before:
        print("w_bl_output input replacement bug occurred!")

    if ("w_bl_output_attacked" in example):
        w_bl_attacked_before = example["w_bl_output_attacked"]
        example["w_bl_output_attacked"] = w_bl_attacked_before.replace(example["truncated_input"][:-1],"")
        if example["w_bl_output_attacked"] != w_bl_attacked_before:
            print("w_bl_output_attacked input replacement bug occurred!")

    ##########

    # preprocess for model generation/completion
    inputs = tokenize_and_truncate({"text":example["truncated_input"]}, 
                                    completion_length=0, 
                                    hf_model_name=hf_model_name,
                                    tokenizer=tokenizer)["inputs"]

    baseline_outputs = tokenize_and_truncate({"text":example["baseline_completion"]}, 
                                        completion_length=0, 
                                        hf_model_name=hf_model_name,
                                        tokenizer=tokenizer)["inputs"][:,1:]

    no_bl_outputs = tokenize_and_truncate({"text":example["no_bl_output"]}, 
                                        completion_length=0, 
                                        hf_model_name=hf_model_name,
                                        tokenizer=tokenizer)["inputs"][:,1:]

    w_bl_outputs = tokenize_and_truncate({"text":example["w_bl_output"]}, 
                                        completion_length=0, 
                                        hf_model_name=hf_model_name,
                                        tokenizer=tokenizer)["inputs"][:,1:]
    if "w_bl_output_attacked" in example:
        w_bl_attacked_outputs = tokenize_and_truncate({"text":example["w_bl_output_attacked"]}, 
                                            completion_length=0, 
                                            hf_model_name=hf_model_name,
                                            tokenizer=tokenizer)["inputs"][:,1:]
    else:
        w_bl_attacked_outputs = None

    if limit_output_tokens > 0:
        example["orig_baseline_completion"] = example["baseline_completion"]
        example["orig_real_completion_length"] = example["real_completion_length"]
        baseline_outputs = baseline_outputs[:,:limit_output_tokens]
        example["real_completion_length"] = baseline_outputs.shape[1]
        example["baseline_completion"] = tokenizer.batch_decode(baseline_outputs, skip_special_tokens=True)[0]

        example["orig_no_bl_output"] = example["no_bl_output"]
        example["orig_no_bl_num_tokens_generated"] = example["no_bl_num_tokens_generated"]
        no_bl_outputs = no_bl_outputs[:,:limit_output_tokens]
        example["no_bl_num_tokens_generated"] = no_bl_outputs.shape[1]
        example["no_bl_output"] = tokenizer.batch_decode(no_bl_outputs, skip_special_tokens=True)[0]

        example["orig_w_bl_output"] = example["w_bl_output"]
        example["orig_w_bl_num_tokens_generated"] = example["w_bl_num_tokens_generated"]
        w_bl_outputs = w_bl_outputs[:,:limit_output_tokens]
        example["w_bl_num_tokens_generated"] = w_bl_outputs.shape[1]
        example["w_bl_output"] = tokenizer.batch_decode(w_bl_outputs, skip_special_tokens=True)[0]

        example["orig_spike_entropies"] = example["spike_entropies"]
        example["spike_entropies"] = [example["spike_entropies"][0][:limit_output_tokens]]

        if "w_bl_output_attacked" in example:
            # raise NotImplementedError("Havent thought what to do yet for this")
            example["orig_w_bl_output_attacked"] = example["w_bl_output_attacked"]
            # example["orig_w_bl_attacked_num_tokens_generated"] = example["w_bl_attacked_num_tokens_generated"]
            w_bl_attacked_outputs = w_bl_attacked_outputs[:,:limit_output_tokens]
            example["w_bl_attacked_num_tokens_generated"] = w_bl_attacked_outputs.shape[1]
            example["w_bl_output_attacked"] = tokenizer.batch_decode(w_bl_attacked_outputs, skip_special_tokens=True)[0]

    # score the 3 sequence completions/outputs wrt to bl hits
    result = score_sequence(inputs=inputs,
                            outputs=baseline_outputs, # <-- real text completions
                            initial_seed=initial_seed,
                            dynamic_seed=dynamic_seed,
                            bl_proportion=bl_proportion, 
                            tokenizer=tokenizer,
                            use_cuda=use_cuda,
                            record_hits=record_hits,
                            debug=False)
    if record_hits:
        bl_hits, num_toks_gend, hit_list = result
    else:
        bl_hits, num_toks_gend = result
    example.update({"baseline_num_toks_gend_eq_0":(num_toks_gend == 0)})
    # if num_toks_gend < 0.99*example["real_completion_length"]: breakpoint()
    # if len(hit_list) < 0.99*example["real_completion_length"]: breakpoint()

    if num_toks_gend == 0:
        # print("No tokens generated, odd, avoiding div by zero and returning -1's")
        wl_frac = -1
        bl_frac = -1
    else:
        wl_frac = (num_toks_gend-bl_hits)/num_toks_gend
        bl_frac = bl_hits/num_toks_gend
    baseline_stats = {
        "baseline_whitelist_fraction": wl_frac,
        "baseline_blacklist_fraction": bl_frac
    }
    example.update(baseline_stats)
    if record_hits: example.update({"baseline_hit_list":hit_list})
    
    result = score_sequence(inputs=inputs,
                                            outputs=no_bl_outputs, # <-- non-blacklisted version
                                            initial_seed=initial_seed,
                                            dynamic_seed=dynamic_seed,
                                            bl_proportion=bl_proportion, 
                                            tokenizer=tokenizer,
                                            record_hits=record_hits,
                                            debug=False)
    if record_hits:
        bl_hits, num_toks_gend, hit_list = result
    else:
        bl_hits, num_toks_gend = result
    example.update({"no_bl_num_toks_gend_eq_0":(num_toks_gend == 0)})
    # if num_toks_gend < 0.99*example["no_bl_num_tokens_generated"]: breakpoint()
    # if len(hit_list) < 0.99*example["no_bl_num_tokens_generated"]: breakpoint()

    if num_toks_gend == 0:
        # print("No tokens generated, odd, avoiding div by zero and returning -1's")
        wl_frac = -1
        bl_frac = -1
    else:
        wl_frac = (num_toks_gend-bl_hits)/num_toks_gend
        bl_frac = bl_hits/num_toks_gend
    no_bl_stats = {
        "no_bl_whitelist_fraction": wl_frac,
        "no_bl_blacklist_fraction": bl_frac
    }
    example.update(no_bl_stats)
    if record_hits: example.update({"no_bl_hit_list":hit_list})
    
    result = score_sequence(inputs=inputs,
                                            outputs=w_bl_outputs, # <-- blacklisted version
                                            initial_seed=initial_seed,
                                            dynamic_seed=dynamic_seed,
                                            bl_proportion=bl_proportion, 
                                            tokenizer=tokenizer,
                                            record_hits=record_hits,
                                            # breakpoint_on_hit=True, # banging head against wall
                                            debug=False)
    if record_hits:
        bl_hits, num_toks_gend, hit_list = result
    else:
        bl_hits, num_toks_gend = result
    example.update({"w_bl_num_toks_gend_eq_0":(num_toks_gend == 0)})
    # if num_toks_gend < 0.99*example["w_bl_num_tokens_generated"]: breakpoint()
    # if len(hit_list) < 0.99*example["w_bl_num_tokens_generated"]: breakpoint()

    if num_toks_gend == 0:
        # print("No tokens generated, odd, avoiding div by zero and returning -1's")
        wl_frac = -1
        bl_frac = -1
    else:
        wl_frac = (num_toks_gend-bl_hits)/num_toks_gend
        bl_frac = bl_hits/num_toks_gend
    w_bl_stats = {
        "w_bl_whitelist_fraction": wl_frac,
        "w_bl_blacklist_fraction": bl_frac
    }
    example.update(w_bl_stats)
    if record_hits: example.update({"w_bl_hit_list":hit_list})

    if w_bl_attacked_outputs is not None:
        result = score_sequence(inputs=inputs,
                                outputs=w_bl_attacked_outputs, # <-- blacklisted but attacked version
                                initial_seed=initial_seed,
                                dynamic_seed=dynamic_seed,
                                bl_proportion=bl_proportion, 
                                tokenizer=tokenizer,
                                record_hits=record_hits,
                                # breakpoint_on_hit=True, # banging head against wall
                                debug=False)
        if record_hits:
            bl_hits, num_toks_gend, hit_list = result
        else:
            bl_hits, num_toks_gend = result
        example.update({"w_bl_attacked_num_toks_gend_eq_0":(num_toks_gend == 0)})
        # if (num_toks_gend-bl_hits)/(num_toks_gend) < 1.0: breakpoint()

        if num_toks_gend == 0:
            # print("No tokens generated, odd, avoiding div by zero and returning -1's")
            wl_frac = -1
            bl_frac = -1
        else:
            wl_frac = (num_toks_gend-bl_hits)/num_toks_gend
            bl_frac = bl_hits/num_toks_gend
        w_bl_attacked_stats = {
            "w_bl_attacked_num_tokens_generated": num_toks_gend,
            "w_bl_attacked_whitelist_fraction": wl_frac,
            "w_bl_attacked_blacklist_fraction": bl_frac
        }
        example.update(w_bl_attacked_stats)
        if record_hits: example.update({"w_bl_attacked_hit_list":hit_list})
    
    return example



def aggregate_bl_stats(example: dict, idx: int, stat_table: dict):
    
    stat_table["baseline_stats"]["whitelist_fraction"] += example["baseline_stats"]["whitelist_fraction"]
    stat_table["baseline_stats"]["blacklist_fraction"] += example["baseline_stats"]["blacklist_fraction"]
    
    stat_table["w_bl_stats"]["whitelist_fraction"] += example["w_bl_stats"]["whitelist_fraction"]
    stat_table["w_bl_stats"]["blacklist_fraction"] += example["w_bl_stats"]["blacklist_fraction"]

    stat_table["no_bl_stats"]["whitelist_fraction"] += example["no_bl_stats"]["whitelist_fraction"]
    stat_table["no_bl_stats"]["blacklist_fraction"] += example["no_bl_stats"]["blacklist_fraction"]

    stat_table["num_examples"] += 1

    return example


def compute_ppl_single(prefix_and_output_text = None,
                        output_text = None,
                        oracle_model_name = None,
                        oracle_model = None,
                        oracle_tokenizer = None):

    with torch.no_grad():
        tokd_prefix = tokenize_and_truncate({"text":prefix_and_output_text}, completion_length=0, hf_model_name=oracle_model_name, tokenizer=oracle_tokenizer, model_max_seq_len=oracle_model.config.max_position_embeddings)["inputs"]
        tokd_inputs = tokd_prefix
        # if only want to score the "generation" part we need the suffix tokenization length
        tokd_suffix = tokenize_and_truncate({"text":output_text}, completion_length=0, hf_model_name=oracle_model_name, tokenizer=oracle_tokenizer, model_max_seq_len=oracle_model.config.max_position_embeddings)["inputs"]

        tokd_inputs = tokd_inputs.to(oracle_model.device)
        # make labels, mark if not including all positions
        tokd_labels = tokd_inputs.clone().detach()
        tokd_labels[:,:tokd_labels.shape[1]-tokd_suffix.shape[1]+1] = -100

        outputs = oracle_model(input_ids=tokd_inputs, labels=tokd_labels)
        loss = outputs.loss # avg CE loss all positions (except -100, TODO plz check that this is working correctly)
        ppl = torch.tensor(math.exp(loss))
    
    return loss.item(), ppl.item()


def evaluate_generation_fluency(example: dict, 
                                idx: int,
                                oracle_model_name = None,
                                oracle_model = None,
                                oracle_tokenizer = None):

    # pull out the required fields from the pipeline results
    inputs_plus_baseline_output = f"{example['truncated_input']}{example['baseline_completion']}"
    baseline_output = f"{example['baseline_completion']}"

    inputs_plus_no_bl_output = f"{example['truncated_input']}{example['no_bl_output']}"
    no_bl_output = f"{example['no_bl_output']}"

    inputs_plus_w_bl_output = f"{example['truncated_input']}{example['w_bl_output']}"
    w_bl_output = f"{example['w_bl_output']}"

    # add metrics
    loss, ppl = compute_ppl_single(inputs_plus_baseline_output, baseline_output, oracle_model_name, oracle_model, oracle_tokenizer)
    example["baseline_loss"] = loss
    example["baseline_ppl"] = ppl
    loss, ppl = compute_ppl_single(inputs_plus_no_bl_output, no_bl_output, oracle_model_name, oracle_model, oracle_tokenizer)
    example["no_bl_loss"] = loss
    example["no_bl_ppl"] = ppl
    loss, ppl = compute_ppl_single(inputs_plus_w_bl_output, w_bl_output, oracle_model_name, oracle_model, oracle_tokenizer)
    example["w_bl_loss"] = loss
    example["w_bl_ppl"] = ppl

    # del any temp values
    return example


def add_idx(example,idx):
    example.update({"idx":idx})
    return example


def check_input_lengths(example,idx, min_sample_len=0, min_prompt_len=0, min_completion_len=0):
    orig_sample_length = example["orig_sample_length"]
    prompt_length = example["prompt_length"]
    real_completion_length = example["real_completion_length"]

    # breakpoint()

    conds = all([
        orig_sample_length >= min_sample_len,
        prompt_length >= min_prompt_len,
        real_completion_length >= min_completion_len,
    ])
    return conds


def check_output_lengths(example,min_output_len=0):
    no_bl_output_len = example["no_bl_num_tokens_generated"]
    w_bl_output_len = example["w_bl_num_tokens_generated"]
    conds = all([
        no_bl_output_len >= min_output_len,
        w_bl_output_len >= min_output_len,
    ])
    return conds