run_eval.py

import os
os.environ["TOKENIZERS_PARALLELISM"] = "false"
import argparse
import torch
import json
from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
from transformers import AutoTokenizer
from collections import defaultdict
from tqdm import tqdm
import numpy as np
import logging
import random
import time
from torch.utils.data import DataLoader

from arguments import parse_arguments

logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s - %(message)s',
                    datefmt='%m/%d/%Y %H:%M:%S')
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)


def load_model(args):
    model_processor = AutoTokenizer.from_pretrained(
        "EleutherAI/gpt-neox-20b", 
        clean_up_tokenization_spaces=True
    )
    model = MambaLMHeadModel.from_pretrained(
        args.model, 
        device=f'cuda:{args.device}', 
        dtype=torch.bfloat16
    )
    model_name = [args.model.split("/")[-1]]

    if args.config is not None:
        from mamba_ssm.mapping import load_dt_remapping_from_config
        with open(args.config, 'r') as f:
            remap_config = json.load(f)
        
        load_dt_remapping_from_config(model, remap_config, f'cuda:{args.device}')
        model_name.extend(args.config.split('/')[-1].split('.')[:-1])
        model_name = '.'.join(model_name)
    else:
        model_name = model_name[0]+"_vanilla"

    return model_processor, model, model_name


def my_longbench(args=None, only_eval=False):
    from tasks.longbench import longbench_pred, scorer, scorer_e

    model_processor, model, model_name = load_model(args)

    if not only_eval:
        longbench_pred(args, model, model_processor, model_name)
    
    scores = dict()

    if args.long_eval_task == "e":
        path = f"pred_longbench_e/{model_name}/"
    else:
        path = f"pred_longbench/{model_name}/"
    
    all_files = os.listdir(path)
    for filename in all_files:

        if not filename.endswith("jsonl"):
            continue
        
        predictions, answers, lengths = [], [], []
        dataset = filename.split('.')[0]
        
        with open(f"{path}{filename}", "r", encoding="utf-8") as f:
            for line in f:
                data = json.loads(line)
                predictions.append(data["pred"])
                answers.append(data["answers"])
                all_classes = data["all_classes"]
                if "length" in data:
                    lengths.append(data["length"])
        
        if args.long_eval_task == "e":
            score = scorer_e(dataset, predictions, answers, lengths, all_classes)
        else:
            score = scorer(dataset, predictions, answers, all_classes)
        
        scores[dataset] = score

    if args.long_eval_task == "e":
        out_path = f"pred_longbench_e/{model_name}/result.json"
    elif args.long_eval_task == "c":
        out_path = f"pred_longbench/{model_name}/result_c.json"
    else:
        out_path = f"pred_longbench/{model_name}/result.json"
    with open(out_path, "w") as f:
        json.dump(scores, f, ensure_ascii=False, indent=4)


