doc-to-lora/eval.py

import gc
import json
import logging
import os
import re
import string
from argparse import Namespace
from collections import Counter, defaultdict
from collections.abc import Callable
from dataclasses import fields
from functools import partial

import numpy as np
import torch
import yaml
from datasets import disable_caching
from peft import PeftModel
from rouge_score import rouge_scorer
from transformers import (
    EvalPrediction,
    GenerationConfig,
    PreTrainedModel,
    Seq2SeqTrainer,
    Seq2SeqTrainingArguments,
    Trainer,
    set_seed,
)
from transformers.utils import is_liger_kernel_available

from ctx_to_lora.data.definitions import LONGBENCH_E_TASKS, LONGBENCH_TASKS
from ctx_to_lora.data.processing import get_tokenized_dataset
from ctx_to_lora.model_loading import get_model, get_tokenizer
from ctx_to_lora.modeling.hypernet import ModulatedPretrainedModel

logger = logging.getLogger()


# TODO: add negative samples eval (per-sample?)
# TODO: add length info for more fine grain eval


# longbench metrics
def normalize_answer(s):
    """Lower text and remove punctuation, articles and extra whitespace."""

    def remove_articles(text):
        return re.sub(r"\b(a|an|the)\b", " ", text)

    def white_space_fix(text):
        return " ".join(text.split())

    def remove_punc(text):
        exclude = set(string.punctuation)
        return "".join(ch for ch in text if ch not in exclude)

    def lower(text):
        return text.lower()

    return white_space_fix(remove_articles(remove_punc(lower(s))))


def f1_score(prediction: str, ground_truth: str):
    common = Counter(prediction) & Counter(ground_truth)
    num_same = sum(common.values())
    if num_same == 0:
        return 0
    precision = 1.0 * num_same / len(prediction)
    recall = 1.0 * num_same / len(ground_truth)
    f1 = (2 * precision * recall) / (precision + recall)
    return f1


def compute_qa_f1_score(pred_texts: list[str], label_texts: list[str]):
    """
    Word-level F1 score for evaluating question answering systems.
    Order of the words does not matter.
    """
    res = []
    for prediction, label in zip(pred_texts, label_texts):
        normalized_prediction = normalize_answer(prediction)
        normalized_label = normalize_answer(label)

        prediction_words = normalized_prediction.split()
        label_words = normalized_label.split()
        res.append(f1_score(prediction_words, label_words))
    return dict(qa_f1=np.mean(res))


CLOSED_QA_DATASETS = {
    "longbench/narrativeqa",
    "longbench/qasper",
    "longbench/multifieldqa_en",
    "longbench/hotpotqa",
    "longbench/2wikimqa",
    "longbench/musique",
    "hotpot_qa",
    "squad",
    "triviaqa_retrieved",
    "negative_nq",
}

for ds_name in list(CLOSED_QA_DATASETS):
    CLOSED_QA_DATASETS.add(f"{ds_name}_e")


def clear_gpu():
    gc.collect()
    torch.cuda.empty_cache()
    torch.cuda.reset_max_memory_allocated()
    torch.cuda.reset_max_memory_cached()


def add_longbench_tasks(ds_names: list[str]):
    if "longbench" in ds_names:
        ds_names.remove("longbench")
        ds_names += LONGBENCH_TASKS
    if "longbench_e" in ds_names:
        ds_names.remove("longbench_e")
        ds_names += LONGBENCH_E_TASKS


def compute_rouge(pred_texts, label_texts):
    out = defaultdict(list)
    scorer = rouge_scorer.RougeScorer(["rouge1", "rougeL"], use_stemmer=False)
    for pred_text, label_text in zip(pred_texts, label_texts):
        scores = scorer.score(pred_text, label_text)
        for k, v in scores.items():
            out[f"{k}.f1"].append(v.fmeasure)
    for k in out:
        out[k] = np.mean(out[k])
    return out


