Yyk040316/long-context-icl

1import argparse2import logging3from typing import List, Optional4 5import pandas as pd6from transformers import PreTrainedTokenizerBase,AutoConfig7import numpy as np8 9from transformers import LlamaForCausalLM, AutoTokenizer, AutoModelForCausalLM10 11from datasets_loader import DATASET_NAMES2LOADERS, get_loader12from experiment_manager import ExperimentManager13 14from utils import get_max_n_shots, filter_extremely_long_samples, save_results15import os16import torch17from vllm import LLM18 19import google.generativeai as genai20 21 22 23_logger = logging.getLogger(__name__)24logging.basicConfig(level=logging.INFO, format='%(message)s')25 26os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"27 28def get_dataset(dataset: str, tokenizer: PreTrainedTokenizerBase, task: str,token=None, half_seed=None,) -> (pd.DataFrame, pd.DataFrame, List):29    da = get_loader(dataset,task=task)30    # Filter extremely long samples from both train and test samples:31    #_logger.info("filtering test set:")32    #test_df = filter_extremely_long_samples(da.test_df, tokenizer)33    #_logger.info("filtering train set:")34    #train_df = filter_extremely_long_samples(da.train_df, tokenizer)35    test_df = da.test_df36    train_df = da.train_df37    #判断如果dataset的名字里有Multilingual38    if task == 'multilingual':39        #把datasets的名字用_分隔开，并取最后的部分40        language = da.language41        return test_df, train_df, language42    elif task == 'classification':43        return test_df, train_df, da.labels44    else:45        return test_df, train_df46 47def run_experiment(datasets: List[str], models_path: List[str], output_dir: str,48                        n_shots: List[int], n_runs: int,49                       random_seed: int,  gpu_num: int, task: str,subject:str = None,fp16=False,use_retrieval=False) -> None:50    #print(f"subject:{subject}")51    52    base_output_dir = output_dir53    all_records = []54    for model_path in models_path:55        #clean_model_name = model_path.replace('/', '+').replace(' ', '_')56        #把model的名字用/分隔开，取最后的部分57        clean_model_name = model_path.split('/')[-1]58        59        print(f'* Starting with model: {model_path} ({clean_model_name})')60 61        for dataset in datasets: 62            #clean_dataset_name = dataset.replace('/', '+').replace(' ', '_')63            clean_dataset_name = dataset.split('/')[-1]64            if use_retrieval:65                print('Retrieving examples in-window; renamed dataset to avoid confusion')66                clean_dataset_name = f"{clean_dataset_name}-retrieval"67                print(f"New dataset name: {clean_dataset_name}")68                69            print(f'\t- Running with dataset: {dataset} ({clean_dataset_name})')70            #output_dir = os.path.join(base_output_dir, clean_model_name, clean_dataset_name)71            output_dir = os.path.join(base_output_dir, clean_model_name)72 73            test_df, train_df = None, None74 75            records = []76 77            78                79            output_str = task80            output_path = os.path.join(output_dir, f"{output_str}n_shots_results_{'_'.join([str(i) for i in n_shots])}.npy")81            #nshots_file_name = os.path.join(output_dir, f"nspw={nspw}-n_shots.txt")82 83            # TODO - incorporate n_runs in the caching system, so we can easily add additional runs, without running from scratch (or get different number of runs)84            # TODO - also, the name currently contains the number of windows to have, so it's impossible to add more windows and use cache, just more nspw85            os.makedirs(os.path.dirname(output_path), exist_ok=True)86 87            print(f'Running with {output_path}...')88 89            if 'gemini' in model_path:90 91                """92                genai.configure(api_key='api key',transport='rest')93                94                model = genai.GenerativeModel(model_path)#gemini-1.5-pro95                tokenizer = None96                config = genai.get_model(model_path)97                context_window_size = config.input_token_limit98                """99                model = clean_model_name100                tokenizer = None101                context_window_size = 2000000102            elif 'gpt' in model_path:103 104                model = clean_model_name105                tokenizer = None106                context_window_size = 128000107            108            elif 'claude' in model_path:109                model = clean_model_name110                tokenizer = None111                context_window_size = 200000112            113 114                115            116            #elif 'gpt' in model_path:117 118 119            else: 120                model = LLM(model_path,device="cuda",gpu_memory_utilization=0.9,tensor_parallel_size=int(gpu_num))121                config = AutoConfig.from_pretrained(model_path)122                tokenizer = AutoTokenizer.from_pretrained(model_path)123                if fp16:124                    model.half()125                #context_window_size = tokenizer.model_max_length126                if not hasattr(config, "rope_scaling") or config.rope_scaling is None:127                    context_window_size = config.max_position_embeddings128                else:129                    context_window_size = tokenizer.model_max_length130                    131 132            print('Loaded model')133                134            if test_df is None:135                # lazy loading136                if task == 'multilingual':137                    test_df, train_df, language = get_dataset(dataset, tokenizer,task = task)138                    labels = None139                elif task == 'classification':140                    test_df, train_df, labels = get_dataset(dataset, tokenizer,task = task)141                    print(f"labels:{labels}")142                    language = None143                else:144                    test_df, train_df = get_dataset(dataset, tokenizer,task = task)145                    language = None146                    labels = None147                    if subject is not None:148                        test_df = test_df[test_df['category'] == subject]149                        train_df = train_df[train_df['category'] == subject]150                        print('Filtered dataset')151                print('Loaded dataset')152 153 154            em = ExperimentManager(test_df, train_df, model = model, tokenizer=tokenizer, random_seed=random_seed,labels=labels,155                                   datasets_name = dataset,156                                    #subsample_test_set=subsample_test_set,157                                    context_size=context_window_size, 158                                    use_retrieval=use_retrieval,task = task,language = language,subject=subject,model_name = clean_model_name)159 160            accuracies, predictions = em.run_experiment_across_shots(n_shots, n_runs,context_window_size=context_window_size)  # an ndarry of shape (n_runs, len(n_shots))161                162            save_results(dataset, n_shots, accuracies, predictions, output_path, model, plot_results=False)163 164            rows, cols = accuracies.shape165 166            for i in range(rows):167                for j in range(cols):168                    record = {169                        "n_shots": n_shots[i],170                        "accuracy": accuracies[i][j],171                        "run_num": j,172                    }173                    records.append(record)174 175            # assume output dir already contains the model name176            fname = f"{output_dir}/n_shots_results_seed_{random_seed}.csv"177            pd.DataFrame(records).to_csv(fname, index=False)178            print('---------------------------------------------------')179            print(f'Done running model {model} on dataset {dataset}. You can find the results in {fname}')180            181            all_records.extend([r | {'model': model, 'dataset': dataset} for r in records])  # require python 3.9+182    fname = f"{output_dir}/all_results_seed_{random_seed}.csv"183    pd.DataFrame(all_records).to_csv(fname, index=False)184    print('---------------------------------------------------')185    print(f'Done running all models on all datasets. You can find the results in {fname}')186 187 188 189if __name__ == '__main__':190    parser = argparse.ArgumentParser()191 192    # Datasets and model related arguments193    parser.add_argument('--datasets', nargs='+', 194                        help=f'Name of datasets. Supported datasets: {DATASET_NAMES2LOADERS.keys()}')195    parser.add_argument('--models-path', nargs='+',196                        help='HF model names to use, either gpt2 or LLaMa family models')197    parser.add_argument('--fp16', help="use half precision",198                        action='store_true', default=False)199    200    # Directories, caching, and I/O arguments201    parser.add_argument('--output-dir', help="Directory for saving the results", default='./temp', type=str)202    203    # Evaluation and sampling related arguments204    #parser.add_argument('--subsample-test-set', type=int,205                        #help='Size of test set to use to speed up eval. None means using all test set.')206    parser.add_argument('--random-seed', default=42, type=int)207    parser.add_argument('--n-runs', help="Number of times experiments are repeated for every number of windows",208                        type=int, default=1)209 210    # Windowing related arguments211    #parser.add_argument('-n', '--n-windows', nargs='+', help="Number of parallel context windows", type=int)212    parser.add_argument('--n-shots', nargs='+',213                        help="number of examples to fit in each window (can be multiple items). Use -1 for maximum possible",214                        type=int, required=True)215        216        217 218    parser.add_argument('--use-retrieval', help="apply retrieval method",219                        action='store_true', default=False)220    221    parser.add_argument('--gpu-num', help="tensor_parallel_size=gpu_num",222                         default=1)223    224    parser.add_argument('--task', help="task",225                        default='summarization')226    227    parser.add_argument('--subject',type = str,default=None)228 229    230    args = parser.parse_args()231    232    #print('running with token:', args.token)233    run_experiment(**vars(args))234    # Windowing related arguments235