evaluation_pipeline.py

# 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 types import NoneType

from typing import Union
import os
import argparse
from functools import partial
from tqdm import tqdm

import wandb
import torch
import numpy as np
import sklearn.metrics as metrics

from datasets import Dataset, Sequence
from transformers import DataCollatorWithPadding

from utils.submitit import str2bool  # better bool flag type for argparse
from utils.io import read_jsonlines, read_json, write_json, write_jsonlines
from utils.notebooks import filter_text_col_length, infer_length_column

from utils.evaluation import (
    SUPPORTED_METRICS,
    NO_CHECK_ARGS,
    ROC_TEST_STAT_SUFFIXES,
    FILTER_BY_COLUMNS,
    conditional_no_check_args,
    load_oracle_model,
    evaluate_ppl,
    load_detector,
    compute_z_scores,
    compute_windowed_z_scores,
    compute_run_len_chsqrd_stats,
    compute_repetition_diversity,
    compute_p_sp,
    compute_coherence,
    compute_mauve,
    compute_detect_retrieval,
    load_tokenizer,
    concat_rows,
)

print(f"Current huggingface cache dir: {os.environ['HF_HOME']}")

from datasets import disable_caching

disable_caching()


def main(args):
    ###########################################################################
    # Create output dir if it doesn't exist, and warn if it contains metric file
    ###########################################################################
    gen_table_w_metrics_path = f"{args.output_dir}/gen_table_w_metrics.jsonl"
    metrics_meta_path = f"{args.output_dir}/gen_table_w_metrics_meta.json"

    print(f"Output dir for this run: {args.output_dir}")
    # notify if exists
    if os.path.exists(args.output_dir):
        print(f"Output dir for this run already exists!")
        print(f"Contents: {sorted(os.listdir(args.output_dir))}")
        # warn if metrics file exists
        if os.path.exists(gen_table_w_metrics_path):
            if not args.overwrite_output_file:
                print(
                    f"WARNING: Exiting to avoid overwriting output file. "
                    f"Pass the '--overwrite_output_file' flag to ignore this check."
                )
                exit()
            else:
                print(
                    f"WARNING: Found existing generation files with metrics added at this output dir. "
                    f"Overwriting anyway :/"
                )
    else:
        # create the output dir where run artifacts are stored
        os.makedirs(args.output_dir)

    ###########################################################################
    # Parse metrics to log - ppl, zscore, etc
    ###########################################################################

    # check that all metrics are supported
    metric_support = [metric in SUPPORTED_METRICS for metric in args.evaluation_metrics]
    assert all(metric_support), (
        f"Unsupported metric '{args.evaluation_metrics[metric_support.index(False)]}' in"
        f" {args.evaluation_metrics}. Supported metrics are: {SUPPORTED_METRICS}"
    )
    # Hack check that if prefix_lengths exists then the method must be
    # detect-retrieval (for now) because other methods don't support the
    # sparse dataset with Nones all over the place
    if "prefix_lengths" in args.__dict__:
        # assert args.evaluation_metrics == [
        #     "detect-retrieval"
        # ], f"Currently, only the detect-retrieval metric supports the prefix_lengths column. "
        print(
            f"WARNING: Found prefix_lengths column assuming that this is either retireval or detectgpt"
        )

    print(f"Evaluation metrics to compute: {args.evaluation_metrics}")

    ###########################################################################
    # Load generations
    ###########################################################################
    print(f"Input dir for this run: {args.input_dir}")
    print(f"Loading previously generated outputs for evaluation via oracle model and metrics...")

    # check for the "attacked version" of the gen table first
    gen_table_meta_path = f"{args.input_dir}/gen_table_attacked_meta.json"
    gen_table_path = f"{args.input_dir}/gen_table_attacked.jsonl"
    safe_gen_table_path = f"{args.input_dir}/gen_table_attacked_safe.jsonl"
    loaded_attacked = True

    attack_variants_exist = [
        os.path.exists(gen_table_meta_path),
        os.path.exists(gen_table_path),
    ]
    if not all(attack_variants_exist):
        loaded_attacked = False
        gen_table_meta_path = f"{args.input_dir}/gen_table_meta.json"
        gen_table_path = f"{args.input_dir}/gen_table.jsonl"
        safe_gen_table_path = f"{args.input_dir}/gen_table_safe.jsonl"

        assert os.path.exists(
            gen_table_meta_path
        ), f"failed file check for prev generations metadata json file: {gen_table_meta_path}"
        assert os.path.exists(
            gen_table_path
        ), f"failed file check for prev generations jsonl file: {gen_table_path}"

    assert not os.path.exists(safe_gen_table_path), (
        f"failed for safety bc there is a secondary 'safe' marked file",
        f" in this dir indicating a possible issue with the generation step. ",
    )

    cmdline_args = args.__dict__.copy()
    prev_gen_table_meta = read_json(gen_table_meta_path)

    joined_args = prev_gen_table_meta.copy()
    for k, v in cmdline_args.items():
        if v is not None:
            joined_args.update({k: v})
        else:
            print(
                f"cmdline arg {k} is None, leaving it as the value found in the input metadata: {prev_gen_table_meta[k]}"
            )

    # check that the args used to generate the prev generations are the same as
    # the current args, for the intersection of keys
    if not args.overwrite_args:
        # update the no check args based on the current state of args
        current_no_check_args = conditional_no_check_args(
            NO_CHECK_ARGS, args.evaluation_metrics, args
        )

        for key in prev_gen_table_meta.keys():
            if key in current_no_check_args:
                continue
            assert joined_args[key] == prev_gen_table_meta[key], (
                f"failed for safety bc after merging the prev metadata with "
                f"the current cmdline args, values for '{key}' are not the same. "
                f"in metadata: {prev_gen_table_meta[key]}, passed: {cmdline_args[key]}. "
                f"Pass the '--overwrite_args' flag to ignore this check."
            )

    args = argparse.Namespace(**joined_args)
    gen_table = [ex for ex in read_jsonlines(gen_table_path)]
    if args.limit_rows == -1:
        gen_table_ds = Dataset.from_list(gen_table)
    else:
        gen_table_ds = Dataset.from_list(gen_table[: args.limit_rows])

    ###########################################################################
    # Extract the seeding scheme fine grained parameters
    ###########################################################################
    from utils.evaluation import scheme_hparam_extractor

    args.__dict__.update(scheme_hparam_extractor(args.seeding_scheme))

    print(f"seeding_scheme: {args.seeding_scheme}")
    print(f"prf_type: {args.prf_type}")
    print(f"anchored: {args.anchored}")
    print(f"context_width: {args.context_width}")
    print(f"self_salt: {args.self_salt}")

    ###########################################################################
    # Concat logic for multiple generations
    ###########################################################################

    if args.concat_rows != 0:
        assert isinstance(args.concat_rows, int), f"Invalid concat_rows arg: {args.concat_rows}. "

        # set to all rows if -1
        if args.concat_rows == -1:
            args.concat_rows = len(gen_table_ds)

        if args.shuffle_before_concat:
            print(f"Shuffling the gen table before concatenating every {args.concat_rows} rows...")
            gen_table_ds = gen_table_ds.shuffle()

        print(f"Concatenating every {args.concat_rows} rows of the gen table...")

        # we concat all cols in OUTPUT_TEXT_COLUMN_NAMES
        # and update the length col to reflect the new length
        # which means we need to tokenize the new text temporarily
        # to get the new length

        tokenizer = load_tokenizer(args)

        concat_partial = partial(concat_rows, tokenizer=tokenizer, args=args)

        # manually write a btach loop bc hf doesnt support returning fewer rows than input
        concatenated_rows = []
        for i in tqdm(range(0, len(gen_table_ds), args.concat_rows)):
            batch = gen_table_ds[i : i + args.concat_rows]
            concatenated_rows.append(concat_partial(batch))
        gen_table_concated_ds = Dataset.from_list(concatenated_rows)

        # overwrite the args.max_new_tokens to reflect the implicit new target length T
        # which is concat_rows * max_new_tokens
        args.max_new_tokens = args.concat_rows * args.max_new_tokens

        # write the dataset out in the same filename as the original
        # but check that the input dir is different from the output dir
        assert (
            args.input_dir != args.output_dir
        ), f"Input dir and output dir must be different to write out the result of concat rows."

        if loaded_attacked:
            concat_meta_path = f"{args.output_dir}/gen_table_attacked_meta.json"
            concat_gen_table_path = f"{args.output_dir}/gen_table_attacked.jsonl"
        else:
            concat_meta_path = f"{args.output_dir}/gen_table_meta.json"
            concat_gen_table_path = f"{args.output_dir}/gen_table.jsonl"

        write_json(args.__dict__, concat_meta_path, indent=4)
        gen_table_concated_lst = [ex for ex in gen_table_concated_ds]
        write_jsonlines(gen_table_concated_lst, concat_gen_table_path)
    else:
        gen_table_concated_ds = gen_table_ds

    ###########################################################################
    # Additional args setup
    ###########################################################################
    # if target_T is not specified, use max_new_tokens (which will be in the reloaded gen metadata)
    # and potentially overwritten by the concat logic above
    if args.target_T == 0:
        args.target_T = args.max_new_tokens

    # storing slurm info to allow auditing logfiles
    # note this is set after the metadata check to ignore overwriting
    args.SLURM_JOB_ID = os.getenv("SLURM_JOB_ID")
    args.SLURM_ARRAY_JOB_ID = os.getenv("SLURM_ARRAY_JOB_ID")
    args.SLURM_ARRAY_TASK_ID = os.getenv("SLURM_ARRAY_TASK_ID")

    ###########################################################################
    # Start logging, we wait to do this until after loading the generations
    # so that we can log the args used to generate them unioned with the
    # cmdline args
    ###########################################################################
    if args.wandb:
        # start a new wandb run to track this experiment, will send data to it
        run = wandb.init(
            # set the wandb project where this run will be logged
            project=args.wandb_project,
            entity=args.wandb_entity,
            name=f"{args.run_name}",
            # track hyperparameters and run metadata
            config=args,
            tags=args.wandb_tags,
        )

    ###########################################################################
    # Perplexity (PPL) evaluation
    # NOTE: basically requires a model on gpu, or is extremely slow
    ###########################################################################
    if "ppl" in args.evaluation_metrics:
        assert args.oracle_model_name_or_path, "PPL metric requires oracle model."

        # Load the oracle model for PPL measurement
        oracle_model, oracle_tokenizer, _ = load_oracle_model(args)

        # construct the collator
        data_collator = DataCollatorWithPadding(
            tokenizer=oracle_tokenizer, padding=True, pad_to_multiple_of=8
        )

        # construct fluency/ppl partial
        evaluate_ppl_partial = partial(
            evaluate_ppl,
            oracle_model_name=args.oracle_model_name_or_path,
            oracle_model=oracle_model,
            oracle_tokenizer=oracle_tokenizer,
            data_collator=data_collator,
        )

        print(f"Computing metrics on model generations: {gen_table_concated_ds}")

        gen_table_w_ppl_ds = gen_table_concated_ds.map(
            evaluate_ppl_partial,
            batched=True,
            batch_size=args.ppl_batch_size,
            load_from_cache_file=False,
            keep_in_memory=True,
        )

        # clear the model just for fun
        oracle_model = oracle_model.to(torch.device("cpu"))
        del oracle_model
    else:
        gen_table_w_ppl_ds = gen_table_concated_ds

    ###########################################################################
    # Cheap to load, and required for all detectors so load it first
    watermark_detector = load_detector(args)

    # Map setup for all dataset operations:
    map_setup = dict(batched=False, load_from_cache_file=False)
    ###########################################################################
    # z-score evaluation
    # NOTE: requires a gpu because if original source of watermark randomness,
    # RNG, is gpu based, then detector should be on gpu as well
    ###########################################################################
    if "z-score" in args.evaluation_metrics:
        # set up the partial
        compute_z_scores_partial = partial(
            compute_z_scores,
            watermark_detector=watermark_detector,
            args=args,
        )

        gen_table_w_zscore_ds = gen_table_w_ppl_ds.map(
            compute_z_scores_partial, **map_setup, desc="Computing z-scores"
        )

    else:
        gen_table_w_zscore_ds = gen_table_w_ppl_ds

    ###########################################################################
    # Windowed z-score evaluation
    ###########################################################################

    if "windowed-z-score" in args.evaluation_metrics:
        # set up the windowed partial
        compute_windowed_z_scores_partial = partial(
            compute_windowed_z_scores,
            watermark_detector=watermark_detector,
            args=args,
        )

        gen_table_w_windowed_zscore_ds = gen_table_w_zscore_ds.map(
            compute_windowed_z_scores_partial, **map_setup, desc="Computing windowed z-scores"
        )
    else:
        gen_table_w_windowed_zscore_ds = gen_table_w_zscore_ds

    ###########################################################################
    # run-len-chisqrd evaluation
    ###########################################################################
    if "run-len-chisqrd" in args.evaluation_metrics:
        assert "w_wm_output_green_token_mask" in gen_table_w_windowed_zscore_ds.column_names, (
            f"Currently, run-len-chisqrd metric requires the green token masks to be computed previously "
            f"by one of the z-score metrics."
        )
        # this ^ is unused currently, but we will need it to remove the assert condition above

        # set up the run len chisqrd partial
        compute_run_len_chisqrd_partial = partial(
            compute_run_len_chsqrd_stats,
            watermark_detector=watermark_detector,
            args=args,
        )

        gen_table_w_run_len_chisqrd_ds = gen_table_w_windowed_zscore_ds.map(
            compute_run_len_chisqrd_partial, **map_setup, desc="Computing runlength tests"
        )
    else:
        gen_table_w_run_len_chisqrd_ds = gen_table_w_windowed_zscore_ds

    ###########################################################################
    # Diversity and Repetition evaluation
    ###########################################################################

    if "repetition" in args.evaluation_metrics or "diversity" in args.evaluation_metrics:
        # set up the partial
        compute_repetition_partial = partial(
            compute_repetition_diversity,
            include_repetition=("repetition" in args.evaluation_metrics),
            include_diversity=("diversity" in args.evaluation_metrics),
        )

        gen_table_w_repetition_ds = gen_table_w_run_len_chisqrd_ds.map(
            compute_repetition_partial, **map_setup, desc="Computing text repetition and diversity"
        )
    else:
        gen_table_w_repetition_ds = gen_table_w_run_len_chisqrd_ds

    ###########################################################################
    # P-SP evaluation
    ###########################################################################

    if "p-sp" in args.evaluation_metrics:
        print(f"Loading the P-SP model and computing P-SP")
        gen_table_w_p_sp_ds = compute_p_sp(gen_table_w_repetition_ds)
    else:
        gen_table_w_p_sp_ds = gen_table_w_repetition_ds

    ###########################################################################
    # Coherence evaluation
    ###########################################################################

    if "coherence" in args.evaluation_metrics:
        print(f"Computing coherence")
        gen_table_w_coherence_ds = compute_coherence(gen_table_w_p_sp_ds)
    else:
        gen_table_w_coherence_ds = gen_table_w_p_sp_ds

    ###########################################################################
    # Mauve evaluation
    ###########################################################################

    if "mauve" in args.evaluation_metrics:
        print(f"Computing mauve")
        gen_table_w_mauve_ds = compute_mauve(gen_table_w_coherence_ds)
    else:
        gen_table_w_mauve_ds = gen_table_w_coherence_ds

    ###########################################################################
    # Retrieval detection
    ###########################################################################

    if "detect-retrieval" in args.evaluation_metrics:
        print(f"Computing detect retrieval")
        gen_table_w_detect_retrieval_ds = compute_detect_retrieval(gen_table_w_mauve_ds, args=args)
    else:
        gen_table_w_detect_retrieval_ds = gen_table_w_mauve_ds

    if "prefix_length" in gen_table_w_detect_retrieval_ds.features:
        if "no_wm_output_retrieval_score" in gen_table_w_detect_retrieval_ds.features:
            print("Avg scores at each prefix length for no_wm_output:")
            print(
                gen_table_w_detect_retrieval_ds.to_pandas()
                .groupby("prefix_length")["no_wm_output_retrieval_score"]
                .describe()
            )
        if "w_wm_output_retrieval_score" in gen_table_w_detect_retrieval_ds.features:
            print("Avg scores at each prefix length for w_wm_output:")
            print(
                gen_table_w_detect_retrieval_ds.to_pandas()
                .groupby("prefix_length")["w_wm_output_retrieval_score"]
                .describe()
            )
        if "w_wm_output_attacked_retrieval_score" in gen_table_w_detect_retrieval_ds.features:
            print("Avg scores at each prefix length for no_wm_output_attacked:")
            print(
                gen_table_w_detect_retrieval_ds.to_pandas()
                .groupby("prefix_length")["w_wm_output_attacked_retrieval_score"]
                .describe()
            )

    ###########################################################################
    # Detectgpt detection
    ###########################################################################
    if "detectgpt" in args.evaluation_metrics:
        assert args.evaluation_metrics == ["detectgpt"], (
            f"Detectgpt must be run separately from other metrics. "
            f"Found: {args.evaluation_metrics}. "
        )
        # check that the right score column exists
        assert any(
            ["detectgpt_score" in col for col in gen_table_w_detect_retrieval_ds.column_names]
        ), (
            f"Detectgpt metric requires the detectgpt_score column to be computed previously "
            f"but no such cols exist in this file."
        )
        print(
            f"Evaluating detectgpt by simply computing ROC-AUC metrics on the scores that already exist"
        )
        gen_table_w_metrics_ds = gen_table_w_detect_retrieval_ds

        # if we loaded an attack file, since detect gpt only outputs a baseline score col
        # and a no_wm_output score col (which is implcitly the attack col if the file was attacked)
        # we need to add the attacked score col to the dataset, and remove the no_wm score col
        if loaded_attacked:
            for suff in ["100_d", "100_z"]:
                gen_table_w_metrics_ds = gen_table_w_metrics_ds.add_column(
                    f"w_wm_output_attacked_detectgpt_score_{suff}",
                    gen_table_w_metrics_ds[f"no_wm_output_detectgpt_score_{suff}"],
                )
                gen_table_w_metrics_ds = gen_table_w_metrics_ds.remove_columns(
                    [f"no_wm_output_detectgpt_score_{suff}"]
                )
    else:
        ###########################################################################
        # Write the final dataset out to disk in jsonl format
        # with the metrics added
        ###########################################################################

        # last applied metric, NOTE which will of course change as more are added
        gen_table_w_metrics_ds = gen_table_w_detect_retrieval_ds

        # write the metadata file, which is a union of the previous metadata
        # and the current cmdline args
        write_json(args.__dict__, metrics_meta_path, indent=4)

        gen_table_w_metrics_lst = [ex for ex in gen_table_w_metrics_ds]
        write_jsonlines(gen_table_w_metrics_lst, gen_table_w_metrics_path)

    ###########################################################################
    # Log the metric series to wandb
    ###########################################################################
    # log the metrics to wandb
    if args.wandb:
        # find cols that should be logged in a table
        tabular_column_types = ["string", "bool"]
        tabular_column_names = [
            name
            for name, _ in filter(
                lambda tup: tup[1].dtype in tabular_column_types,
                gen_table_w_metrics_ds.features.items(),
            )
        ]
        # the rest should be logged as series
        series_column_names = [
            name
            for name, _ in filter(
                lambda tup: tup[1].dtype not in tabular_column_types,
                gen_table_w_metrics_ds.features.items(),
            )
        ]

        for metric_name in series_column_names:
            # summarize series metrics as mean by default
            wandb.define_metric(metric_name, summary="mean")

        if args.log_raw_series:
            # log the raw series
            for example in tqdm(
                gen_table_w_metrics_ds.remove_columns(tabular_column_names),
                desc="Logging series metrics to wandb",
            ):
                run.log(example)

        if args.log_raw_tabular:
            # log the raw tabular data
            # but also include the dataset index as a column
            series_column_names.remove("idx")
            table = wandb.Table(
                dataframe=gen_table_w_metrics_ds.remove_columns(series_column_names).to_pandas()
            )
            run.log({"output_table": table})

        ###########################################################################
        # Filter rows, then log means to wandb
        ###########################################################################
        assert (
            args.target_T - args.lower_tolerance_T
        ) >= 0, "target_T - lower_tolerance_T must be >= 0"

        target_T = args.target_T
        lower_tolerance = args.lower_tolerance_T
        upper_tolerance = args.upper_tolerance_T
        filtered_table = gen_table_w_metrics_ds.to_pandas()  # explictly convert lists

        for col in args.filter_by_columns:
            length_col_name = infer_length_column(col, filtered_table, args=args)
            filtered_table = filter_text_col_length(
                filtered_table,
                text_col_name=length_col_name,
                count_suffix="",
                upper_T=target_T + upper_tolerance,
                lower_T=target_T - lower_tolerance,
            )

        # Save filtered mean values:
        for metric_name in series_column_names:
            filtered_name = f"f_{target_T}p{upper_tolerance}m{lower_tolerance}_{metric_name}"
            try:
                run.summary[f"{filtered_name}_mean"] = filtered_table[metric_name].mean()
                run.summary[f"{filtered_name}_std"] = filtered_table[metric_name].std()
            except TypeError:
                two_dim_mean = filtered_table[metric_name].apply(np.mean).mean()

        ###########################################################################
        # Compute ROC-AUC and send to wandb
        ###########################################################################
        try:
            test_stats = args.roc_test_stat
            if isinstance(test_stats, str):
                test_stats = [test_stats]
            for test_stat in test_stats:
                for attacked in [True, False]:
                    try:
                        roc_auc, fpr, tpr, thresholds, tpr_at_X_fpr = _roc_metrics_for_wandb(
                            filtered_table, test_stat, attacked=attacked
                        )
                        run.summary[
                            f"{'attacked_' if attacked else ''}{test_stat}_roc_auc"
                        ] = roc_auc
                        run.summary[
                            f"{'attacked_' if attacked else ''}{test_stat}_tpr_at_X_fpr"
                        ] = tpr_at_X_fpr

                        # for tp, fp, thr in tqdm(
                        #     zip(tpr, fpr, thresholds), desc="Logging ROC curve"
                        # ):
                        #     run.log(
                        #         {
                        #             f"{'attacked_' if attacked else ''}{test_stat}_fpr": fp,
                        #             f"{'attacked_' if attacked else ''}{test_stat}_tpr": tp,
                        #             f"{'attacked_' if attacked else ''}thr": thr,
                        #         }
                        #     )
                        data = [[x, y] for (x, y) in zip(fpr, tpr)]
                        table = wandb.Table(data=data, columns=["fpr", "tpr"])
                        run.log(
                            {
                                f"{'attacked_' if attacked else ''}{test_stat}": wandb.plot.line(
                                    table,
                                    "fpr",
                                    "tpr",
                                    title=f"ROC ({test_stat}{',attacked' if attacked else ',clean'})",
                                )
                            }
                        )
                        print(f"Successfully logged ROC-AUC metrics for {test_stat}.")

                    except Exception as e:
                        if args.verbose:
                            print(e)
                            print(
                                f"Failed to log ROC-AUC metrics for {'attacked output' if attacked else ''} {test_stat}."
                                f"Metric probably was not computed and or attack col not present."
                            )
        except Exception as e:
            if args.verbose:
                print(f"Exception: {e}")
                print(
                    f"Failed to log ROC-AUC metrics. ",
                    f"Make sure the test statistic required for detection ({test_stat}) has been computed!",
                )

        ################################################################################
        # NOTE we do that ^^^ basic ROC logic first because it's faster
        # as well as the manual prefix lengths at T logic bc that's also faster
        ################################################################################

        # Handle z @ T but for the retrieval and detectgpt scores that are evaluated
        # manually at each prefix length.  Use groupby to compute the mean and std
        # for each prefix length for any of the feats that have retrieval_score in them,
        # then log those pairs to wandb.
        at_T_df = gen_table_w_metrics_ds.to_pandas()

        for name, feat in gen_table_w_metrics_ds.features.items():
            if "retrieval_score" in name and "prefix_length" in at_T_df.columns:
                # compute the mean and std for each prefix length
                # and log those pairs to wandb
                df_view = at_T_df.groupby("prefix_length")[name].describe()[["mean", "std"]]
                T_indices = df_view.index

                # for idx, (mean, std) in df_view.iterrows():
                #     run.log(data={f"{name}_mean": mean, f"{name}_std": std, "idx_T": idx})
                # log this triple as a table instead like the ROC curve above
                # where the first two are plotted and the third is the x axis
                data = [[x, y, z] for x, (y, z) in df_view.iterrows()]
                table = wandb.Table(data=data, columns=["idx_T", "mean", "std"])
                # compute stderr from std
                table.add_column(
                    "stderr",
                    [
                        std / np.sqrt(len(at_T_df[at_T_df["prefix_length"] == idx]))
                        for idx, std in zip(T_indices, df_view["std"])
                    ],
                )
                # first log mean
                run.log({f"{name}": wandb.plot.line(table, "idx_T", "mean", title=f"{name} mean")})
                # then log std err
                run.log(
                    {
                        f"{name}_stderr": wandb.plot.line(
                            table, "idx_T", "stderr", title=f"{name} stderr"
                        )
                    }
                )

                # also compute an AUC at each prefix len idx by treating the name col as the positives
                # and the baseline_completion_retrieval_score as the negatives
                # then log those pairs to wandb
                if name != "baseline_completion_retrieval_score":
                    pos_negs_at_T = at_T_df.groupby("prefix_length")[
                        [name, "baseline_completion_retrieval_score"]
                    ]
                    # auc_at_T = []
                    # tpr_at_X_fpr = []
                    all_aucs, all_tpr_at_X_fpr = [], []
                    for idx, sub_df in pos_negs_at_T:
                        pos = sub_df[name]
                        neg = sub_df["baseline_completion_retrieval_score"]
                        # convert to arrays and remove nans
                        pos = pos.to_numpy()[~np.isnan(pos.to_numpy())]
                        neg = neg.to_numpy()[~np.isnan(neg.to_numpy())]

                        fpr, tpr, thresholds = metrics.roc_curve(
                            np.concatenate([np.ones_like(pos), np.zeros_like(neg)]),  # labels
                            np.concatenate([pos, neg]),  # scores
                            pos_label=1,
                        )
                        auc = metrics.auc(fpr, tpr)
                        try:
                            tpr_at_X_fpr = tpr[np.where(fpr < 1e-3)[0][-1]]
                        except IndexError:
                            tpr_at_X_fpr = float("NaN")
                        all_aucs.append(auc)
                        all_tpr_at_X_fpr.append(tpr_at_X_fpr)

                        # run.log(data={f"{name}_auc_at_T": auc, "idx_T": idx})
                    # log this triple as a table instead like the AUC and tpr at X fpr below
                    # where the first two are plotted and the third is the x axis
                    data = [
                        [x, y, z] for x, (y, z) in zip(T_indices, zip(all_aucs, all_tpr_at_X_fpr))
                    ]
                    table = wandb.Table(data=data, columns=["idx_T", "aucs", "tpr_at"])
                    run.log(
                        {
                            f"{name}_aucs": wandb.plot.line(
                                table, "idx_T", "aucs", title=f"{name} aucs"
                            )
                        }
                    )
                    run.log(
                        {
                            f"{name}_tpr_at": wandb.plot.line(
                                table, "idx_T", "tpr_at", title=f"{name} tpr_at"
                            )
                        }
                    )

            elif "detectgpt_score" in name and "prefix_length" in at_T_df.columns:
                # this covers detectgpt_score_100_d and variants
                # compute the mean and std for each prefix length
                # and log those pairs to wandb
                df_view = at_T_df.groupby("prefix_length")[name].describe()[["mean", "std"]]
                T_indices = df_view.index

                # for idx, (mean, std) in df_view.iterrows():
                #     run.log(data={f"{name}_mean": mean, f"{name}_std": std, "idx_T": idx})
                # log this triple as a table instead like the ROC curve above
                # where the first two are plotted and the third is the x axis
                data = [[x, y, z] for x, (y, z) in df_view.iterrows()]
                table = wandb.Table(data=data, columns=["idx_T", "mean", "std"])

                # compute stderr from std
                table.add_column(
                    "stderr",
                    [
                        std / np.sqrt(len(at_T_df[at_T_df["prefix_length"] == idx]))
                        for idx, std in zip(T_indices, df_view["std"])
                    ],
                )
                # first log mean
                run.log({f"{name}": wandb.plot.line(table, "idx_T", "mean", title=f"{name} mean")})
                # then log std err
                run.log(
                    {
                        f"{name}_stderr": wandb.plot.line(
                            table, "idx_T", "stderr", title=f"{name} stderr"
                        )
                    }
                )

                # also compute an AUC at each prefix len idx by treating the name col as the positives
                # and the baseline_completion_retrieval_score as the negatives
                # then log those pairs to wandb
                if "baseline_completion_detectgpt_score" not in name:
                    # check which suffix this is in ["_100_d", "_100_z"]
                    # and use that to set the baseline/falst col
                    if name.endswith("_100_d"):
                        baseline_col = "baseline_completion_detectgpt_score_100_d"
                    elif name.endswith("_100_z"):
                        baseline_col = "baseline_completion_detectgpt_score_100_z"
                    pos_negs_at_T = at_T_df.groupby("prefix_length")[[name, baseline_col]]
                    # auc_at_T = []
                    # tpr_at_X_fpr = []
                    all_aucs, all_tpr_at_X_fpr = [], []
                    for idx, sub_df in pos_negs_at_T:
                        pos = sub_df[name]
                        neg = sub_df[baseline_col]
                        # convert to arrays and remove nans
                        pos = pos.to_numpy()[~np.isnan(pos.to_numpy())]
                        neg = neg.to_numpy()[~np.isnan(neg.to_numpy())]

                        fpr, tpr, thresholds = metrics.roc_curve(
                            np.concatenate([np.ones_like(pos), np.zeros_like(neg)]),  # labels
                            np.concatenate([pos, neg]),  # scores
                            pos_label=1,
                        )
                        auc = metrics.auc(fpr, tpr)
                        try:
                            tpr_at_X_fpr = tpr[np.where(fpr < 1e-3)[0][-1]]
                        except IndexError:
                            tpr_at_X_fpr = float("NaN")
                        all_aucs.append(auc)
                        all_tpr_at_X_fpr.append(tpr_at_X_fpr)

                        # run.log(data={f"{name}_auc_at_T": auc, "idx_T": idx})
                    # log this triple as a table instead like the AUC and tpr at X fpr below
                    # where the first two are plotted and the third is the x axis
                    data = [
                        [x, y, z] for x, (y, z) in zip(T_indices, zip(all_aucs, all_tpr_at_X_fpr))
                    ]
                    table = wandb.Table(data=data, columns=["idx_T", "aucs", "tpr_at"])
                    run.log(
                        {
                            f"{name}_aucs": wandb.plot.line(
                                table, "idx_T", "aucs", title=f"{name} aucs"
                            )
                        }
                    )
                    run.log(
                        {
                            f"{name}_tpr_at": wandb.plot.line(
                                table, "idx_T", "tpr_at", title=f"{name} tpr_at"
                            )
                        }
                    )

        ###########################################################################
        # Compute our @ T detection metrics and send to wandb
        ###########################################################################

        # Merge z_at_T and other sequence metrics so they can be shown in wandb:
        for name, feat in gen_table_w_metrics_ds.features.items():
            if isinstance(feat, Sequence):
                max_feat_seq_len = max([len(l) for l in gen_table_w_metrics_ds[name]])
                merging_seq = np.zeros(max_feat_seq_len)
                counts = np.zeros(max_feat_seq_len)
                proto_variance = np.zeros(max_feat_seq_len)
                for entry in gen_table_w_metrics_ds[name]:
                    len_seq = len(entry)
                    delta = entry * counts[:len_seq] - merging_seq[:len_seq]
                    # Accumulate ragged sum over entries:
                    counts[:len_seq] += 1
                    merging_seq[:len_seq] += entry[: len(merging_seq)]
                    # Compute ragged, running variance via Welford:
                    gamma = entry * counts[:len_seq] - merging_seq[:len_seq]
                    proto_variance[:len_seq] += (delta / counts[:len_seq]) * (
                        gamma / counts[:len_seq]
                    )

                mask = counts != 0
                averaged_seq = merging_seq.copy()
                averaged_seq[mask] /= counts
                averaged_seq[~mask] = float("NaN")

                seq_stderr = proto_variance.copy()
                seq_stderr[counts > 1] = np.sqrt(
                    proto_variance[counts > 1] / (counts[counts > 1] - 1)
                ) / np.sqrt(counts[counts > 1])
                seq_stderr[counts <= 1] = float("NaN")
                # for idx, (avg, stderr) in enumerate(zip(averaged_seq[mask], seq_stderr[mask])):
                #     run.log(data={f"{name}_avg": avg, f"{name}_stderr": stderr, "idx_T": idx})
                # log this triple as a table instead like the ROC curve above
                # where the first two are plotted and the third is the x axis
                data = [
                    [x, y, z]
                    for (x, y, z) in zip(
                        averaged_seq[mask], seq_stderr[mask], range(len(averaged_seq[mask]))
                    )
                ]
                table = wandb.Table(data=data, columns=["avg", "stderr", "idx_T"])

                # first plot avg
                run.log({f"{name}": wandb.plot.line(table, "idx_T", "avg", title=f"{name} avg")})
                # then plot stderr
                run.log(
                    {
                        f"{name}_stderr": wandb.plot.line(
                            table, "idx_T", "stderr", title=f"{name} stderr"
                        )
                    }
                )

        # Compute AUC_at_T
        # For now we'll just do a dumb loop over scipy.roc_curve, but this could be batched
        test_stats = args.roc_test_stat
        if isinstance(test_stats, str):
            test_stats = [test_stats]

        for test_stat in test_stats:
            for attacked in [True, False]:
                base_col = f"baseline_completion_{test_stat}_at_T"
                w_wm_col = f"w_wm_output{'_attacked' if attacked else ''}_{test_stat}_at_T"
                name = f"w_wm{'_attacked' if attacked else ''}_{test_stat}_at_T"

                if w_wm_col in gen_table_w_metrics_ds.features.keys():  # metric was computed
                    print(f"Computing AUC at T for {name}.")
                    max_length = min(
                        max([len(l) for l in gen_table_w_metrics_ds[base_col]]),
                        max([len(l) for l in gen_table_w_metrics_ds[w_wm_col]]),
                    )

                    all_aucs, all_tpr_at_X_fpr = [], []
                    for T in range(1, max_length):
                        w_wm_stats = np.array(
                            [t[T] for t in gen_table_w_metrics_ds[w_wm_col] if len(t) > T]
                        )

                        baseline_stats = np.array(
                            [t[T] for t in gen_table_w_metrics_ds[base_col] if len(t) > T]
                        )[: len(w_wm_stats)]
                        all_scores = np.concatenate([baseline_stats, w_wm_stats])

                        baseline_labels = np.zeros_like(baseline_stats)
                        attacked_labels = np.ones_like(w_wm_stats)
                        all_labels = np.concatenate([baseline_labels, attacked_labels])

                        if len(np.unique(all_labels)) < 2:
                            roc_auc = float("NaN")
                            tpr_at_X_fpr = float("NaN")
                        else:
                            fpr, tpr, thresholds = metrics.roc_curve(
                                all_labels, all_scores, pos_label=1
                            )
                            roc_auc = metrics.auc(fpr, tpr)
                            try:
                                tpr_at_X_fpr = tpr[np.where(fpr < 1e-3)[0][-1]]
                            except IndexError:
                                tpr_at_X_fpr = float("NaN")

                        all_aucs.append(roc_auc)
                        all_tpr_at_X_fpr.append(tpr_at_X_fpr)
                    # for idx, (aucs, tpr_at) in enumerate(zip(all_aucs, all_tpr_at_X_fpr)):
                    #     run.log(data={f"{name}_aucs": aucs, f"{name}_tpr_at": tpr_at, "idx_T": idx})
                    # log these two separately using a table
                    data = [
                        [x, y, z]
                        for (x, y, z) in zip(all_aucs, all_tpr_at_X_fpr, range(len(all_aucs)))
                    ]
                    table = wandb.Table(data=data, columns=["aucs", "tpr_at", "idx_T"])
                    run.log(
                        {
                            f"{name}_aucs": wandb.plot.line(
                                table, "idx_T", "aucs", title=f"{name} aucs"
                            )
                        }
                    )
                    run.log(
                        {
                            f"{name}_tpr_at": wandb.plot.line(
                                table, "idx_T", "tpr_at", title=f"{name} tpr_at"
                            )
                        }
                    )

        # finish the wandb run
        run.finish()

    return