def compute_per_token_acc(shift_logits, shift_labels, valid_masks):
    indices = torch.where(valid_masks)
    # acc = (shift_logits.argmax(-1) == shift_labels)[indices].float().mean().item()
    # return {"per_token_acc": acc}
    acc = (shift_logits.argmax(-1) == shift_labels)[indices].float()
    return {
        "per_token_accs": acc.flatten().tolist(),
        "n_per_token_accs": valid_masks.sum().item(),
    }


def compute_prefix_matching(shift_logits, shift_labels, valid_masks):
    lengths = valid_masks.sum(dim=1)

    is_wrong = (shift_logits.argmax(-1) != shift_labels) * valid_masks
    is_correct = (shift_logits.argmax(-1) == shift_labels) * valid_masks
    # NOTE: not reliable for multi-turn conversations
    # ie, all tokens in the following user's turn will be correct
    # still monotonically correlate with perf though
    wrong_pos = torch.argmax(is_wrong, dim=1) - torch.argmax(valid_masks, dim=1)
    perf = wrong_pos / lengths

    # if all tokens are correct, set to 1
    perf = torch.where(is_correct.sum(dim=1) == lengths, 1, perf)
    # return {"prefix_matching": perf.mean().item()}
    return {
        "prefix_matchings": perf.tolist(),
        "n_prefix_matchings": valid_masks.shape[0],
    }


@torch.no_grad()
def compute_perplexity(shift_logits, shift_labels, valid_masks):
    loss_fct = torch.nn.CrossEntropyLoss(reduction="none")
    loss = loss_fct(shift_logits.transpose(1, 2), shift_labels)
    loss = (loss * valid_masks).sum(dim=1) / valid_masks.sum(dim=1)
    # perplexity = torch.exp(loss).mean().item()
    # return {"perplexity": perplexity}
    preplexities = torch.exp(loss)
    return {
        "perplexities": preplexities.tolist(),
        "n_perplexities": valid_masks.shape[0],
    }


class Evaluator:
    def __init__(self, metric_fns: list[Callable]):
        self.metric_fns = metric_fns
        self.reset()

    def reset(self):
        self.accum_metrics = defaultdict(list)
        self.count = defaultdict(list)

    def update(self, shift_logits, shift_labels, valid_masks):
        for metric_fn in self.metric_fns:
            metric = metric_fn(shift_logits, shift_labels, valid_masks)
            for k, v in metric.items():
                if k.startswith("n_"):
                    self.count[k[2:]].append(v)
                else:
                    self.accum_metrics[k] += v

    def compute(self):
        # Get result across entire eval set
        result = {
            k: np.sum(v) / np.sum(self.count[k]) for k, v in self.accum_metrics.items()
        }
        # Reset batch statistics
        self.reset()
        return result


@torch.inference_mode()
def compute_metrics(
    eval_pred: EvalPrediction,
    compute_result: bool,
    evaluator: Evaluator,
):
    logits, labels = eval_pred.predictions, eval_pred.label_ids
    shift_logits = logits[..., :-1, :]
    shift_labels = labels[..., 1:]
    valid_masks = torch.where(shift_labels != -100, 1, 0)
    evaluator.update(shift_logits, shift_labels, valid_masks)
    if compute_result:
        return evaluator.compute()


def save_generated_text(samples, output_dir, split):
    os.makedirs(output_dir, exist_ok=True)
    # Create any necessary subdirectories if split contains path separators
    if "/" in split:
        split_dir = os.path.join(output_dir, os.path.dirname(split))
        os.makedirs(split_dir, exist_ok=True)

    with open(f"{output_dir}/{split}_generated_text.jsonl", "w") as f:
        for sample in samples:
            f.write(json.dumps(sample) + "\n")