def special_input_ppl(args):
    tokenizer, model, model_name = load_model(args)

    with open(args.sample_path, 'r', encoding='utf-8') as file:
        inputs = file.read()
    inputs = tokenizer(inputs, return_tensors="pt").to(f"cuda:{args.device}")

    ppls = []
    perplexities = {}
    max_amount_of_windows = 50
    window_sizes = [1e3, 4e3, 10e3, 20e3, 40e3, 60e3, 80e3, 100e3, 120e3]
    
    for window_size in window_sizes:
        window_size = int(window_size)
        nlls = []
        seq_len = inputs['input_ids'].size(1)
        
        if seq_len < window_size:
            print(f'skipping sample seq_len = {seq_len//1000}K < window_size = {window_size//1000}K')
            continue
        
        stride = max(10, (seq_len-window_size)//max_amount_of_windows)

        for begin_loc in range(0, seq_len-window_size, stride):
            end_loc = begin_loc + window_size
            input_ids = inputs['input_ids'][:, begin_loc:end_loc].to(f"cuda:{args.device}")
            target_ids = input_ids.clone()

            with torch.no_grad():
                target_ids = target_ids[:, -100:]
                outputs = model(input_ids, num_last_tokens=100+1)
                logits = outputs.logits
                ce_loss = torch.nn.CrossEntropyLoss()
                neg_log_likelihood = ce_loss(logits.squeeze()[:-1], target_ids.squeeze())
            nlls.append(neg_log_likelihood)
            
        ppl = torch.exp(torch.stack(nlls).mean()).cpu().to(torch.float)
        print(f'{int(window_size/1e3)}k calculated perplexity: {ppl:.2f}')
        ppls.append(ppl)
        perplexities[f'{int(window_size/1e3)}k'] = f'{ppl:.4f}'
        
    avg_ppl = torch.stack(ppls).mean().item()
    perplexities["average"] = f'{avg_ppl:.4f}'
    
    output_dir = f'pred_ppl/{args.sample_path.split(".txt")[0]}'
    os.makedirs(output_dir, exist_ok=True)
    with open(f'{output_dir}/{model_name}.json', "w") as f:
        json.dump(perplexities, f, ensure_ascii=False, indent=4)


def my_helmet(args):
    from tasks.helmet import run_test

    tokenizer, model, model_name = load_model(args)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
        tokenizer.pad_token_id = tokenizer.eos_token_id
    tokenizer.truncation_side = "left"
    tokenizer.padding_side = "left"

    stop_token_ids = tokenizer.eos_token_id
    stop_token_ids = [stop_token_ids] if not isinstance(stop_token_ids, list) else stop_token_ids
    if args.stop_newline:
        stop = list(set(["\n", "Ċ", "ĊĊ", "<0x0A>"]))
        stop_token_ids = list(set([tokenizer.convert_tokens_to_ids(stop_token) for stop_token in stop] + stop_token_ids))
        stop_token_ids = [x for x in stop_token_ids if x is not None]

    datasets = args.datasets.split(",")
    test_files = args.test_files.split(",")
    demo_files = args.demo_files.split(",")
    max_lengths = ([int(args.input_max_length)] * len(datasets)) if isinstance(args.input_max_length, int) or len(args.input_max_length.split(",")) == 1 else [int(l) for l in args.input_max_length.split(",")]
    gen_lengths = ([int(args.generation_max_length)] * len(datasets)) if isinstance(args.generation_max_length, int) or len(args.generation_max_length.split(",")) == 1 else [int(l) for l in args.generation_max_length.split(",")]
    assert len(test_files) == len(demo_files)

    for dataset, test_file, demo_file, max_length, gen_length in zip(datasets, test_files, demo_files, max_lengths, gen_lengths):
        args.datasets = dataset
        args.test_files = test_file
        args.demo_files = demo_file

        try: 
            output_path = run_test(args, model, tokenizer, dataset, test_file, demo_file, max_length, gen_length, stop_token_ids, model_name, logger)
            if "alce" in dataset and not args.count_tokens and (not os.path.exists(output_path+".score") or args.overwrite):
                from tasks import helmet_eval_alce as eval_alce
                logger.info("running eval_alce.py...")
                cli_args = ["--f", output_path]
                if not "nocite" in dataset:
                    cli_args.append("--citations")
                if "asqa" in dataset:
                    cli_args.append("--mauve")
                elif "eli5" in dataset:
                    cli_args += ["mauve", "--claims_nli"]
                eval_alce.main(cli_args)

        except Exception as e:
            logger.exception(e)
            logger.error(f"Error in {dataset}, continuing...")
            if args.debug:
                raise e


if __name__ == '__main__':
    args = parse_arguments()
    args.remove_last_space = not args.keep_last_space
    
    if args.long_eval_task != "no":
        print("Evaluating on LongBench datasets...")
        my_longbench(args, only_eval=args.only_eval)
    if args.helmet_config is not None:
        print("Evaluating on Helemt datasets...")
        my_helmet(args)
    if args.ppl:
        print("Calculating perplexity on special input...")
        special_input_ppl(args)