Yyk040316/long-context-icl

1import logging2import os3from typing import List, Tuple4 5import numpy as np6import pandas as pd7from matplotlib import pyplot as plt8from numpy import typing as npt9from torch import distributed as dist10from transformers import PreTrainedTokenizerBase, LlamaTokenizer, LlamaTokenizerFast11from retriv import SparseRetriever12import re13 14import random15 16from constants import TEXT_BETWEEN_SHOTS17 18_logger = logging.getLogger(__name__)19logging.basicConfig(level=logging.INFO, format='%(message)s')20 21 22def get_max_n_shots(train_df: pd.DataFrame, test_df: pd.DataFrame, tokenizer: PreTrainedTokenizerBase,23                    prompt_size: int) -> int:24    # this is nice info-- let's log this even if we don't need to use it 25    longest_test_prompt = test_df[N_TOKENS].max()26    _logger.info(f"longest_test_prompt = {longest_test_prompt}")27 28    n_tokens_between_shots = n_tokens_in_prompt(tokenizer, TEXT_BETWEEN_SHOTS)29    shot_lengths = train_df[N_TOKENS] + n_tokens_between_shots30    prompt_length_percentile = shot_lengths.quantile(0.9)31    print(f"Median length of demonstration: {shot_lengths.quantile(0.5)}")32    print(f"Mean length of demonstration: {sum(shot_lengths)/len(shot_lengths)}")33 34    max_possible_shots_length = prompt_size - longest_test_prompt35    return int(np.floor(max_possible_shots_length / prompt_length_percentile))36 37 38 39 40def synchronize_examples_across_dfs(df1: pd.DataFrame, df2: pd.DataFrame, comp_column: str = "text"):41    df1 = df1.loc[df1[comp_column].isin(df2[comp_column])]42    df2 = df2.loc[df2[comp_column].isin(df1[comp_column])]43    return df1, df244 45def filter_extremely_long_samples(df: pd.DataFrame, tokenizer: PreTrainedTokenizerBase) -> pd.DataFrame:46    df[N_TOKENS] = df[PROMPTS].map(lambda x: n_tokens_in_prompt(tokenizer, x))47    mask = df[N_TOKENS] <= df[N_TOKENS].quantile(0.99)48    _logger.info(f"filtered {sum(~mask)} from  dataset due to extreme length")49    df = df.loc[mask].copy()50    _logger.info(f"longest remaining prompt according to tokenizer: {df[N_TOKENS].max()}")51    return df52 53 54def n_tokens_in_prompt(tokenizer: PreTrainedTokenizerBase, prompt: str, add_special_tokens=False) -> int:55    return len(tokenizer.encode(prompt, add_special_tokens=add_special_tokens))56 57 58def plot_results_graph(results, dataset_name, n_shots, model='') -> None:59    plt.figure()60    plt.errorbar(n_shots, np.mean(results, axis=1), np.std(results, axis=1), fmt='*')61    plt.xlabel("# shots")62    plt.xticks(n_shots)63    metric = 'Accuracy'64    plt.ylabel(f"{dataset_name} {metric}")65    plt.title(f"{metric} {dataset_name} {model}")66 67 68def load_results(dataset_name: str, output_dir: str, plot=False) -> Tuple[npt.NDArray[float], List[int]]:69    all_results = os.listdir(output_dir)70    results_path = [r for r in all_results if r.startswith(f'{dataset_name}_')]71    if len(results_path) != 1:72        raise ValueError(f"Found {len(results_path)} results!")73    results_path = results_path[0]74    results = np.load(os.path.join(output_dir, results_path))75    n_shots = [int(d) for d in results_path.split('.')[-2].split('_') if d.isdigit()]76    if plot:77        plot_results_graph(results, dataset_name, n_shots)78    return results, n_shots79 80def save_results(dataset: str, n_shots: List[int], results: np.ndarray[int], predictions: List[str], outpath: str,81                 model: str = '', plot_results: bool = True,shuffle = False,noisy = False,reinforce = False) -> None:82    if plot_results:83        plot_results_graph(results, dataset, n_shots, model)84        plt.show()85    if not dist.is_initialized() or dist.get_rank() == 0:86        # in case we use multiple GPUs - we only save one file87        np.save(outpath, results)88        with open(outpath.split(".")[0] + "-outputs.pkl", 'wb') as f:89            import pickle90            pickle.dump(predictions, f)91        clean_name = outpath.split(".")[0].split('/')[-1]92        for num, nshots in enumerate(n_shots):93            if num >= len(predictions):94                break95            for i, rep in enumerate(predictions[num]):96                # need to add id and output columns 97                rep['id'] = rep.index98                rep['n_shots'] = nshots99                if shuffle == True:100                    rep['shuffle_number'] = i101                    with open(os.path.dirname(outpath) + "/" + clean_name.split("n_shots_")[0]+"+n_shots="+str(nshots)+"+shuffle="+str(i)+".csv", 'w',encoding="utf-8") as f:102                        rep.to_csv(f)103 104                elif noisy == True:105                    rep['noisy_level'] = i106                    with open(os.path.dirname(outpath) + "/" + clean_name.split("n_shots_")[0]+"+n_shots="+str(nshots)+"+noisy_level="+str(i)+".csv", 'w',encoding="utf-8") as f:107                        rep.to_csv(f)108                elif reinforce == True:109                    rep['reinforce_number'] = i110                    with open(os.path.dirname(outpath) + "/" + clean_name.split("n_shots_")[0]+"+n_shots="+str(nshots)+"+reinforce_number="+str(i)+".csv", 'w',encoding="utf-8") as f:111                        rep.to_csv(f)112                else:113                    rep['run_number'] = i114                    with open(os.path.dirname(outpath) + "/" + clean_name.split("n_shots_")[0]+"+n_shots="+str(nshots)+"+run="+str(i)+".csv", 'w',encoding="utf-8") as f:115                        rep.to_csv(f)116 117def encode_labels(tokenizer: PreTrainedTokenizerBase, labels: List[str]) -> List[List[int]]:118    if isinstance(tokenizer, LlamaTokenizer):119        # sentence piece - adds a space at the beginning of the sentence120        return [tokenizer.encode(f'{label.lstrip()}', add_special_tokens=False) for label in labels]121 122    return [tokenizer.encode(f' {label.lstrip()}', add_special_tokens=False) for label in labels]123 124 125def encode_stop_seq(tokenizer: PreTrainedTokenizerBase, stop_seq: str) -> int:126    stop_seq_token_id = tokenizer.encode(stop_seq, add_special_tokens=False)127    if isinstance(tokenizer, LlamaTokenizer) or isinstance(tokenizer, LlamaTokenizerFast):128        assert len(stop_seq_token_id) == 2129    else:130        assert len(stop_seq_token_id) == 1131    return stop_seq_token_id[-1]132 133 134def extract_again(text):135    pattern = r"[aA]nswer is \(?([A-J])\)?"136    match = re.search(pattern, text)137    if match:138        return match.group(1)139    else:140        #print("1st answer extract failed\n" + text)141        return extract_final(text)142 143 144def extract_answer(text):145    match = re.search(r'.*[aA]nswer:\s*\(?([A-J])\)?', text)146    if match:147        return match.group(1)148    else:149        #print(" 2nd answer extract failed\n")150        return extract_again(text)151    152 153 154def extract_final(text):155    pattern = r"\(?\b([A-J])\b\)?(?!.*\(?\b([A-J])\b\)?)"156    match = re.search(pattern, text, re.DOTALL)157    if match:158        return match.group(0)159    else:160        #print("failed to extract answer\n")161        return None162    163def extract_answer_math(text):164    """165    从字符串中匹配:166    1. "final answer is" (大小写不敏感)167    2. 任意内容（非贪婪）168    3. 直到 ". I hope" (大小写不敏感)169    170    返回从 "final answer is" 后开始，到 ". I hope" 前结束的部分。171    如果找不到匹配，返回 None。172    """173    pattern = re.compile(174        # (?i) 表示忽略大小写175        r'(?i)final\s*answer\s*is\s*(.*?)\.\s*i\s*hope'176    )177    match = pattern.search(text)178    if match:179        # group(1) 即捕获到的答案内容180        # strip() 用来去掉首尾空格181        return match.group(1).strip()182    #return extract_again_math(text)183    return None184 185 186def extract_again_math(text):187    pattern = r"[aA]nswer\s*:\s*(.+?)(?:\.?\s*[Aa]nswer|$)"188    match = re.search(pattern, text)189    if match:190        return match.group(1)191    else:192        #print("1st answer extract failed\n" + text)193        return extract_final_math(text)194 195 196def extract_final_math(text):197    index = text.find('\\boxed{')198    if index == -1:199        return None200    index += len('\\boxed{')201    brace_count = 1202    content = ''203    while index < len(text):204        char = text[index]205        if char == '{':206            brace_count += 1207        elif char == '}':208            brace_count -= 1209            if brace_count == 0:210                break211        content += char212        index += 1213    return content if content != '' else None214 215def extract_answer_gsm8k(text):216    # 匹配 "####" 后的数字部分（允许包含逗号）217    pattern = r"####\s*([\d,]+)"218    match = re.search(pattern, text)219    if match:220        return match.group(1)  # 返回匹配的数字部分221    else:222        # 如果没有匹配到，则打印提示并返回 None223        # print("No final math result found\n" + text)224        return extract_again_math(text)225 226 227    228 229 230 231def _fix_fracs(string):232    substrs = string.split("\\frac")233    new_str = substrs[0]234    if len(substrs) > 1:235        substrs = substrs[1:]236        for substr in substrs:237            new_str += "\\frac"238            if substr[0] == "{":239                new_str += substr240            else:241                try:242                    assert len(substr) >= 2243                except:244                    return string245                a = substr[0]246                b = substr[1]247                if b != "{":248                    if len(substr) > 2:249                        post_substr = substr[2:]250                        new_str += "{" + a + "}{" + b + "}" + post_substr251                    else:252                        new_str += "{" + a + "}{" + b + "}"253                else:254                    if len(substr) > 2:255                        post_substr = substr[2:]256                        new_str += "{" + a + "}" + b + post_substr257                    else:258                        new_str += "{" + a + "}" + b259    string = new_str260    return string261 262def _fix_a_slash_b(string):263    if len(string.split("/")) != 2:264        return string265    a = string.split("/")[0]266    b = string.split("/")[1]267    try:268        a = int(a)269        b = int(b)270        assert string == "{}/{}".format(a, b)271        new_string = "\\frac{" + str(a) + "}{" + str(b) + "}"272        return new_string273    except:274        return string275 276def _remove_right_units(string):277    # "\\text{ " only ever occurs (at least in the val set) when describing units278    if "\\text{ "in string:279        splits = string.split("\\text{ ")280        assert len(splits) == 2281        return splits[0]282    if "\\text{" in string:283        splits = string.split("\\text{")284        assert len(splits) == 2285        return splits[0]286    else:287        return string288 289def _fix_sqrt(string):290    if "\\sqrt" not in string:291        return string292    splits = string.split("\\sqrt")293    new_string = splits[0] 294    for split in splits[1:]:295        if split[0] != "{":296            a = split[0]297            new_substr = "\\sqrt{" + a + "}" + split[1:]298        else:299            new_substr = "\\sqrt" + split300        new_string += new_substr301    return new_string302 303def _replace_frac(string):304    # 将 \frac{a}{b} 替换为 a/b305    pattern = r'\\frac\{([^{}]+)\}\{([^{}]+)\}'306    repl = r'\1/\2'307    string = re.sub(pattern, repl, string)308    return string309 310def _strip_string(string):311    # linebreaks  312    string = string.replace("\n", "")313    #print(string)314 315    string = string.replace("\(", "")316    string = string.replace("\)", "")317 318    string = string.replace("\\,", "")319    string = string.replace("\,", "")320    string = string.replace(",", "")321 322    # remove inverse spaces323    string = string.replace("\\!", "")324    #print(string)325 326    # replace \\ with \327    string = string.replace("\\\\", "\\")328    #print(string)329 330    # replace tfrac and dfrac with frac331    string = string.replace("tfrac", "frac")332    string = string.replace("dfrac", "frac")333    #print(string)334 335    # remove \left and \right336    string = string.replace("\\left", "")337    string = string.replace("\\right", "")338    #print(string)339    340    # Remove circ (degrees)341    string = string.replace("^{\\circ}", "")342    string = string.replace("^\\circ", "")343 344    # remove dollar signs345    string = string.replace("\\$", "")346    string = string.replace("\$", "")347    string = string.replace("$", "")348    349    # remove units (on the right)350    string = _remove_right_units(string)351 352    # remove percentage353    string = string.replace("\\%", "")354    string = string.replace("\%", "")355 356    # " 0." equivalent to " ." and "{0." equivalent to "{." Alternatively, add "0" if "." is the start of the string357    string = string.replace(" .", " 0.")358    string = string.replace("{.", "{0.")359 360    361    # if empty, return empty string362    if len(string) == 0:363        return string364    if string[0] == ".":365        string = "0" + string366 367    # to consider: get rid of e.g. "k = " or "q = " at beginning368    if len(string.split("=")) == 2:369        if len(string.split("=")[0]) <= 2:370            string = string.split("=")[1]371 372    # fix sqrt3 --> sqrt{3}373    string = _fix_sqrt(string)374 375    # remove spaces376    string = string.replace(" ", "")377 378    # \frac1b or \frac12 --> \frac{1}{b} and \frac{1}{2}, etc. Even works with \frac1{72} (but not \frac{72}1). Also does a/b --> \\frac{a}{b}379    string = _fix_fracs(string)380 381    # manually change 0.5 --> \frac{1}{2}382    if string == "0.5":383        string = "\\frac{1}{2}"384 385    # NOTE: X/Y changed to \frac{X}{Y} in dataset, but in simple cases fix in case the model output is X/Y386    string = _fix_a_slash_b(string)387    string = _replace_frac(string)388 389    #如果string是一个数字390    if string.isdigit():391        #如果是3.0这类的整数但是多了一个.0，去掉.0392        if string[-2:] == ".0":393            string = string[:-2]394 395 396 397    return string398 399def is_equiv(str1, str2, verbose=False):400    if str1 is None and str2 is None:401        print("WARNING: Both None")402        return True403    if str1 is None or str2 is None:404        return False405 406    try:407        ss1 = _strip_string(str1)408        ss2 = _strip_string(str2)409        if verbose:410            print(ss1, ss2)411        return ss1 == ss2412    except:413        return str1 == str2414    415 416def retrieve_context(train_df: pd.DatetimeIndex, index: SparseRetriever, curr_example: str, n_examples: int, split_text, shuffle_seed=None):417    retrieved = index.search(418        query=curr_example,    # What to search for        419        return_docs=False,          # Default value, return the text of the documents420        cutoff=n_examples,                # Default value, number of results to return421    )422    inds = [int(d) for d in retrieved]423    424    if len(inds) < n_examples:425        print(f"WARNING: sampling {n_examples - len(inds)} examples randomly to fill window")426        #从train_df中随机抽取n_examples - len(inds)个样本，并且与inds里已经检索得到的不一样427        inds.extend(train_df['id'].sample(n_examples - len(inds)).loc[~train_df['id'].isin(inds)])428    429    dps = list(train_df.loc[train_df['id'].isin(inds)]['prompt'])430    if shuffle_seed:431        import random432        prev_state = random.getstate()433        random.seed(shuffle_seed)434        random.shuffle(dps)435        random.setstate(prev_state)436        437    text = split_text.join(dps)438    return text439 440def create_retriever(train_df):441    sr = SparseRetriever(442        index_name="training-examples",443        model="bm25",444        min_df=1,445        tokenizer="whitespace",446        stemmer="english",447        stopwords="english",448        do_lowercasing=True,449        do_ampersand_normalization=True,450        do_special_chars_normalization=True,451        do_acronyms_normalization=True,452        do_punctuation_removal=True,453    )454    import random455    filename = f"__temp_index_file_{random.randint(1,5888)}_{random.randint(1,5999)}.csv"456    train_df['id'] = train_df.index457    from pathlib import Path458    import os459    if os.path.exists(filename):460        Path.unlink(Path(filename))461    462    train_df.to_csv(filename)463    sr.index_file(path=filename, 464        show_progress=True,  465        callback=lambda doc: {      # Callback defaults to None.466            "id": doc["id"],467            "text": doc["prompt"]},          468    )469    Path.unlink(Path(filename))470 471    return sr472 473def add_noisy(df,task,noisy_level,noisy_idx,residue_df,labels = None):474    if noisy_level == 0:475        return df476    df_idx = df.index.tolist()477    #先从df_inx里去除noisy_idx那部分478    candidate_idx = [idx for idx in df_idx if idx not in noisy_idx]479 480 481    if task == 'summarization' or task == 'multilingual':482        # 复制 solution 列到 shuffle_solution 列483        df['noisy_solution'] = df['solution']484 485        # 随机选取 noisy_level 个样本的索引486        487        noisy_indices = random.sample(candidate_idx, noisy_level - len(noisy_idx))488 489        all_noisy_indices = noisy_idx + noisy_indices490        #从residue_df里随机选取noisy_level个样本的solution，然后把df里all_noisy_indices对应的solution替换成residue_df里对应的solution491        residue_values = residue_df.loc[random.sample(residue_df.index.tolist(),noisy_level), 'solution'].tolist()492        df.loc[all_noisy_indices, 'noisy_solution'] = residue_values493        494        495        496        """497        original_values = df.loc[noisy_indices, 'solution'].tolist()498        #确保value里没有重复的，否则中断程序499        assert len(original_values) == len(set(original_values)), "Noisy values contain duplicates"500        # 确保完全乱序，直到没有元素保持原位501        if len(noisy_indices) == 1:502            #直接把noisy_indices对应的shuffle_solution替换成除去这个indices对应的solution以外的任意一个值503            shuffled_values = random.sample([v for v in df['solution'].tolist() if v != original_values[0]],1)504            assert shuffled_values[0] != original_values[0], "Shuffled value is the same as original value"505        else:506            iter = 0507            while True:508                shuffled_values = random.sample(original_values, len(original_values))509                iter += 1510                if all(o != s for o, s in zip(original_values, shuffled_values)):511                    break512                if iter > 100:513                    print("WARNING: could not shuffle noisy examples after 100 iterations")514            # 更新 shuffle_solution 列515        """516        #df.loc[noisy_indices, 'solution'] = shuffled_values517 518        519 520        #生成prompt_new列521        df['prompt_new'] = df['problem'] + df['noisy_solution'] + '\n'522 523        return df,all_noisy_indices524    elif task == 'classification':525        df['shuffle_solution'] = df['solution']526        noisy_indices = random.sample(df.index.tolist(), noisy_level)527        original_values = df.loc[noisy_indices, 'solution'].tolist()528        529        values = set(original_values)530 531        if len(values) > 1:532            for i in range(len(original_values)):533                #把值替换成values不同于原值的值534                iter = 0535                while True:536                    new_value = random.choice(list(values))537                    iter += 1538                    if new_value != original_values[i]:539                        break540                    if iter > len(labels):541                        print("WARNING: could not shuffle noisy examples after all labels were tried")542                        assert False543        else:544            #如果只有一个值，就从labels里随机选一个不同于原值的值545            for i in range(len(original_values)):546                iter = 0547                while True:548                    new_value = random.choice(labels)549                    iter += 1550                    if new_value != original_values[i]:551                        break552                    if iter > len(labels):553                        print("WARNING: could not shuffle noisy examples after all labels were tried")554                        assert False555 556        df.loc[noisy_indices, 'shuffle_solution'] = shuffled_values557        df['prompt_new'] = df['problem'] + df['shuffle_solution'] + '\n'558        return df,noisy_indices559    elif task == 'qa':560        # 复制 solution 列到 shuffle_solution 列561        df['shuffle_solution'] = df['solution']562        df['shuffle_answer'] = df['answer']563 564        # 随机选取 noisy_level 个样本的索引565        noisy_indices = random.sample(df.index.tolist(), noisy_level)566        567        original_values = df.loc[noisy_indices, 'solution'].tolist()568        original_answers = df.loc[noisy_indices, 'answer'].tolist()569 570        #确保value里没有重复的，否则中断程序571        assert len(original_values) == len(set(original_values)), "Noisy values contain duplicates"572        # 确保完全乱序，直到没有元素保持原位573        iter = 0574        while True:575            shuffled_indices = random.sample(range(len(original_values)), len(original_values))576            shuffled_values = [original_values[i] for i in shuffled_indices]577            shuffled_answers = [original_answers[i] for i in shuffled_indices]578 579            iter += 1580            # 检查是否满足完全乱序581            if all(original_values[i] != shuffled_values[i] for i in range(len(original_values))):582                break583            if iter > 100:584                print("WARNING: could not shuffle noisy examples after 100 iterations")585        # 更新 shuffle_solution 列586        df.loc[noisy_indices, 'shuffle_solution'] = shuffled_values587        df.loc[noisy_indices, 'shuffle_answer'] = shuffled_answers588 589        #生成prompt_new列590        df['prompt_new'] = df['problem'] + 'Solution:\n' + df['shuffle_solution'] + '\n' + "Answer: " +  "(" + df['shuffle_answer'].apply(lambda x: x["answer"].rstrip()) + ")" + '\n'591 592        return df,noisy_indices593 594        595        596