def _roc_metrics_for_wandb(
    gen_table_ds, test_stat="z_score", prefix="", attacked=False, remove_nan=True
):
    # In theory, we actually should be filtering the attacked column too, but we know these
    # end up very short sometimes. So, to make sure the logic works, we just
    # filter for any rows where the test metrics are NaN and note the damage

    baseline_col_name = f"{prefix}baseline_completion_{test_stat}"
    if "retrieval" in test_stat:
        if attacked:
            w_wm_col_name = f"{prefix}w_wm_output_attacked_retrieval_score"
        else:
            w_wm_col_name = f"{prefix}{args.retrieval_db_column}_retrieval_score"
    elif "detectgpt" in test_stat:
        if attacked:
            w_wm_col_name = f"{prefix}w_wm_output_attacked_{test_stat}"
        else:
            w_wm_col_name = f"{prefix}no_wm_output_{test_stat}"
    else:
        w_wm_col_name = f"{prefix}w_wm_output{'_attacked' if attacked else ''}_{test_stat}"

    # drop nans in either column
    if remove_nan:
        orig_length = len(gen_table_ds)
        gen_table_ds = gen_table_ds.dropna(subset=[baseline_col_name, w_wm_col_name])
        if orig_length != len(gen_table_ds):
            print(
                f"NOTE: During ROC calculation, dropped {orig_length - len(gen_table_ds)} rows due to NaNs in {baseline_col_name} or {w_wm_col_name}"
            )

    baseline_stats = gen_table_ds[baseline_col_name].values
    w_wm_stats = gen_table_ds[w_wm_col_name].values
    all_scores = np.concatenate([baseline_stats, w_wm_stats])

    baseline_labels = np.zeros_like(baseline_stats)
    attacked_labels = np.ones_like(w_wm_stats)
    all_labels = np.concatenate([baseline_labels, attacked_labels])

    fpr, tpr, thresholds = metrics.roc_curve(all_labels, all_scores, pos_label=1)
    roc_auc = metrics.auc(fpr, tpr)
    try:
        tpr_at_X_fpr = tpr[np.where(fpr < 1e-3)[0][-1]]
    except IndexError:
        tpr_at_X_fpr = float("NaN")
    return roc_auc, fpr, tpr, thresholds, tpr_at_X_fpr


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Run evaluation pipeline for watermark detection")
    parser.add_argument(
        "--evaluation_metrics",
        type=str,
        default="all",
        help="Comma separated list of columns to remove from the dataset before generation.",
    )
    parser.add_argument(
        "--compute_scores_at_T",
        type=str2bool,
        default=True,
        help="Whether to compute (applicable) metrics at each T index in the output/text columns.",
    )
    parser.add_argument(
        "--overwrite_args",
        type=str2bool,
        default=False,
        help="Whether to overwrite the shared args in the metadata file with the current, runtime args.",
    )
    parser.add_argument(
        "--oracle_model_name_or_path",
        type=str,
        default="facebook/opt-6.7b",
        help="Oracle model, path to pretrained model or model identifier from huggingface.co/models.",
    )
    parser.add_argument(
        "--load_fp16",
        type=str2bool,
        default=None,
        help=(
            "Whether to run model (for ppl) in float16 precsion, note, will overwrite error as a reminder that "
            "generation was run in other mode, even though there's no hard requirement that these match."
        ),
    )
    parser.add_argument(
        "--ppl_batch_size",
        type=int,
        default=1,
        help="Batch size for ppl eval.",
    )
    parser.add_argument(
        "--seeding_scheme",
        type=Union[str, NoneType],
        default=None,
        help="Seeding scheme to use to generate the greenlists at each generation and verification step.",
    )
    parser.add_argument(
        "--gamma",
        type=Union[float, NoneType],
        default=None,
        help="The fraction of the vocabulary to partition into the greenlist at each generation and verification step.",
    )
    parser.add_argument(
        "--normalizers",
        type=Union[str, NoneType],
        default=None,
        help="Single or comma separated list of the preprocessors/normalizer names to use when performing watermark detection.",
    )
    parser.add_argument(
        "--ignore_repeated_ngrams",
        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(
        "--return_green_token_mask",
        type=str2bool,
        default=True,
        help="Whether to return the mask marking which tokens are green from the watermark detector.",
    )
    parser.add_argument(
        "--window_settings",
        type=str,
        default="20,40,max",  # can also be "20" or "20,40,max"
        help="Comma separated list of window sizes to use for watermark detection. Only used if 'windowed-z-score' is in the evaluation metrics list.",
    )
    parser.add_argument(
        "--run_len_chisqrd_variant",
        type=str,
        default="F_succ_T_runs",
        choices=["F_succ_T_runs", "T_and_F_runs"],
        help="The variant of the run length test to use for watermark detection.",
    )
    parser.add_argument(
        "--run_len_chisqrd_bin_spec",
        type=str,
        default="max_plus_1",
        choices=["max", "max_plus_1"],
        help="The binning specification to use for the run length test.",
    )
    parser.add_argument(
        "--run_len_chisqrd_mask_zeros",
        type=str2bool,
        default=True,
        help="Whether to mask zeros in the run length test.",
    )
    parser.add_argument(
        "--run_len_chisqrd_mask_leading_bins",
        type=int,
        default=0,
        help="The number of leading bins to mask in the run length test.",
    )
    parser.add_argument(
        "--run_len_chisqrd_lambda",
        type=str,
        default="pearson",
        choices=["pearson", "g_test", "cressie_read"],
        help="The lambda_ param to use for the run length test.",
    )
    parser.add_argument(
        "--retrieval_technique",
        type=str,
        default="bm25",
        choices=["bm25", "sim"],
        help="The retrieval technique to use for retrieval detection.",
    )
    parser.add_argument(
        "--retrieval_db_column",
        type=str,
        default="no_wm_output",
        choices=["w_wm_output", "no_wm_output"],
        help="The column to populate the db/index with use for retrieval detection.",
    )
    parser.add_argument(
        "--retrieval_db_load_all_prefixes",
        type=str2bool,
        default=False,
        help="Whether to load all prefixes into the retrieval db, or just the longest for each unique entry.",
    )
    parser.add_argument(
        "--roc_test_stat",
        type=str,
        default="all",
        help="The comma separated list of test statistics to use for the ROC-AUC metric.",
    )
    parser.add_argument(
        "--target_T",
        type=int,
        default=0,
        help="The target generation length to use when dropping rows before ROC-AUC evaluation.",
    )
    parser.add_argument(
        "--lower_tolerance_T",
        type=int,
        default=25,
        help="The lower tolerance to use when dropping rows before ROC-AUC evaluation.",
    )
    parser.add_argument(
        "--upper_tolerance_T",
        type=int,
        default=25,
        help="The upper tolerance to use when dropping rows before ROC-AUC evaluation.",
    )
    parser.add_argument(
        "--filter_by_columns",
        type=str,
        default="all",
        help="The comma separated list of columns to filter by before ROC-AUC evaluation.",
    )
    parser.add_argument(
        "--wandb",
        type=str2bool,
        default=False,
        help="Whether to log to wandb.",
    )
    parser.add_argument(
        "--wandb_project",
        type=str,
        default="lm-watermarking",
        help="The name of the wandb project.",
    )
    parser.add_argument(
        "--wandb_entity",
        type=str,
        default="jwkirchenbauer",
        help="The wandb entity/user for the project.",
    )
    parser.add_argument(
        "--wandb_tags",
        type=str,
        default="",
        help="The comma separated list of tags to add to the wandb run.",
    )
    parser.add_argument(
        "--run_name",
        type=str,
        default="",
        help="The unique name for the run.",
    )
    parser.add_argument(
        "--input_dir",
        type=str,
        default="./input",
        help="The directory containing the input files.",
    )
    parser.add_argument(
        "--output_dir",
        type=str,
        default="",
        help=(
            "The directory in which to write out the dataset after adding the metrics. "
            "If not specified, will use the input_dir. Note, if the output_dir already "
            "contains the metric-enriched file, it will be overwritten :/"
        ),
    )
    parser.add_argument(
        "--overwrite_output_file",
        type=str2bool,
        default=False,
        help="Whether to overwrite the output file if it already exists.",
    )
    parser.add_argument(
        "--limit_rows",
        type=int,
        default=-1,
        help="The number of rows to limit the dataset to. Useful for debugging.",
    )
    parser.add_argument(
        "--concat_rows",
        type=int,
        default=0,
        help="The number of rows to concatenate into a single row. Result is a mangled dataset, be careful",
    )
    parser.add_argument(
        "--shuffle_before_concat",
        type=str2bool,
        default=False,
        help="Whether to shuffle the dataset before concatenating rows.",
    )
    parser.add_argument(
        "--verbose",
        type=str2bool,
        default=None,
        help="Whether to verbosely print things here and there.",
    )
    parser.add_argument(
        "--log_raw_series",
        type=str2bool,
        default=True,
        help="Whether to log the raw series metric data to wandb.",
    )
    parser.add_argument(
        "--log_raw_tabular",
        type=str2bool,
        default=True,
        help="Whether to log the raw tabular metric data to wandb.",
    )
    args = parser.parse_args()

    ###########################################################################
    # Argument validation and conditional setting
    ###########################################################################

    # convert evaluation metrics to list
    assert args.evaluation_metrics, "evaluation_metrics list must be specified"
    args.evaluation_metrics = args.evaluation_metrics.split(",")

    if args.evaluation_metrics == ["all"]:
        all_metrics = SUPPORTED_METRICS
        all_metrics.remove("ppl")  # by default not running this anymore
        all_metrics.remove("detectgpt")  # can't run this with other metrics
        args.evaluation_metrics = all_metrics
    if args.evaluation_metrics == ["all_w_ppl"]:
        args.evaluation_metrics = SUPPORTED_METRICS

    # if no output dir specified, use the input dir
    if args.output_dir == "":
        args.output_dir = args.input_dir

    # check limit_rows
    assert (args.limit_rows == -1) or (
        (args.limit_rows > 0) and isinstance(args.limit_rows, int)
    ), "limit_rows must be -1 or > 0"

    # convert normalizers to list
    if args.normalizers:
        args.normalizers = args.normalizers.split(",")
    else:
        args.normalizers = []

    # convert roc_test_stat to list
    args.roc_test_stat = args.roc_test_stat.split(",")

    if args.roc_test_stat == ["all"]:
        args.roc_test_stat = ROC_TEST_STAT_SUFFIXES

    # convert filter_by_columns to list
    args.filter_by_columns = args.filter_by_columns.split(",")
    if args.filter_by_columns == ["all"]:
        args.filter_by_columns = FILTER_BY_COLUMNS

    # split wandb tags
    if args.wandb_tags != "":
        args.wandb_tags = args.wandb_tags.split(",")
    else:
        args.wandb_tags = []

    # split window settings
    args.window_settings = args.window_settings.split(",")

    main(args)