def decode_test_result(test_dataset, test_result, tokenizer, ctx_tokenizer):
    out = []
    for sample, pred_toks in zip(test_dataset, test_result.predictions):
        # sample and labels are not padded
        # pred_toks are padded though
        d = dict()
        if "labels" in sample:
            start_idx = np.argmax(sample["labels"] != -100)
            label_toks = sample["labels"][start_idx:]
            # labels are padded with -100, so we need to
            # replace them with the pad token id
            label_toks = np.where(
                label_toks == -100, tokenizer.pad_token_id, label_toks
            )
            label_text = tokenizer.decode(label_toks, skip_special_tokens=True)
            d["label"] = label_text.strip()

        # remove the label part
        input_toks = sample["input_ids"][:start_idx]

        gen_toks = pred_toks[np.argmax(pred_toks != tokenizer.pad_token_id) :]
        gen_toks = gen_toks[start_idx:]
        gen_toks = np.where(gen_toks == -100, tokenizer.pad_token_id, gen_toks)

        d["input"] = tokenizer.decode(input_toks, skip_special_tokens=True)
        d["generated"] = tokenizer.decode(gen_toks, skip_special_tokens=True).strip()
        if "ctx_ids" in sample:
            d["context"] = ctx_tokenizer.decode(
                sample["ctx_ids"], skip_special_tokens=True
            )
        for k in sample:
            if k.endswith("_len"):
                d[k] = sample[k].item()
        out.append(d)
    # sort samples by length if possible
    len_key = "ctx_ids_len" if "ctx_ids_len" in sample else "input_ids_len"
    sorted(out, key=lambda x: x[len_key])

    return out


@torch.inference_mode()
def eval_generation(
    eval_trainer, tokenizer, ctx_tokenizer, datasets, split, remove_context, gen_kwargs
):
    if not isinstance(datasets, dict):
        datasets = {"": datasets}
    out = {}
    for ds_name, ds in datasets.items():
        print(f"Evaluating: {ds_name}")
        split_name = f"{split}_{ds_name}" if ds_name else split
        if remove_context:
            split_name += "_no_context"
        eval_result = eval_trainer.predict(
            ds,
            metric_key_prefix=split_name,
            **gen_kwargs,
        )
        decoded_txts = decode_test_result(ds, eval_result, tokenizer, ctx_tokenizer)
        pred_texts = [txt["generated"] for txt in decoded_txts]
        label_texts = [txt["label"] for txt in decoded_txts]
        rouge_metrics = compute_rouge(pred_texts, label_texts)
        for k, v in rouge_metrics.items():
            eval_result.metrics[f"{split_name}_{k}"] = v

        if ds_name in CLOSED_QA_DATASETS:
            print("Computing QA F1 Score")
            qa_f1_metric = compute_qa_f1_score(
                [txt["generated"] for txt in decoded_txts],
                [txt["label"] for txt in decoded_txts],
            )
            for k, v in qa_f1_metric.items():
                eval_result.metrics[f"{split_name}_{k}"] = v

        # Group by input length and compute metrics for each group
        length_bins = [(0, 511), (512, 1023), (1024, 2047), (2048, 4095), (4096, 8192)]
        # Ensure all keys for length metrics are present, even if a bin is empty
        for low, high in length_bins:
            eval_result.metrics[f"{split_name}_rouge1.f1_len_{low}-{high}"] = "None"
            eval_result.metrics[f"{split_name}_rougeL.f1_len_{low}-{high}"] = "None"
            if ds_name in CLOSED_QA_DATASETS:
                eval_result.metrics[f"{split_name}_qa_f1_len_{low}-{high}"] = "None"

        grouped_texts = defaultdict(lambda: {"generated": [], "label": [], "count": 0})
        for txt in decoded_txts:
            len_key = "ctx_ids_len" if "ctx_ids_len" in txt else "input_ids_len"
            input_len = txt[len_key]
            for low, high in length_bins:
                if low <= input_len <= high:
                    group_key = f"{low}-{high}"
                    grouped_texts[group_key]["generated"].append(txt["generated"])
                    grouped_texts[group_key]["label"].append(txt["label"])
                    grouped_texts[group_key]["count"] += 1
                    break

        for group_key, data in grouped_texts.items():
            if data["count"] > 0:
                group_rouge_metrics = compute_rouge(data["generated"], data["label"])
                for k, v in group_rouge_metrics.items():
                    eval_result.metrics[f"{split_name}_{k}_len_{group_key}"] = v

                if ds_name in CLOSED_QA_DATASETS:
                    group_qa_f1_metric = compute_qa_f1_score(
                        data["generated"], data["label"]
                    )
                    for k, v in group_qa_f1_metric.items():
                        eval_result.metrics[f"{split_name}_{k}_len_{group_key}"] = v

        save_generated_text(
            decoded_txts,
            split=split_name,
            output_dir=eval_trainer.args.output_dir,
        )
        out[split_name] = eval_result.metrics
        eval_trainer.log_metrics(split_name, eval_result.metrics)
        eval_trainer.save_metrics(split_name, eval_result.metrics)
        clear_gpu()
    return out


@torch.no_grad()
def eval_teacher_forcing(eval_trainer, datasets, split, remove_context):
    if not isinstance(datasets, dict):
        datasets = {"": datasets}
    out = {}
    for ds_name, ds in datasets.items():
        split_name = f"{split}_{ds_name}" if ds_name else split
        if remove_context:
            split_name += "_no_context"
        metrics = eval_trainer.evaluate(ds, metric_key_prefix=split_name)
        out[split_name] = metrics
        eval_trainer.log_metrics(split_name, metrics)
        eval_trainer.save_metrics(split_name, metrics)
        clear_gpu()
    return out


def train_collator(inp_list, tokenizer):
    # input is a list of tokenized sequences
    padding_kwargs = dict(
        padding=True,
        padding_side="right",
        pad_to_multiple_of=8,
        return_tensors="pt",
    )

    ctx_ids = None
    if "ctx_ids" in inp_list[0]:
        # have to be manual since it has [ctx_len, features] shape
        # pad to the longest ctx_len in the batch
        # which can have a different length from the input_ids, attn_mask, labels

        ctx_ids = [example.pop("ctx_ids") for example in inp_list]
        ctx_ids = torch.nn.utils.rnn.pad_sequence(
            ctx_ids,
            batch_first=True,
            padding_value=0,
        )
        # exotic keys won't be padded, so we need to pad them as well
        ctx_attn_mask = [example.pop("ctx_attn_mask") for example in inp_list]
        ctx_attn_mask = torch.nn.utils.rnn.pad_sequence(
            ctx_attn_mask,
            batch_first=True,
            padding_value=0,
        )

    chat_ids = None
    if "chat_ids" in inp_list[0]:
        chat_ids = [x.pop("chat_ids") for x in inp_list]
        chat_ids = torch.nn.utils.rnn.pad_sequence(
            chat_ids,
            batch_first=True,
            padding_value=0,
        )
        chat_attn_mask = [x.pop("chat_attn_mask") for x in inp_list]
        chat_attn_mask = torch.nn.utils.rnn.pad_sequence(
            chat_attn_mask,
            batch_first=True,
            padding_value=0,
        )
        chat_labels = [x.pop("chat_labels") for x in inp_list]
        chat_labels = torch.nn.utils.rnn.pad_sequence(
            chat_labels,
            batch_first=True,
            padding_value=-100,
        )
        chat_labels = torch.where(chat_attn_mask == 0, -100, chat_labels)

    labels = [x.pop("labels") for x in inp_list]
    padded_seq = tokenizer.pad(inp_list, **padding_kwargs)

    # hacky explicit padding since the labels are not padded by default
    labels = tokenizer.pad({"input_ids": labels}, **padding_kwargs)["input_ids"]
    labels = torch.where(padded_seq["attention_mask"] == 0, -100, labels)
    out = {**padded_seq, "labels": labels}
    if ctx_ids is not None:
        out["ctx_ids"] = ctx_ids
        out["ctx_attn_mask"] = ctx_attn_mask
    if chat_ids is not None:
        out["chat_ids"] = chat_ids
        out["chat_attn_mask"] = chat_attn_mask
        out["chat_labels"] = chat_labels
    return out


def generation_collator(inp_list, tokenizer):
    padding_kwargs = dict(padding=True, padding_side="left", return_tensors="pt")
    input_ids = [x.pop("input_ids") for x in inp_list]
    attn_mask = [x.pop("attention_mask") for x in inp_list]
    labels = [x.pop("labels") for x in inp_list]
    for i, label in enumerate(labels):
        # remove the response tokens
        idx = np.argmax(label != -100)
        input_ids[i] = input_ids[i][:idx]
        attn_mask[i] = attn_mask[i][:idx]
    out = tokenizer.pad(
        {"input_ids": input_ids, "attention_mask": attn_mask}, **padding_kwargs
    )

    # we don't include the labels in the output
    # since we don't want to compute the loss on the labels
    # during generation
    if "ctx_ids" in inp_list[0]:
        # have to be manual since it has [ctx_len, features] shape
        # pad to the longest ctx_len in the batch
        # which can have a different length from the input_ids, attn_mask, labels
        ctx_ids = [example.pop("ctx_ids") for example in inp_list]
        ctx_ids = torch.nn.utils.rnn.pad_sequence(
            ctx_ids,
            batch_first=True,
            padding_value=0,
        )
        # exotic keys won't be padded, so we need to pad them as well
        ctx_attn_mask = [example.pop("ctx_attn_mask") for example in inp_list]
        ctx_attn_mask = torch.nn.utils.rnn.pad_sequence(
            ctx_attn_mask,
            batch_first=True,
            padding_value=0,
        )
        out["ctx_ids"] = ctx_ids
        out["ctx_attn_mask"] = ctx_attn_mask
    return out


def evaluate(
    checkpoint_path,
    model_name_or_path,
    eval_batch_size,
    args,
    split,
    generative,
):
    assert split in ["validation", "test"]
    ctx_name = None
    if model_name_or_path is None:
        state_dict = torch.load(checkpoint_path, weights_only=False)
        ctx_name = state_dict["ctx_encoder_args"].ctx_encoder_model_name_or_path
        model = ModulatedPretrainedModel.from_state_dict(
            state_dict,
            train=False,
            use_flash_attn=True,
        )
        # if generative:
        #     model = model.to(torch.bfloat16)
    else:
        model_kwargs = dict(attn_implementation="flash_attention_2")
        model = get_model(
            model_name_or_path,
            train=False,
            requires_grad=False,
            model_kwargs=model_kwargs,
        )
    # NOTE: there is still some randomness in the eval result
    # despite all the deterministic settings
    if is_liger_kernel_available():
        from liger_kernel.transformers import _apply_liger_kernel_to_instance

        if isinstance(model, ModulatedPretrainedModel):
            print("Applying liger-kernel to ModulatedPretrainedModel")
            if isinstance(model.base_model, PeftModel):
                _apply_liger_kernel_to_instance(model=model.base_model.base_model.model)
            else:
                _apply_liger_kernel_to_instance(model=model.base_model.model)
            if ctx_name is not None:
                print("Applying liger-kernel to ctx_encoder_model")
                _apply_liger_kernel_to_instance(model=model.ctx_encoder.base_model)
        elif isinstance(model, PeftModel):
            print("Applying liger-kernel to PeftModel")
            _apply_liger_kernel_to_instance(model=model.base_model.model)
        elif isinstance(model, PreTrainedModel):
            print("Applying liger-kernel to PretrainedModel")
            _apply_liger_kernel_to_instance(model=model.model)

    tokenizer = get_tokenizer(args.model_name_or_path, train=False)
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token_id = tokenizer.eos_token_id
    base_model = (
        model.base_model if isinstance(model, ModulatedPretrainedModel) else model
    )
    base_model.config.pad_token_id = tokenizer.pad_token_id
    base_model.generation_config.pad_token_id = tokenizer.pad_token_id

    ctx_tokenizer = tokenizer
    if ctx_name:
        ctx_tokenizer = get_tokenizer(ctx_name, train=False)
        if ctx_tokenizer.pad_token_id is None:
            ctx_tokenizer.pad_token_id = ctx_tokenizer.eos_token_id

    tokenizer_kwargs = {}
    ctx_tokenizer_kwargs = {}
    add_ctx_to_chat = (
        not isinstance(model, ModulatedPretrainedModel) and not args.remove_context
    )
    ctx_model_max_len = (
        model.ctx_encoder.config.max_position_embeddings
        if isinstance(model, ModulatedPretrainedModel)
        else None
    )

    _get_tokenized_dataset = partial(
        get_tokenized_dataset,
        base_model_max_len=model.base_model.config.max_position_embeddings,
        tokenizer=tokenizer,
        tokenizer_kwargs=tokenizer_kwargs,
        ctx_model_max_len=ctx_model_max_len,
        ctx_tokenizer=ctx_tokenizer,
        ctx_tokenizer_kwargs=ctx_tokenizer_kwargs,
        add_ctx_to_chat=add_ctx_to_chat,
        add_repeat_prompt=False,
        add_negative_prompt=False,
        use_kl_loss=False,
        set_format="pt",
    )

    datasets = dict()
    ds_names = args.val_ds_names if split == "validation" else args.test_ds_names
    add_longbench_tasks(ds_names)
    for ds_name in ds_names:
        datasets[ds_name] = _get_tokenized_dataset(ds_name, split)
    print(datasets)

    gen_kwargs = dict(
        do_sample=False,
        max_new_tokens=256,  # max_new_tokens=args.max_new_tokens
    )

    eval_trainer_args = {}

    # Copy only necessary attributes from training_args to eval_trainer_args
    seq2seq_training_args_fields = {f.name for f in fields(Seq2SeqTrainingArguments)}
    for attr, value in dict(**vars(args)).items():
        if attr in seq2seq_training_args_fields and not attr.startswith("_"):
            eval_trainer_args[attr] = value

    eval_trainer_args["eval_strategy"] = "no"
    eval_trainer_args["save_strategy"] = "no"
    eval_trainer_args["overwrite_output_dir"] = True
    eval_trainer_args["batch_eval_metrics"] = True
    eval_trainer_args["per_device_eval_batch_size"] = eval_batch_size

    eval_trainer_args = Seq2SeqTrainingArguments(
        **eval_trainer_args,
        predict_with_generate=generative,
        generation_config=GenerationConfig(**gen_kwargs),
    )

    # Seq2SeqTrainer is actually just the same as Trainer
    # (although it uses a different data collator, i.e., explicit prompt/answer separation)
    # it just allows `predict_with_generate`
    # allowing us to compute metrics on the generated outputs
    # no clue why they call this seq2seq...

    print("=" * 80 + "\n" + "Evaluating model..." + "\n" + "=" * 80)
    print(f"checkpoint_path: {checkpoint_path}")

    model.eval()

    collator = generation_collator if generative else train_collator

    trainer_kwargs = {
        "model": model,
        "args": eval_trainer_args,
        "data_collator": partial(collator, tokenizer=tokenizer),
    }
    out = {}
    if not generative:
        trainer_kwargs["compute_metrics"] = partial(
            compute_metrics,
            evaluator=Evaluator(
                [compute_per_token_acc, compute_prefix_matching, compute_perplexity]
            ),
        )
        # this is insane
        # but i don't know why calling trainer.evaluate() on different datasets
        # always gives the same numbers across datasets...
        # spents a few hours on this but couldn't find the reason
        for ds_name, ds in datasets.items():
            eval_trainer = Trainer(**trainer_kwargs)
            metrics = eval_teacher_forcing(
                eval_trainer, {ds_name: ds}, split, args.remove_context
            )
            out.update(metrics)
    else:
        eval_trainer = Seq2SeqTrainer(**trainer_kwargs)
        metrics = eval_generation(
            eval_trainer,
            tokenizer,
            ctx_tokenizer,
            datasets,
            split,
            args.remove_context,
            gen_kwargs,
        )
        out.update(metrics)

    clear_gpu()
    return out


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Evaluate a checkpoint")
    parser.add_argument(
        "--model_name_or_path",
        type=str,
        default=None,
        help="Evaluate a base model from HuggingFace Hub, without loading checkpoint",
    )
    parser.add_argument(
        "--checkpoint_path",
        type=str,
        default=None,
        help="Path to the checkpoint to evaluate",
    )
    parser.add_argument(
        "--split",
        type=str,
        choices=["validation", "test"],
        default="validation",
        help="Which split to evaluate on",
    )
    parser.add_argument(
        "--datasets",
        type=str,
        nargs="+",
        help=(
            "Specific datasets to evaluate on."
            "If not provided, uses default from args.yaml"
        ),
    )
    parser.add_argument(
        "--eval_batch_size",
        type=int,
        default=8,
        help="Eval batch size for teacher forcing",
    )
    parser.add_argument(
        "--eval_batch_size_gen",
        type=int,
        default=32,
        help="Eval batch size for generation",
    )
    parser.add_argument(
        "--remove_context",
        action="store_true",
        help="Remove context when evaluating the base model.",
    )

    cli_args = parser.parse_args()
    # setup_logging(output_dir, debug=os.getenv("DEBUG", False))

    assert bool(cli_args.model_name_or_path) ^ bool(cli_args.checkpoint_path), (
        "Either --model_name_or_path or --checkpoint_path must be provided"
    )

    os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
    os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "true"
    os.environ["FLASH_ATTENTION_DETERMINISTIC"] = "1"
    os.environ["WANDB_MODE"] = "disabled"
    disable_caching()
    set_seed(42)
    torch.use_deterministic_algorithms(True, warn_only=True)
    torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False
    torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
    torch.backends.cudnn.benchmark = False
    torch.backends.cuda.matmul.allow_tf32 = False
    torch.backends.cudnn.allow_tf32 = False

    if cli_args.checkpoint_path:
        checkpoint_dir = "/".join(cli_args.checkpoint_path.split("/")[:-1])
        run_dir = "/".join(cli_args.checkpoint_path.split("/")[:-2])
        cur_it = int(cli_args.checkpoint_path.split("checkpoint-")[1].split("/")[0])
        args = Namespace(**yaml.unsafe_load(open(f"{run_dir}/args.yaml")))
        print(f"checkpoint_path: {cli_args.checkpoint_path}")
        print(f"run_dir: {run_dir}")

        args.output_dir = f"{run_dir}/eval-results-{cur_it}"
        args.logging_dir = f"{run_dir}/eval-results-{cur_it}"
        args.run_name = run_dir.split("/")[-1]
        args.remove_context = False  # modulated model doesn't see ctx by default
    else:
        args = Namespace(
            model_name_or_path=cli_args.model_name_or_path,
            output_dir=f"eval_results/{cli_args.model_name_or_path}",
            logging_dir=f"eval_results/{cli_args.model_name_or_path}",
            run_name=f"eval_results/{cli_args.model_name_or_path}",
            # max_base_len=2**13,
            # max_ctx_len=-1,  # not used
            val_ds_names=[],
            test_ds_names=[],
            remove_context=cli_args.remove_context,
        )

    # Override dataset names if provided via CLI
    if cli_args.datasets:
        if cli_args.split == "validation":
            args.val_ds_names = cli_args.datasets
        else:
            args.test_ds_names = cli_args.datasets

    # evaluate(
    #     cli_args.checkpoint_path,
    #     cli_args.model_name_or_path,
    #     cli_args.eval_batch_size,
    #     args,
    #     split=cli_args.split,
    #     generative=False,
    # )
    evaluate(
        cli_args.checkpoint_path,
        cli_args.model_name_or_path,
        cli_args.eval_batch_size_gen,
        args,
        split=cli_args.split,
        generative=True,
    )