First add.

Long_LLM/activation_beacon/new/main/eval_passkey.py

+import re
+import os
+import json
+import torch
+import datasets
+import numpy as np
+
+import pandas as pd
+import seaborn as sns
+import matplotlib.pyplot as plt
+from matplotlib.colors import LinearSegmentedColormap
+
+
+from typing import List, Optional
+from tqdm import tqdm
+from fuzzywuzzy import fuzz
+from accelerate import Accelerator
+from transformers import HfArgumentParser
+from transformers.utils import logging
+from dataclasses import dataclass, field, asdict
+
+from src import ModelArgs, DefaultDataCollator, FileLogger, get_model_and_tokenizer, makedirs, apply_chat_template
+
+logger = logging.get_logger(__name__)
+
+
+@dataclass
+class Args(ModelArgs):
+    output_dir: str = field(
+        default="data/results/passkey/",
+        metadata={'help': 'The base directory for saving results and logs.'}
+    )
+    result_dir: Optional[str] = field(
+        default=None,
+        metadata={'help': 'The directory relative to output_dir for saving results.'}
+    )
+
+    min_length: int = field(
+        default=8192,
+        metadata={'help': 'Minimum context length in evaluation.'}
+    )
+    max_length: int = field(
+        default=131072,
+        metadata={'help': 'Maximum context length in evaluation.'}
+    )
+    num_length_interval: int = field(
+        default=20,
+        metadata={'help': 'Number of invervals between min_length and max_length.'}
+    )
+    test_length: List[int] = field(
+        default=None,
+        metadata={'help': 'Specified evaluation lengths.'}
+    )
+
+    min_depth: float = field(
+        default=0,
+        metadata={'help': 'Minimum pass key depth in the context.'}
+    )
+    max_depth: float = field(
+        default=100,
+        metadata={'help': 'Maximum pass key depth in the context.'}
+    )
+    num_depth_interval: int = field(
+        default=10,
+        metadata={'help': 'Number of invervals between min_depth and max_depth.'}
+    )
+    test_depth: List[int] = field(
+        default=None,
+        metadata={'help': 'Specified evaluation depths.'}
+    )
+
+    passkey_length: int = field(
+        default=5,
+        metadata={'help': 'How many numbers are in the passkey?'}
+    )
+    seed: int = field(
+        default=123,
+        metadata={'help': 'Random seed.'}
+    )
+
+    do_sample: bool = False
+    max_new_tokens: int = 50
+
+
+def generate_sample(tokenizer, chat_template, context_length, passkey_depth, passkey_length, rng:np.random.Generator=np.random.default_rng(42)):
+    passkey = str(rng.integers(10**(passkey_length - 1), 10**passkey_length))
+    description = "There is an important infomation hidden in the following context. Find the information and memorize it. I will quiz you about the important information there.\n"
+    noises = "The grass is green. The sky is blue. The sun is yellow. Here we go. There and back again." * (context_length // 10)
+    information = f"\n\nThe pass key is {passkey}. Remember it. {passkey} is the pass key.\n\n"
+    prompt = "\n\nWhat is the pass key?"
+
+    # these inputs are used only once
+    description_input_ids = tokenizer.encode(description, add_special_tokens=False)
+    information_input_ids = tokenizer.encode(information, add_special_tokens=False)
+    prompt_input_ids = tokenizer.encode(prompt, add_special_tokens=False)
+
+    description_length = len(description_input_ids)
+    information_length = len(information_input_ids)
+    prompt_length = len(prompt_input_ids)
+
+    # must leave room for information and prompt
+    minimum_pos = description_length
+    maximum_pos = context_length - prompt_length - information_length - 1
+    if minimum_pos > context_length or maximum_pos < 0:
+        raise ValueError(f"The length {context_length} is too small. Please increase interval!")
+
+    passkey_pos = minimum_pos + round((maximum_pos - minimum_pos) * passkey_depth / 100)
+
+    # DEBUG
+    # information_pos = description_length
+    # information_pos = rng.integers(minimum_pos, min(maximum_pos, 1000))
+    # information_pos = rng.integers(1024, min(maximum_pos, 2000))
+
+    prefix_noise = tokenizer.encode(noises, max_length=passkey_pos - description_length, truncation=True, add_special_tokens=False)
+    suffix_noise = tokenizer.encode(noises, max_length=context_length - passkey_pos - information_length - prompt_length, truncation=True, add_special_tokens=False)
+
+    input_ids = sum([description_input_ids, prefix_noise, information_input_ids, suffix_noise, prompt_input_ids], [])
+    inputs = tokenizer.decode(input_ids)
+
+    inputs = apply_chat_template(chat_template, messages=[{'role': 'user', 'content': inputs}], tokenizer=tokenizer, add_generation_prompt=True).raw
+
+    return inputs, prompt, passkey
+
+
+@torch.no_grad()
+def main():
+    parser = HfArgumentParser([Args])
+    args: Args = parser.parse_args_into_dataclasses()[0]
+
+    accelerator = Accelerator(cpu=args.cpu)
+    model, tokenizer = get_model_and_tokenizer(args, device=accelerator.device)
+
+    if args.test_length is None:
+        test_lengths = np.linspace(args.min_length, args.max_length, args.num_length_interval, endpoint=True).astype(int).tolist()
+    else:
+        test_lengths = args.test_length
+
+    if args.test_depth is None:
+        test_depths = np.linspace(args.min_depth, args.max_depth, args.num_depth_interval, endpoint=True).astype(int).tolist()
+    else:
+        test_depths = args.test_depth
+
+    rng_state = np.random.default_rng(args.seed)
+
+    all_inputs = []
+    for length in tqdm(test_lengths, desc="Constructing Data"):
+        for depth in test_depths:
+            inputs, prompt, passkey = generate_sample(
+                tokenizer=tokenizer, 
+                chat_template=args.chat_template, 
+                context_length=length, 
+                passkey_depth=depth, 
+                passkey_length=args.passkey_length,
+                rng=rng_state
+            )
+            all_inputs.append({'inputs': inputs, 'prompt': prompt, 'passkey': passkey, 'length': length, 'depth': depth})
+
+    dataset = datasets.Dataset.from_list(all_inputs)
+    dataloader = torch.utils.data.DataLoader(
+        # length and depth are useless in forward computation
+        dataset.remove_columns(['length', 'depth', 'passkey']), 
+        batch_size=args.batch_size, 
+        collate_fn=DefaultDataCollator(tokenizer),
+        pin_memory=not args.cpu,
+    )
+
+    # NOTE: prepare dataloader so the data moves to GPU automatically
+    dataloader = accelerator.prepare(dataloader)
+
+    accelerator.wait_for_everyone()
+
+    all_outputs = []
+
+    for x in tqdm(dataloader, desc="Evaluating"):
+        prompt = x.pop("prompt")
+        inputs = x.pop("inputs")
+        # TODO: retrieval
+
+        # NOTE: important to reset memory for every batch
+        if hasattr(model, "memory"):
+            model.memory.reset()
+
+        inputs = tokenizer(inputs, return_tensors="pt").to(model.device)
+
+        output = model.generate(**inputs)
+
+        if isinstance(output, torch.Tensor):
+            # 1, max_new_tokens
+            output = output[:, inputs['input_ids'].shape[1]:]
+            output = tokenizer.batch_decode(output, skip_special_tokens=True)
+        elif isinstance(output, list):
+            pass
+
+        if accelerator.num_processes > 1:
+            output = accelerator.gather_for_metrics(output)
+
+        all_outputs.extend(output)
+
+    if accelerator.process_index == 0:
+        accuracy = {l: {d: [] for d in test_depths} for l in test_lengths}
+        fuzzy_score = {l: {d: [] for d in test_depths} for l in test_lengths}
+        results = {l: {d: [] for d in test_depths} for l in test_lengths}
+
+        for l, d, p, o in zip(dataset['length'], dataset['depth'], dataset['passkey'], all_outputs):
+            # extract numbers
+            o = re.search("\d+", o)
+            if o:
+                o = o.group()
+            else:
+                o = ""
+            results[l][d].append({'target': p, 'prediction': o})
+
+            acc = float(p == o)
+            score = round(fuzz.ratio(o, p) / 100, 2)
+
+            accuracy[l][d].append(acc)
+            fuzzy_score[l][d].append(score)
+
+        for l, lv in accuracy.items():
+            for d, dv in lv.items():
+                accuracy[l][d] = round(sum(dv) / len(dv), 2)
+
+        for l, lv in fuzzy_score.items():
+            for d, dv in lv.items():
+                fuzzy_score[l][d] = round(sum(dv) / len(dv), 2)
+        
+        result_dir = os.path.join(args.output_dir, args.result_dir)
+        with open(makedirs(os.path.join(result_dir, "results.json")), "w", encoding='utf-8') as f:
+            json.dump(results, f)
+        # also save config
+        args.save(os.path.join(result_dir, "config.json"))
+
+        metrics = {'accuracy': accuracy, 'fuzz': fuzzy_score}
+        file_logger = FileLogger(makedirs(os.path.join(args.output_dir, "metrics.log")))
+        file_logger.log(metrics, Args=asdict(args))
+
+        for metric_key, metric_value in metrics.items():
+            # Copied from https://github.com/gkamradt/LLMTest_NeedleInAHaystack/blob/main/viz/CreateVizFromLLMTesting.ipynb
+            cmap = LinearSegmentedColormap.from_list("custom_cmap", ["#F0496E", "#EBB839", "#0CD79F"])
+            # Create the heatmap with better aesthetics
+            sns.set(rc={"figure.figsize": (17.5, 8), "axes.titlesize":14, "axes.labelsize":12}, style="whitegrid", palette="colorblind")
+            data = pd.DataFrame(metric_value)
+            ax = sns.heatmap(
+                data,
+                cmap=cmap,
+                vmin=0,
+                vmax=1,
+                fmt="g",
+                linewidth=.5,
+            )
+            cbar = ax.collections[0].colorbar
+            cbar.set_label(metric_key, size=14)
+
+            # More aesthetics
+            plt.title('Passkey Retrieval')  # Adds a title
+            plt.xlabel('Context Length', fontsize=14)  # X-axis label
+            plt.ylabel('Depth Percent', fontsize=14)  # Y-axis label
+            plt.xticks(rotation=45, fontsize=10)  # Rotates the x-axis labels to prevent overlap
+            plt.yticks(rotation=0, fontsize=10)  # Ensures the y-axis labels are horizontal
+            plt.tight_layout()  # Fits everything neatly into the figure area
+            # save to result_dir
+            plt.savefig(os.path.join(result_dir, f"{metric_key}.png"), format='png', bbox_inches='tight')
+            plt.close()
+
+
+if __name__ == "__main__":
+    main()

Long_LLM/activation_beacon/new/main/eval_topic.py

+# modified based on https://github.com/DachengLi1/LongChat/blob/longeval/longeval/eval.py
+
+import os
+import json
+import torch
+import datasets
+import numpy as np
+from tqdm import tqdm
+from functools import partial
+from typing import List, Optional
+from accelerate import Accelerator
+from transformers import HfArgumentParser
+from transformers.utils import logging
+from torch.utils.data import DataLoader
+from dataclasses import dataclass, field, asdict
+from collections import defaultdict
+
+from src import ModelArgs, DefaultDataCollator, FileLogger, get_model_and_tokenizer, makedirs, split_file_dir_name_ext, apply_chat_template
+from .longbench_utils import qa_f1_score
+
+logger = logging.get_logger(__name__)
+
+
+@dataclass
+class Args(ModelArgs):
+    eval_data: str = field(
+        default="long-llm:longeval/topic_retrieval.json",
+        metadata={'help': 'Evaluation json data.'}
+    )
+    output_dir: str = field(
+        default="data/results/topic_retrieval/",
+        metadata={'help': 'The base directory for saving results and logs.'}
+    )
+    result_dir: Optional[str] = field(
+        default=None,
+        metadata={'help': 'The directory relative to output_dir for saving results.'}
+    )
+    num_topic: List[int] = field(
+        default_factory=lambda: [5, 10, 15, 20, 25, 30, 40, 50, 60, 70],
+        metadata={'help': 'How many topics to in the conversation?'}
+    )
+    adapt_window: bool = field(
+        default=False,
+        metadata={'help': 'Dynamically change the beacon window so that the input is always compressed?'}
+    )
+    target_topic: str = field(
+        default="first",
+        metadata={'help': 'Which topic to evaluate?'}
+    )
+
+    do_sample: bool = False
+    max_new_tokens: int = 50
+
+def process_topic_retrieval(data, tokenizer, chat_template, num_topic, target_topic):
+    outputs = {'input_ids': [], 'attention_mask': [], 'target': [], 'length': [], 'num': []}
+    
+    for context, question, topics, num in zip(data['context'], data['question'], data['topics'], data['num_topics']):
+        # filter out samples that do not have proper number of topics/lines
+        if num not in num_topic:
+            continue
+
+        if num == 1:
+            context = context.split(" \n USER: Great, this is the end of our discussion")[0]
+            context = context + " Now the record ends."
+
+        if target_topic == "first":
+            question = f"What is the first topic we have discussed? Only give me the topic name. Do not summarize yourself."
+            target = topics[0]
+        elif target_topic == "random":
+            target_idx = np.random.randint(0, num)
+            question = f"What is the No.{target_idx} topic we have discussed? Only give me the topic name. Do not summarize yourself."
+            target = topics[target_idx]
+        else:
+            raise NotImplementedError
+
+        prompt = " ".join([context, question])
+        # the question always asks for the first topic
+
+        encoded = apply_chat_template(chat_template, [{'role': 'user', 'content': prompt}], tokenizer=tokenizer, add_generation_prompt=True).encoded
+
+        encoded["target"] = target
+        encoded["length"] = len(encoded.input_ids)
+        encoded["num"] = num
+
+        for k, v in encoded.items():
+            if k in outputs:
+                outputs[k].append(v)
+
+    return outputs
+
+
+@torch.no_grad()
+def main():
+    parser = HfArgumentParser([Args])
+    args: Args = parser.parse_args_into_dataclasses()[0]
+
+    accelerator = Accelerator(cpu=args.cpu)
+    model, tokenizer = get_model_and_tokenizer(args, device=accelerator.device)
+
+    with accelerator.main_process_first():
+        process_fn = partial(process_topic_retrieval,
+            tokenizer=tokenizer,
+            chat_template=args.chat_template,
+            num_topic=args.num_topic,
+            target_topic=args.target_topic,
+        )
+
+        raw_dataset = datasets.load_dataset("json", data_files=args.eval_data, cache_dir=args.dataset_cache_dir, split="train")
+        dataset = raw_dataset.map(process_fn, batched=True, num_proc=32, remove_columns=raw_dataset.column_names)
+        # group instances of the same number of topics together, so that their lengths are approximately equal
+        groupby_dataset = dataset.to_pandas().groupby("num")
+
+    data_collator = DefaultDataCollator(tokenizer=tokenizer)
+    
+    accuracy = {}
+    f1_score = {}
+    results = defaultdict(list)
+    # used for adapt_window
+    if args.adapt_window:
+        beacon_window = getattr(model.config, "beacon_window", None)
+
+    for num, dataset in groupby_dataset:
+        dataset = datasets.Dataset.from_pandas(groupby_dataset.get_group(num), preserve_index=False)
+        all_targets = dataset["target"]
+        # remove unnecessary columns
+        dataset = dataset.remove_columns(["target", "num"])
+
+        dataloader = DataLoader(
+            dataset, 
+            batch_size=args.batch_size, 
+            collate_fn=data_collator,
+            # only pin memory when no gpu
+            pin_memory=not args.cpu,
+        )
+
+        # NOTE: prepare dataloader so the data moves to GPU automatically
+        dataloader = accelerator.prepare(dataloader)
+
+        all_lengths = []
+        all_outputs = []
+        for i, x in enumerate(tqdm(dataloader, desc=f"Evaluating {num} Topics")):
+            # NOTE: important to reset memory for every batch
+            if hasattr(model, "memory"):
+                if args.adapt_window:
+                    length = x['length'][0].item()
+                    if length < beacon_window:
+                        beacon_window = (length // 256) * 256
+                        beacon_stride = beacon_window
+                        model.memory.set(
+                            beacon_window=beacon_window,
+                            beacon_stride=beacon_stride,
+                        )
+
+                model.memory.reset()
+
+            length = x.pop("length").tolist()
+
+            output = model.generate(**x)
+
+            if isinstance(output, torch.Tensor):
+                # 1, max_new_tokens
+                output = output[:, x['input_ids'].shape[1]:]
+                output = tokenizer.batch_decode(output, skip_special_tokens=True)
+            elif isinstance(output, list):
+                pass
+
+            if accelerator.num_processes > 1:
+                output = accelerator.gather_for_metrics(output)
+                length = accelerator.gather_for_metrics(length)
+            
+            all_outputs.extend(output)
+            all_lengths.extend(length)
+
+        length = int(sum(all_lengths) / len(all_lengths))
+
+        acc = 0
+        f1 = 0
+        for output, target in zip(all_outputs, all_targets):
+            if target.lower() in output.lower():
+                acc += 1
+            else:
+                acc += 0
+            f1 += qa_f1_score(output, target)
+            results[length].append({"target": target, "prediction": output})
+
+        acc /= len(all_outputs)
+        f1 /= len(all_outputs)
+
+        accuracy[length] = acc
+        f1_score[length] = round(f1, 4)
+    
+    if accelerator.process_index == 0:
+        result_dir = os.path.join(args.output_dir, args.result_dir) if args.result_dir is not None else args.output_dir
+        with open(makedirs(os.path.join(result_dir, "results.json")), "w", encoding='utf-8') as f:
+            json.dump(results, f)
+        # also save config
+        args.save(os.path.join(result_dir, "config.json"))
+
+        file_logger = FileLogger(makedirs(os.path.join(args.output_dir, "metrics.log")))
+        file_logger.log({'accuracy': accuracy, 'f1': f1_score}, Args=asdict(args))
+
+
+if __name__ == "__main__":
+    main()

Long_LLM/activation_beacon/new/main/infbench_utils.py

+import json
+import re
+import string
+
+from pathlib import Path
+from collections import Counter, defaultdict
+from tqdm import tqdm
+from rouge import Rouge
+
+
+
+def normalize_answer(s: str) -> str:
+    """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 normalize_zh_answer(s: str) -> str:
+    """Chinese version. Lower text and remove punctuation, extra whitespace."""
+
+    def white_space_fix(text):
+        return "".join(text.split())
+
+    def remove_punc(text):
+        cn_punctuation = "！？｡。＂＃＄％＆＇（）＊＋，－／：；＜＝＞＠［＼］＾＿｀｛｜｝～｟｠｢｣､、〃》「」『』【】〔〕〖〗〘〙〚〛〜〝〞〟〰〾〿–—‘’‛“”„‟…‧﹏."  # noqa
+        all_punctuation = set(string.punctuation + cn_punctuation)
+        return "".join(ch for ch in text if ch not in all_punctuation)
+
+    def lower(text):
+        return text.lower()
+
+    return white_space_fix(remove_punc(lower(s)))
+
+
+def f1_score(prediction, ground_truth) -> tuple[float, float, float]:
+    common = Counter(prediction) & Counter(ground_truth)
+    num_same = sum(common.values())
+    if num_same == 0:
+        return 0, 0, 0
+    precision = 1.0 * num_same / len(prediction)
+    recall = 1.0 * num_same / len(ground_truth)
+    f1 = (2 * precision * recall) / (precision + recall)
+    return f1, precision, recall
+
+
+def qa_f1_score(pred: str, ground_truths) -> float:
+    """Computes the F1, recall, and precision."""
+    f1 = 0
+    prec = 0
+    recall = 0
+    for ground_truth in ground_truths:
+        normalized_prediction = normalize_answer(pred)
+        normalized_ground_truth = normalize_answer(ground_truth)
+
+        prediction_tokens = normalized_prediction.split()
+        ground_truth_tokens = normalized_ground_truth.split()
+        scores = f1_score(prediction_tokens, ground_truth_tokens)
+        this_f1, this_prec, this_recall = scores
+        f1 = max(f1, this_f1)
+        prec = max(prec, this_prec)
+        recall = max(recall, this_recall)
+    return f1
+
+
+def qa_f1_score_zh(pred: str, ground_truths: list[str]) -> float:
+    """
+    QA F1 score for chinese.
+    """
+    f1 = 0
+    prec = 0
+    recall = 0
+    for ground_truth in ground_truths:
+        norm_pred = normalize_zh_answer(pred)
+        norm_label = normalize_zh_answer(ground_truth)
+
+        # One character one token.
+        pred_tokens = list(norm_pred)
+        label_tokens = list(norm_label)
+        scores = f1_score(pred_tokens, label_tokens)
+        this_f1, this_prec, this_recall = scores
+        f1 = max(f1, this_f1)
+        prec = max(prec, this_prec)
+        recall = max(recall, this_recall)
+    return f1
+
+
+def load_json(fname):
+    return json.load(open(fname))
+
+
+def iter_jsonl(fname, cnt=None):
+    i = 0
+    with open(fname, "r", encoding="utf8") as fin:
+        for line in fin:
+            if line.strip() == "":  # Skip empty lines
+                continue
+            if i == cnt:
+                break
+            if line.strip() == "":  # Skip empty lines
+                continue
+            yield json.loads(line)
+            i += 1
+
+def first_int_match(prediction):
+    pred_list = re.split("[^0-9]", prediction)
+    pred_value = ""
+    for item in pred_list:
+        if item != "":
+            pred_value = item
+            break
+    return pred_value
+
+
+def split_retrieval_answer(pred: str):
+    for c in ["\n", ":", '"', "'", ".", ",", "?", "!", "{", "}"]:
+        pred = pred.replace(c, " ")
+    words = pred.split()
+    return words
+
+
+def get_score_one_kv_retrieval(pred, label, model_name: str) -> bool:
+    for c in ['\n', ':', '\"', '\'', '.', ',', '?', '!', '{', '}']:
+        pred = pred.replace(c, ' ')
+    words = pred.split()
+    return label in words
+
+
+def get_score_one_passkey(pred, label, model_name: str) -> bool:
+    if isinstance(label, list):
+        label = label[0]
+    return label == first_int_match(pred)
+
+
+def get_score_one_number_string(pred, label, model_name: str) -> bool:
+    if isinstance(label, list):
+        label = label[0]
+    return label == first_int_match(pred)
+
+
+def get_score_one_code_run(pred, label, model_name: str) -> bool:
+    """
+    Returns the score of one example in Code.Run.
+    """
+    if isinstance(label, list):
+        label = label[0]
+    pred = pred.strip()
+    for c in ["\n", ".", "`", "'", '"', ":"]:
+        pred = pred.replace(c, " ")
+    words = pred.split()
+    if len(words) == 0:
+        return False
+    try:
+        pred = int(words[-1])
+        return label == pred
+    except Exception:
+        return False
+
+
+def get_score_one_code_debug(pred, label, model_name: str) -> bool:
+    """
+    Returns the score of one example in Code.Debug.
+    """
+    label_c = label[1]
+    fn_name = label[0]
+    if pred[:2] in [f"{label_c}.", f"{label_c}:"]:
+        return True
+
+    ans_prefixes = [
+        "answer is:",
+        # "answer is",
+        # "error is",
+        "is:",
+        "answer:",
+    ]
+    pred = pred.strip()
+    for c in ["\n", "`", "'", '"', "-", "*", "Option", "option"]:
+        pred = pred.replace(c, " ")
+    while "  " in pred:
+        pred = pred.replace("  ", " ")
+    for prefix in ans_prefixes:
+        idx = pred.find(prefix)
+        if idx == -1:
+            continue
+        # The prediction ends with this prefix
+        if len(pred) < idx + len(prefix) + 1:
+            return False
+        pred = pred[idx + len(prefix) + 1 :]
+        for s in [label_c, fn_name]:
+            if pred.startswith(s):
+                return True
+        return False
+    return False
+
+
+def get_score_one_math_find(pred, label, model_name: str) -> bool:
+    if isinstance(label, list):
+        # In math_find, there is always only one label.
+        label = label[0]
+    if isinstance(label, int):
+        # Find first int or float
+        first_num = re.search(r"\d+\.\d+|\d+", pred)
+        if first_num is None:
+            return False
+        first_num = first_num.group(0).strip()
+        return int(first_num) == label
+    elif isinstance(label, float):
+        # Find first float or int
+        first_float = re.search(r"\d+\.\d+|\d+", pred)
+        if first_float is None:
+            return False
+        first_float = first_float.group(0).strip()
+        return float(first_float) == label
+    else:
+        raise TypeError(f"Expected int or float, got {type(label)}")
+
+
+def get_score_one_longdialogue_qa_eng(pred, label, model_name: str) -> bool:
+    label = label[0]
+    for c in ["\n", ":", '"', "'", ".", ",", "?", "!", "{", "}"]:
+        pred = pred.replace(c, " ")
+    words = pred.split()
+    words = [x.upper() for x in words]
+    return label in words
+
+
+def get_score_one_longbook_choice_eng(pred, label, model_name: str) -> bool:
+    # Just use the first letter as the prediction
+    pred = pred.strip()
+    if pred == "":
+        return False
+    if pred[0] in "ABCD":
+        return pred[0] in label
+    if pred in label:
+        return True
+    # Find a answer prefix
+    for c in ["\n", '"', "'", ".", ",", "?", "!", "{", "}"]:
+        pred = pred.replace(c, " ")
+    while "  " in pred:
+        pred = pred.replace("  ", " ")
+    ans_prefixes = [
+        "answer is:",
+        "answer:",
+        "answer is",
+        "option is",
+    ]
+    for prefix in ans_prefixes:
+        idx = pred.find(prefix)
+        if idx == -1:
+            continue
+        # The prediction ends with this prefix
+        if len(pred) < idx + len(prefix) + 1:
+            return False
+        after_prefix = pred[idx + len(prefix) + 1 :]
+        for s in label:
+            if after_prefix.startswith(s):
+                return True
+        return False
+
+    # Finally, just find the first occurrence of A, B, C, or D.
+    words = pred.split()
+    for word in words:
+        if word in "ABCD":
+            return word in label
+    return False
+
+
+def get_score_one_longbook_qa_eng(pred, label, model_name: str) -> float:
+    return qa_f1_score(pred, label)
+
+
+def get_score_one_longbook_sum_eng(
+    pred: str, label: str, model_name: str
+) -> float:
+    rouge = Rouge()
+    if pred == "":
+        pred = "THIS_IS_A_NULL_STRING"
+    try:
+        scores = rouge.get_scores([pred], label, avg=True)
+        return scores["rouge-l"]["f"]
+    except:
+        return 0
+
+
+def get_score_one_longbook_qa_chn(pred, label, model_name: str) -> float:
+    return qa_f1_score_zh(pred, label)
+
+
+def get_score_one_math_calc(pred, label, model_name: str) -> float:
+    assert isinstance(label, list), f"Expected list, got {type(label)}"
+    # assert isinstance(pred, list), f"Expected list, got {type(pred)}"
+    pred_nums = []
+    pred_list = re.split("[^0-9]", pred)
+    for item in pred_list:
+        if item != "":
+            pred_nums.append(int(item))
+
+    # Our prompts makes GPT4 always output the first number as the first value
+    # in the predicted answer.
+    if model_name == "gpt4":
+        pred_nums = pred_nums[1:]
+
+    cnt = 0
+    for i in range(len(label)):
+        if i >= len(pred_nums):
+            break
+        if label[i] == pred_nums[i]:
+            cnt += 1
+        else:
+            break
+    return cnt / len(label)
+
+
+def get_score_one(
+    pred: str, label: str, task_name: str, model_name: str
+) -> float:
+    """
+    Computes the score for one prediction.
+    Returns one float (zero and one for boolean values).
+    """
+    NAME_TO_SCORE_GETTER = {
+        # Retrieve
+        "kv_retrieval": get_score_one_kv_retrieval,
+        "kv_retrieval_prefix": get_score_one_kv_retrieval,
+        "kv_retrieval_both": get_score_one_kv_retrieval,
+
+        "passkey": get_score_one_passkey,
+        "number_string": get_score_one_number_string,
+        # Code
+        "code_run": get_score_one_code_run,
+        "code_debug": get_score_one_code_debug,
+        # Longbook
+        "longdialogue_qa_eng": get_score_one_longdialogue_qa_eng,
+        "longbook_qa_eng": get_score_one_longbook_qa_eng,
+        "longbook_sum_eng": get_score_one_longbook_sum_eng,
+        "longbook_choice_eng": get_score_one_longbook_choice_eng,
+        "longbook_qa_chn": get_score_one_longbook_qa_chn,
+        # Math
+        "math_find": get_score_one_math_find,
+        "math_calc": get_score_one_math_calc,
+    }
+    assert task_name in NAME_TO_SCORE_GETTER, f"Invalid task name: {task_name}"
+    score = NAME_TO_SCORE_GETTER[task_name](pred, label, model_name)
+    return float(score)
+
+
+def get_labels(preds: list) -> list[str]:
+    possible_label_keys = ["ground_truth", "label"]
+    for label_key in possible_label_keys:
+        if label_key in preds[0]:
+            return [x.get(label_key, "XXXXXXXXXX") for x in preds]
+    raise ValueError(f"Cannot find label in {preds[0]}")
+
+
+def get_preds(preds: list, data_name: str) -> list[str]:
+    pred_strings = []
+    possible_pred_keys = ["prediction", "pred"]
+    for pred in preds:
+        this_pred = "NO PREDICTION"
+        for pred_key in possible_pred_keys:
+            if pred_key in pred:
+                this_pred = pred[pred_key]
+                break
+        else:
+            raise ValueError(f"Cannot find prediction in {pred}")
+        pred_strings.append(this_pred)
+    return pred_strings
+
+
+def get_score(
+    labels: list, preds: list, data_name: str, model_name: str
+) -> float:
+    """
+    Computes the average score for a task.
+    """
+    assert len(labels) == len(preds)
+    scores = []
+    for label, pred in tqdm(zip(labels, preds)):
+        score = get_score_one(pred, label, data_name, model_name)
+        scores.append(score)
+    return sum(scores) / len(scores)
+
+
+def compute_scores(preds_path, data_name: str, model_name: str):
+    print("Loading prediction results from", preds_path)
+    preds = list(iter_jsonl(preds_path))
+    labels = get_labels(preds)
+    preds = get_preds(preds, data_name)
+
+    acc = get_score(labels, preds, data_name, model_name)
+    print(acc)
+
+
+def create_prompt(eg: dict, data_name: str, prompt_template: str) -> str:
+    """
+    Create prompt for a given example.
+
+    Args:
+        eg: example dict
+        data_name: name of the dataset/task
+    """
+    # if model_name == "gpt4":
+    #     # Math.Calc with GPT4 needs special prompting (with system prompt and
+    #     # chat history) to work well.
+    #     if data_name == "math_calc":
+    #         return eg["context"]
+
+    templates = MODEL_TO_PROMPT_TEMPLATE[prompt_template]
+    template = templates[data_name]
+    # ================= Code tasks
+    if data_name == "code_run":
+        find_result = re.findall(r"func_[0-9]+\(\-?[0-9]+\)", eg['input'])
+        func_call = find_result[0]
+        func = func_call.split("(")[0]
+        return template.format(
+            func=func,
+            func_call=func_call,
+            context=eg["context"],
+        )
+    elif data_name in ["code_debug", "code_debug_qa"]:
+        # Load source code
+        code = eg["context"]
+        if data_name == "code_debug":
+            return template.format(
+                context=code,
+                OPTION_A=eg["options"][0],
+                OPTION_B=eg["options"][1],
+                OPTION_C=eg["options"][2],
+                OPTION_D=eg["options"][3],
+            )
+        return template.format(
+            context=code,
+        )
+    # ================= Code tasks
+    elif data_name == "longdialogue_qa_eng":
+        script = eg["context"]
+        prompt = template.format(context=script)
+        return prompt
+    # ==================== Long book tasks
+    elif data_name in [
+        "longbook_choice_eng",
+        "longbook_qa_eng",
+        "longbook_sum_eng",
+        "longbook_qa_chn",
+    ]:
+        book = eg["context"]
+        if data_name == "longbook_choice_eng":
+            return template.format(
+                question=eg["input"],
+                context=book,
+                OPTION_A=eg["options"][0],
+                OPTION_B=eg["options"][1],
+                OPTION_C=eg["options"][2],
+                OPTION_D=eg["options"][3],
+            )
+        elif data_name == "longbook_qa_eng":
+            return template.format(
+                question=eg["input"],
+                context=book,
+            )
+        elif data_name == "longbook_sum_eng":
+            return template.format(
+                context=book,
+            )
+        elif data_name == "longbook_qa_chn":
+            return template.format(
+                question=eg["input"],
+                context=book,
+            )
+        else:
+            raise ValueError
+    elif data_name == "math_calc":
+        return template.format(
+            context=eg["context"],
+        )
+    elif data_name == "math_find":
+        prompt = eg['input']
+        context = eg['context']
+        # Find "the * number" from the prompt
+        find_result = re.findall(r"The .+ of", prompt)
+        assert find_result, f"Cannot find the target number in {prompt}"
+        target_number = find_result[0].lower()[:-3]
+        # Replace the number with the answer
+        prefix = f"What is {target_number} in the following list?"
+        return template.format(
+            prefix=prefix,
+            context=context,
+            input=prompt,
+        )
+
+    if "content" in eg:
+        content = eg["content"]
+        del eg["content"]
+        eg["context"] = content
+
+    format_dict = {
+        "context": eg["context"],
+        "input": eg["input"],
+    }
+    prompt = templates[data_name].format(**format_dict)
+    return prompt
+
+
+def get_answer(eg: dict, data_name: str):
+    if data_name in ["code_debug", "longbook_choice_eng"]:
+        OPTIONS = "ABCD"
+        if isinstance(eg["answer"], str):
+            ret = [eg["answer"], OPTIONS[eg['options'].index(eg["answer"])]]
+        elif isinstance(eg["answer"], list):
+            if len(eg["answer"]) == 1:
+                ret = [eg["answer"][0], OPTIONS[eg['options'].index(eg["answer"][0])]]
+            elif len(eg["answer"]) == 2 and eg["answer"][1] in ['A', 'B', 'C', 'D']:
+                ret = eg['answer']
+            else:
+                raise ValueError
+        else:
+            raise ValueError
+        return ret
+
+    return eg["answer"]
+
+
+ALL_TASKS = [
+    "passkey",
+    "number_string",
+    "kv_retrieval",
+    "longdialogue_qa_eng",
+    "longbook_sum_eng",
+    "longbook_choice_eng",
+    "longbook_qa_eng",
+    "longbook_qa_chn",
+    "math_find",
+    "math_calc",
+    "code_run",
+    "code_debug",
+]
+
+
+TASK_TO_PATH = {
+    # Retrieval tasks
+    "passkey": "passkey.jsonl",
+    "number_string": "number_string.jsonl",
+    "kv_retrieval": "kv_retrieval.jsonl",
+    # Book tasks
+    "longbook_sum_eng": "longbook_sum_eng.jsonl",
+    "longbook_choice_eng": "longbook_choice_eng.jsonl",
+    "longbook_qa_eng": "longbook_qa_eng.jsonl",
+    "longbook_qa_chn": "longbook_qa_chn.jsonl",
+    # "book_qa_eng": "longbook_eng/longbook_qa_eng.jsonl",
+    "longdialogue_qa_eng": "longdialogue_qa_eng.jsonl",
+    # Math tasks
+    "math_find": "math_find.jsonl",
+    "math_calc": "math_calc.jsonl",
+    # Code tasks
+    "code_run": "code_run.jsonl",
+    "code_debug": "code_debug.jsonl",
+}
+
+TASK_TO_MAX_NEW_TOKENS = {
+    "passkey": 6,
+    "number_string": 12,
+    "kv_retrieval": 50,
+    "longbook_sum_eng": 1200,
+    "longbook_choice_eng": 40,
+    "longbook_qa_eng": 40,
+    "longbook_qa_chn": 40,
+    "longdialogue_qa_eng": 40,
+    "math_find": 3,
+    "math_calc": 30000,
+    "code_run": 5,
+    "code_debug": 5,
+}
+
+gpt4_templates = {
+    "passkey": "There is an important info hidden inside a lot of irrelevant text. Find it and memorize them. I will quiz you about the important information there.\n\n{context}\n\n{input}",  # noqa
+    "number_string": "There is an important info hidden inside a lot of irrelevant text. Find it. I will quiz you about the important information there.\n\n{context}\n\n{input}",  # noqa
+    "kv_retrieval": "Extract the value corresponding to the specified key in the JSON object below.\n\n{context}\n\n{input}",  # noqa
+    # "longbook_sum_eng": "Summarize the book below:\n\n{context}",  # noqa
+    "longbook_qa_eng": "Read the book below and answer a question.\n\n{context}\n\nQuestion: {question}\n\nBe very concise.",  # noqa
+    "longbook_choice_eng": "Read the book and answer the question.\n\n{context}\n\nQuestion: {question}\n\nOnly one of the following options is correct, tell me the answer using one single letter (A, B, C, or D). Don't say anything else.\nA. {OPTION_A}\nB. {OPTION_B}\nC. {OPTION_C}\nD. {OPTION_D}",  # noqa
+    "longbook_sum_eng": "Summarize the following book.\n\n{context}",  # noqa
+    "longbook_qa_chn": "请根据以下书籍回答我的问题。\n\n{context}\n\n问题：{question}\n请尽量简短地回答。",  # noqa
+    "math_find": "{prefix}\n\n{context}\n\n{input}",
+    "math_calc": "Compute the intermediate values in the following long expression.\n\n{context}",  # noqa
+    "code_run": "Following is a set of Python functions. There is a function called named {func}.\n\n{context}\n\nPlease give me the exact number of the return value of {func_call}. Be concise. Your response must end with the final returned value.",  # noqa
+    "code_debug": "There is ONLY ONE function in the large project that is deliberately made to include an obvious error. Please find the function that contains the most obvious errors. I will give you four options to narrow your scope. You can inspect the options and think. Eventually, tell me the answer using one single letter (A, B, C, or D).\n\n{context}\n\nWhich funtion has deliberate error?\nA. {OPTION_A}\nB. {OPTION_B}\nC. {OPTION_C}\nD. {OPTION_D}\n\nYou should first find the functions in the options. Repeat their content, inspect through code, and at last give me your answer for the function that has the deliberate and obvious error in A, B, C, or D.",  # noqa
+    "longdialogue_qa_eng": "Below is a dialogue script where one random occurrence of a character name is replaced with \"$$MASK$$\", and you should try to guess who that character is.\n\nThe dialogue:\n\n---\n\n{context}\n\n---\n\nEnd of dialogue.\n\nWhich character is most likely \"$$MASK$$\"? Just say the name used by the scriptwriter (before the colon marks) of one single character and nothing else.",  # noqa
+}
+
+yarn_mistral_templates = {
+    "passkey": "There is an important info hidden inside a lot of irrelevant text. Find it and memorize it. I will quiz you about the important information.\n\n{context}\n\n{input}\n\nThe pass key is",  # noqa
+    "number_string": "There is an important info hidden inside a lot of irrelevant text. Find it. I will quiz you about the important information there.\n\n{context}\n\n{input}\n\nThe sequence of digits is",  # noqa
+    "kv_retrieval": "Extract the value corresponding to the specified key in the JSON object below.\n\n{context}\n\n{input}",  # noqa
+    "longbook_sum_eng": "Summarize the book below.\n\n{context}\n\nSummary:",  # noqa
+    "longbook_choice_eng": "Read the book and answer the question.\n\n{context}\n\nQuestion: {question}\nA. {OPTION_A}\nB. {OPTION_B}\nC. {OPTION_C}\nD. {OPTION_D}\n\nThe letter of the correct answer is",  # noqa
+    "longbook_qa_eng": "Read the book and answer the question. Be very concise in your answer.\n\n{context}\n\nQuestion: {question}\nAnswer:",  # noqa
+    "longbook_qa_chn": "阅读以下书籍然后回答问题。\n\n{context}\n\n问题：{question}\n答案：",  # noqa
+    "math_find": "{prefix}\n\n{context}\n\n{input}",
+    "math_calc": "Let us calculate the intermediate values of an expression.\n\nExpression: 1 + 3 + 4\nValues: [1, 4, 8]\n\nExpression: 8 - 3 + 2 - 4\nValues: [8, 5, 7, 3]\n\nExpression: {context}\nValues:",  # noqa
+    "code_run": "There is a function called {func} in the following Python code.\n\n{context}\n\nPlease compute the exact value of {func_call}. The value of {func_call} is",  # noqa
+    "code_debug": "Following is a Python code where exactly one of the functions/methods has a deliberate error that makes it crash.\n\n{context}\n\nOptions:\nA. {OPTION_A}\nB. {OPTION_B}\nC. {OPTION_C}\nD. {OPTION_D}\n\nThe correct option is:",  # noqa
+    "longdialogue_qa_eng": "Below is a dialogue script where one random occurrence of a character name is replaced with \"$$MASK$$\", and you should try to guess who that character is.\n\n{context}\n\nThe name that has been replaced with $$MASK$$ is likely",  # noqa
+}
+
+claude2_templates = {
+    "passkey": "There is an important info hidden inside a lot of irrelevant text. Find it and memorize them. I will quiz you about the important information there.\n\n{context}\n{input}\nThe pass key is",
+    "number_string": "There is an important info hidden inside a lot of irrelevant text. Find it. I will quiz you about the important information there.\n\n{context}\n{input}\nThe sequence of digits is",  # noqa
+    "kv_retrieval": "There is an important info hidden inside a lot of irrelevant text. Find it. I will quiz you about the important information there.\n\n{context}\n{input}",
+    "longbook_sum_eng": "Summarize the following book.\n\n{context}",  # noqa
+    "longbook_choice_eng": "Read the book and answer the question.\n\n{context}\n\nQuestion: {question}\n\nOnly one of the following options is correct, tell me the answer using one single letter (A, B, C, or D). Don't say anything else.\nA. {OPTION_A}\nB. {OPTION_B}\nC. {OPTION_C}\nD. {OPTION_D}",  # noqa
+    "longbook_qa_eng": "Read the novel below and answer a question:\n\n{context}\n\n{input}\nPlease answer as short as possible. The answer is: ",  # noqa
+    "longbook_qa_chn": "请根据以下书籍回答我的问题。\n\n{context}\n\n问题：{question}\n请尽量简短地回答。",  # noqa
+    "math_find": "{prefix}\n\n{context}\n\n{input}",
+    "math_calc": "Let us calculate the intermediate values of an expression.\nExpression: 1 + 3 + 4\nValues: [1, 4, 8]\n\nExpression: 8 - 3 + 2 - 4\nValues: [8, 5, 7, 3]\n\nExpression: {context}\nValues:",  # noqa
+    "code_run": "In the file functions_module.py, there is a function called ${func}.\n\n\nHere is the content of functions_module.py:\n{context}\n\nPlease give me the exact number of the return value of {func_call}. Your response should end with the sentence \'The return value is:\'.",  # noqa
+    "code_debug": "There is ONLY ONE function in the large project that is deliberately made to include an obvious error. Please find the function that contains the most obvious errors. I will give you four options to narrow your scope. You can inspect through the options and think. Eventually, tell me the answer using one single letter (A, B, C, or D).\n\n{context}\n\nWhich funtion has deliberate error?\nA. {OPTION_A}\nB. {OPTION_B}\nC. {OPTION_C}\nD. {OPTION_D}\n\nYou should first find the functions in the options. Repeat their content, inspect through code, and at last give me your answer for the function that has the deliberate and obvious error in A, B, C, or D.",  # noqa
+    "longdialogue_qa_eng": "Below is a dialogue script where one random occurrence of a character name is replaced with \"$$MASK$$\", and you should try to guess who that character is.\n\nThe dialogue:\n\n---\n\n{context}\n\n---\n\nEnd of dialogue.\n\nWhich character is most likely \"$$MASK$$\"? Just say the name used by the scriptwriter (before the colon marks) of one single character and nothing else.",  # noqa
+}
+
+kimi_templates = {
+    "passkey": "There is an important info hidden inside a lot of irrelevant text. Find it and memorize them. I will quiz you about the important information there.\n\n{context}\n{input}\nThe pass key is",  # noqa
+    "number_string": "There is an important info hidden inside a lot of irrelevant text. Find it. I will quiz you about the important information there.\n\n{context}\n{input}\nThe sequence of digits is",  # noqa
+    "kv_retrieval": "Extract the value corresponding to the specified key in the JSON object below.\n\n{context}\n{input}",  # noqa
+    "longbook_sum_eng": "Summarize the book below:\n\n{file:{context}}",  # noqa
+    "longbook_choice_eng": "Read the book and answer the question.\n\nQuestion: {question}\n\nOnly one of the following options is correct, tell me the answer using one single letter (A, B, C, or D). Don't say anything else.\nA. {OPTION_A}\nB. {OPTION_B}\nC. {OPTION_C}\nD. {OPTION_D}" + "{file:{document}}",  # noqa
+    "longbook_qa_eng": "Read the book below and answer a question.\n\nQuestion: {question}\n\nBe very concise." + "{file:{context}}",  # noqa
+    "longbook_qa_chn": "阅读以下书籍然后回答问题。\n\n问题：{question}\n答案：" + "{file:{context}}",  # noqa
+    "math_find": "{prefix}\n\n{context}\n\n{input}",
+    "math_calc": "Let us calculate the intermediate values of an expression.\nExpression: 1 + 3 + 4\nValues: [1, 4, 8]\n\nExpression: 8 - 3 + 2 - 4\nValues: [8, 5, 7, 3]\n\nExpression: {context}\nValues:",  # noqa
+    "code_run": "In the file functions_module.py, there is a function called ${func}.\n\n\nHere is the content of functions_module.py:\n\nPlease give me the exact number of the return value of ${func_call}. Your response should end with the sentence 'The return value is:'." + "{context}",  # noqa
+    "code_debug": "Below is a code repository where there is one single function with bugs that causes an error. Please tell me the name of that function.\nWhich function has bugs? Give me the final answer in this format: \"[FINAL ANSWER: XXX]\". Don't say anything else." + "{fcontext}",  # noqa
+    # "longdialogue_qa_eng": "Below is a dialogue script where one random occurrence of a character name is replaced with \"$$MASK$$\", and you should try to guess who that character is.\n\nThe name that has been replaced with $$MASK$$ is likely" + "{context}",  # noqa
+    "longdialogue_qa_eng": "Below is a dialogue script where one random occurrence of a character name is replaced with \"$$MASK$$\", and you should try to guess who that character is. Give me the answer using the name before the colons, don't say anything else.\n\n{context}",  # noqa
+}
+
+MODEL_TO_PROMPT_TEMPLATE = {
+    "gpt4": gpt4_templates,
+    "claude2": claude2_templates,
+    "kimi": kimi_templates,
+    "mistral": yarn_mistral_templates,
+}

Long_LLM/activation_beacon/new/main/longbench_utils.py

+import re
+import string
+import jieba
+import difflib
+import numpy as np
+from fuzzywuzzy import fuzz
+from typing import List
+from collections import Counter
+from rouge import Rouge
+
+
+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 normalize_zh_answer(s):
+    """Lower text and remove punctuation, extra whitespace."""
+
+    def white_space_fix(text):
+        return "".join(text.split())
+
+    def remove_punc(text):
+        cn_punctuation = "！？｡。＂＃＄％＆＇（）＊＋，－／：；＜＝＞＠［＼］＾＿｀｛｜｝～｟｠｢｣､、〃》「」『』【】〔〕〖〗〘〙〚〛〜〝〞〟〰〾〿–—‘’‛“”„‟…‧﹏."
+        all_punctuation = set(string.punctuation + cn_punctuation)
+        return "".join(ch for ch in text if ch not in all_punctuation)
+
+    def lower(text):
+        return text.lower()
+
+    return white_space_fix(remove_punc(lower(s)))
+
+def count_score(prediction, ground_truth, **kwargs):
+    numbers = re.findall(r"\d+", prediction)
+    right_num = 0
+    for number in numbers:
+        if str(number) == str(ground_truth):
+            right_num += 1
+    final_score = 0.0 if len(numbers) == 0 else right_num / len(numbers)
+    return float(final_score)
+
+def retrieval_score(prediction, ground_truth, **kwargs):
+    pattern = r'Paragraph (\d+)'
+    matches = re.findall(pattern, ground_truth)
+    ground_truth_id = matches[0]
+    numbers = re.findall(r"\d+", prediction)
+    right_num = 0
+    for number in numbers:
+        if str(number) == str(ground_truth_id):
+            right_num += 1
+    final_score = 0.0 if len(numbers) == 0 else right_num / len(numbers)
+    return float(final_score)
+
+def retrieval_zh_score(prediction, ground_truth, **kwargs):
+    pattern = r'段落(\d+)'
+    matches = re.findall(pattern, ground_truth)
+    ground_truth_id = matches[0]
+    numbers = re.findall(r"\d+", prediction)
+    right_num = 0
+    for number in numbers:
+        if str(number) == str(ground_truth_id):
+            right_num += 1
+    final_score = 0.0 if len(numbers) == 0 else right_num / len(numbers)
+    return float(final_score)
+
+def code_sim_score(prediction, ground_truth, **kwargs):
+    all_lines = prediction.lstrip('\n').split('\n')
+    prediction = ""
+    for line in all_lines:
+        if ('`' not in line) and ('#' not in line) and ('//' not in line):
+            prediction = line
+            break
+    return (fuzz.ratio(prediction, ground_truth) / 100)
+
+def classification_score(prediction, ground_truth, **kwargs):
+    em_match_list = []
+    all_classes = kwargs["all_classes"]
+    for class_name in all_classes:
+        if class_name in prediction:
+            em_match_list.append(class_name)
+    for match_term in em_match_list:
+        if match_term in ground_truth and match_term != ground_truth:
+            em_match_list.remove(match_term)
+    if em_match_list != 0:
+        if ground_truth in em_match_list:
+            score = (1.0 / len(em_match_list))
+        else:
+            score = 0.0
+    else:
+        best_match = None
+        highest_similarity = 0
+        for string in all_classes:
+            similarity = difflib.SequenceMatcher(None, string, prediction).ratio()
+            if similarity > highest_similarity:
+                highest_similarity = similarity
+                best_match = string
+        score = float(best_match == ground_truth)
+    return score
+    
+def rouge_score(prediction, ground_truth, **kwargs):
+    rouge = Rouge()
+    try:
+        scores = rouge.get_scores([prediction], [ground_truth], avg=True)
+    except:
+        return 0.0
+    return scores["rouge-l"]["f"]
+
+def rouge_score_zh(prediction, ground_truth, **kwargs):
+    prediction = " ".join(list(jieba.cut(prediction, cut_all=False)))
+    ground_truth = " ".join(list(jieba.cut(ground_truth, cut_all=False))) 
+    score = rouge_score(prediction, ground_truth)
+    return score
+
+def f1_score(prediction, ground_truth, **kwargs):
+    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 qa_f1_score(prediction, ground_truth, **kwargs):
+    normalized_prediction = normalize_answer(prediction)
+    normalized_ground_truth = normalize_answer(ground_truth)
+
+    prediction_tokens = normalized_prediction.split()
+    ground_truth_tokens = normalized_ground_truth.split()
+    return f1_score(prediction_tokens, ground_truth_tokens)
+
+
+def qa_f1_score_zh(prediction, ground_truth, **kwargs):
+    prediction_tokens = list(jieba.cut(prediction, cut_all=False))
+    ground_truth_tokens = list(jieba.cut(ground_truth, cut_all=False))
+    prediction_tokens = [normalize_zh_answer(token) for token in prediction_tokens]
+    ground_truth_tokens = [normalize_zh_answer(token) for token in ground_truth_tokens]
+    prediction_tokens = [token for token in prediction_tokens if len(token) > 0]
+    ground_truth_tokens = [token for token in ground_truth_tokens if len(token) > 0]
+    return f1_score(prediction_tokens, ground_truth_tokens)
+
+def scorer(dataset, predictions, answers, all_classes):
+    total_score = 0.
+    for (prediction, ground_truths) in zip(predictions, answers):
+        score = 0.
+        if dataset in ["trec", "triviaqa", "samsum", "lsht"]:
+            prediction = prediction.lstrip('\n').split('\n')[0]
+        for ground_truth in ground_truths:
+            score = max(score, DATASET2METRIC[dataset](prediction, ground_truth, all_classes=all_classes))
+        total_score += score
+    return round(100 * total_score / len(predictions), 2)
+
+
+DATASET2PROMPT = {
+    "narrativeqa": "You are given a story, which can be either a novel or a movie script, and a question. Answer the question asconcisely as you can, using a single phrase if possible. Do not provide any explanation.\n\nStory: {context}\n\nNow, answer the question based on the story asconcisely as you can, using a single phrase if possible. Do not provide any explanation.\n\nQuestion: {input}\n\nAnswer:",
+    "qasper": "You are given a scientific article and a question. Answer the question as concisely as you can, using a single phrase or sentence if possible. If the question cannot be answered based on the information in the article, write \"unanswerable\". If the question is a yes/no question, answer \"yes\", \"no\", or \"unanswerable\". Do not provide any explanation.\n\nArticle: {context}\n\n Answer the question based on the above article as concisely as you can, using a single phrase or sentence if possible. If the question cannot be answered based on the information in the article, write \"unanswerable\". If the question is a yes/no question, answer \"yes\", \"no\", or \"unanswerable\". Do not provide any explanation.\n\nQuestion: {input}\n\nAnswer:",
+    "multifieldqa_en": "Read the following text and answer briefly.\n\n{context}\n\nNow, answer the following question based on the above text, only give me the answer and do not output any other words.\n\nQuestion: {input}\nAnswer:",
+    "multifieldqa_zh": "阅读以下文字并用中文简短回答：\n\n{context}\n\n现在请基于上面的文章回答下面的问题，只告诉我答案，不要输出任何其他字词。\n\n问题：{input}\n回答：",
+    "hotpotqa": "Answer the question based on the given passages. Only give me the answer and do not output any other words.\n\nThe following are given passages.\n{context}\n\nAnswer the question based on the given passages. Only give me the answer and do not output any other words.\n\nQuestion: {input}\nAnswer:",
+    "2wikimqa": "Answer the question based on the given passages. Only give me the answer and do not output any other words.\n\nThe following are given passages.\n{context}\n\nAnswer the question based on the given passages. Only give me the answer and do not output any other words.\n\nQuestion: {input}\nAnswer:",
+    "musique": "Answer the question based on the given passages. Only give me the answer and do not output any other words.\n\nThe following are given passages.\n{context}\n\nAnswer the question based on the given passages. Only give me the answer and do not output any other words.\n\nQuestion: {input}\nAnswer:",
+    "dureader": "请基于给定的文章回答下述问题。\n\n文章：{context}\n\n请基于上述文章回答下面的问题。\n\n问题：{input}\n回答：",
+    "gov_report": "You are given a report by a government agency. Write a one-page summary of the report.\n\nReport:\n{context}\n\nNow, write a one-page summary of the report.\n\nSummary:",
+    "qmsum": "You are given a meeting transcript and a query containing a question or instruction. Answer the query in one or more sentences.\n\nTranscript:\n{context}\n\nNow, answer the query based on the above meeting transcript in one or more sentences.\n\nQuery: {input}\nAnswer:",
+    "multi_news": "You are given several news passages. Write a one-page summary of all news. \n\nNews:\n{context}\n\nNow, write a one-page summary of all the news.\n\nSummary:",
+    "vcsum": "下面有一段会议记录，请你阅读后，写一段总结，总结会议的内容。\n会议记录：\n{context}\n\n会议总结：",
+    "trec": "Please determine the type of the question below. Here are some examples of questions.\n\n{context}\n{input}",
+    "triviaqa": "Answer the question based on the given passage. Only give me the answer and do not output any other words. The following are some examples.\n\n{context}\n\n{input}",
+    "samsum": "Summarize the dialogue into a few short sentences. The following are some examples.\n\n{context}\n\n{input}",
+    "lsht": "请判断给定新闻的类别，下面是一些例子。\n\n{context}\n{input}",
+    "passage_count": "There are some paragraphs below sourced from Wikipedia. Some of them may be duplicates. Please carefully read these paragraphs and determine how many unique paragraphs there are after removing duplicates. In other words, how many non-repeating paragraphs are there in total?\n\n{context}\n\nPlease enter the final count of unique paragraphs after removing duplicates. The output format should only contain the number, such as 1, 2, 3, and so on.\n\nThe final answer is: ",
+    "passage_retrieval_en": "Here are 30 paragraphs from Wikipedia, along with an abstract. Please determine which paragraph the abstract is from.\n\n{context}\n\nThe following is an abstract.\n\n{input}\n\nPlease enter the number of the paragraph that the abstract is from. The answer format must be like \"Paragraph 1\", \"Paragraph 2\", etc.\n\nThe answer is: ",
+    "passage_retrieval_zh": "以下是若干段落文字，以及其中一个段落的摘要。请确定给定的摘要出自哪一段。\n\n{context}\n\n下面是一个摘要\n\n{input}\n\n请输入摘要所属段落的编号。答案格式必须是\"段落1\"，\"段落2\"等格式\n\n答案是：",
+    "lcc": "Please complete the code given below. \n{context}Next line of code:\n",
+    "repobench-p": "Please complete the code given below. \n{context}{input}Next line of code:\n"
+}
+
+DATASET2MAXNEWTOKENS = {
+    "narrativeqa": 128,
+    "qasper": 128,
+    "multifieldqa_en": 64,
+    "multifieldqa_zh": 64,
+    "hotpotqa": 32,
+    "2wikimqa": 32,
+    "musique": 32,
+    "dureader": 128,
+    "gov_report": 512,
+    "qmsum": 512,
+    "multi_news": 512,
+    "vcsum": 512,
+    "trec": 64,
+    "triviaqa": 32,
+    "samsum": 128,
+    "lsht": 64,
+    "passage_count": 32,
+    "passage_retrieval_en": 32,
+    "passage_retrieval_zh": 32,
+    "lcc": 64,
+    "repobench-p": 64
+}
+
+DATASET2METRIC = {
+    "narrativeqa": qa_f1_score,
+    "qasper": qa_f1_score,
+    "multifieldqa_en": qa_f1_score,
+    "multifieldqa_zh": qa_f1_score_zh,
+    "hotpotqa": qa_f1_score,
+    "2wikimqa": qa_f1_score,
+    "musique": qa_f1_score,
+    "dureader": rouge_score_zh,
+    "gov_report": rouge_score,
+    "qmsum": rouge_score,
+    "multi_news": rouge_score,
+    "vcsum": rouge_score_zh,
+    "trec": classification_score,
+    "triviaqa": qa_f1_score,
+    "samsum": rouge_score,
+    "lsht": classification_score,
+    "passage_retrieval_en": retrieval_score,
+    "passage_count": count_score,
+    "passage_retrieval_zh": retrieval_zh_score,
+    "lcc": code_sim_score,
+    "repobench-p": code_sim_score,
+}
+
+DATASET2CATEGORY = {
+    "narrativeqa": "EN Single-Doc QA",
+    "qasper": "EN Single-Doc QA",
+    "multifieldqa_en": "EN Single-Doc QA",
+    "multifieldqa_zh": "CN Single-Doc QA",
+    "hotpotqa": "EN Multi-Doc QA",
+    "2wikimqa": "EN Multi-Doc QA",
+    "musique": "EN Multi-Doc QA",
+    "dureader": "CN Multi-Doc QA",
+    "gov_report": "EN Summarization",
+    "qmsum": "EN Summarization",
+    "multi_news": "EN Summarization",
+    "vcsum": "CN Summarization",
+    "trec": "EN Few-Shot Learning",
+    "triviaqa": "EN Few-Shot Learning",
+    "samsum": "EN Few-Shot Learning",
+    "lsht": "CN Few-Shot Learning",
+    "passage_retrieval_en": "EN Synthetic Task",
+    "passage_count": "EN Synthetic Task",
+    "passage_retrieval_zh": "CN Synthetic Task",
+    "lcc": "Code Completion",
+    "repobench-p": "Code Completion",
+}
\ No newline at end of file

Long_LLM/activation_beacon/new/main/pretrain_data.py

+import os
+import json
+import random
+import math
+import datasets
+from tqdm import tqdm
+from typing import List
+from datetime import timedelta
+from accelerate import Accelerator, InitProcessGroupKwargs
+from dataclasses import dataclass, field
+from transformers import HfArgumentParser
+from transformers.utils import logging
+from transformers.tokenization_utils import PreTrainedTokenizer
+
+from src import split_file_dir_name_ext, get_model_and_tokenizer, format_numel_str, ModelArgs
+
+
+logger = logging.get_logger(__name__)
+
+
+@dataclass
+class Args(ModelArgs):
+    config: str = field(
+        default="data/config/slimpajama.json",
+        metadata={'help': 'Configuration json path for standard pretraining (concatenating multiple documents to form instances of equal lengths).'}
+    )
+    train_data: str = field(
+        default="long-llm:slimpajama",
+        metadata={'help': 'Directory of training data (multiple json files whose name correspond to the ones in config).'}
+    )
+    output_dir: str = field(
+        default="data/pretrain/llama-8K_2B",
+        metadata={'help': 'Output directory for results and logs.'}
+    )
+
+    num_token: List[str] = field(
+        default_factory=lambda: ["8192:2B"],
+        metadata={'help': 'How many tokens to use for a specified length? (T/t for trillion, B/b for billion, M/m for million)'}
+    )
+    add_bos: bool = field(
+        default=True,
+        metadata={'help': 'Add bos at the end of each document?'}
+    )
+    add_eos: bool = field(
+        default=True,
+        metadata={'help': 'Add eos at the end of each document?'}
+    )
+    seed: int = field(
+        default=123,
+        metadata={'help': 'Random seed.'}
+    )
+
+
+def prepare_pretrain_data(data_files, tokenizer: PreTrainedTokenizer, config: dict, length_2_num_token: dict, add_bos:bool=True, add_eos:bool=True, seed=42, cache_dir=None, load_from_cache_file=None):
+    random.seed(seed)
+
+    if isinstance(data_files, list):
+        data_files = data_files[0]
+
+    assert os.path.isdir(data_files), f"Make sure the data_files parameter is a directory containing the pretraining data json files! Found {data_files}."
+        
+    def _process(data):
+        input_ids = tokenizer(data["text"], add_special_tokens=False)["input_ids"]
+        return {"input_ids": input_ids}
+
+    num_token_avg_per_source = config["num_tokens_avg"]
+    mixture = config["mixture"]
+
+    # concatenate all input_ids and partiton them according to num_instances
+    outputs = {"input_ids": [], "attention_mask": [], "labels": [], "length": []}
+
+    for file_name in os.listdir(data_files):
+        file_path = os.path.join(data_files, file_name)
+        dataset_name = split_file_dir_name_ext(file_path)[1]
+
+        if dataset_name not in mixture:
+            continue
+
+        mix_portion = mixture[dataset_name] / 100
+        
+        if mix_portion == 0:
+            continue
+
+        num_token_this_dataset = {k: math.ceil(v * mix_portion) for k, v in length_2_num_token.items()}
+        num_instances_this_dataset = {k: math.ceil(v / k) for k, v in num_token_this_dataset.items()}
+        info = {k: format_numel_str(v) for k, v in num_token_this_dataset.items()}
+        logger.info(f"processing {dataset_name} dataset, generating {info} tokens...")
+
+        # tokenize all records
+        dataset = datasets.load_dataset("json", data_files=file_path, split="train", cache_dir=cache_dir)
+        dataset = dataset.map(_process, batched=True, num_proc=32, remove_columns=dataset.column_names, batch_size=100, load_from_cache_file=load_from_cache_file)
+
+        tqdm_bar = tqdm(total=sum(num_instances_this_dataset.values()))
+
+        max_length_candidates = [k for k, v in num_instances_this_dataset.items() if v > 0]
+        max_length = random.choice(max_length_candidates)
+
+        input_ids = []
+        for x in dataset:
+            sample_input_ids = x["input_ids"]
+            if add_bos:
+                assert tokenizer.bos_token_id is not None, f"Make sure the bos_token_id exists when enable add_eos."
+                sample_input_ids = [tokenizer.bos_token_id] + sample_input_ids
+            if add_eos:
+                assert tokenizer.eos_token_id is not None, f"Make sure the eos_token_id exists when enable add_eos."
+                sample_input_ids = sample_input_ids + [tokenizer.eos_token_id]
+            # add input_ids
+            input_ids.extend(sample_input_ids)
+            
+            if len(input_ids) >= max_length:
+                cursor = 0
+                while cursor + max_length <= len(input_ids):
+                    instance_input_ids = input_ids[cursor: cursor + max_length].copy()
+                    instance_attention_mask = [1 for _ in instance_input_ids]
+                    instance_labels = instance_input_ids.copy()
+
+                    # move the cursor
+                    cursor += max_length
+
+                    # add to final data
+                    outputs["input_ids"].append(instance_input_ids)
+                    outputs["attention_mask"].append(instance_attention_mask)
+                    outputs["labels"].append(instance_labels)
+                    outputs["length"].append(max_length)
+
+                    # update num_instances
+                    num_instances_this_dataset[max_length] -= 1
+                    tqdm_bar.update(1)
+
+                    # sample new max_length
+                    max_length_candidates = [k for k, v in num_instances_this_dataset.items() if v > 0]
+                    if len(max_length_candidates) == 0:
+                        # all needed data have been collected
+                        break
+                    elif len(max_length_candidates) == 1:
+                        max_length = max_length_candidates[0]
+                    else:
+                        max_length = random.choice(max_length_candidates)
+
+                # remove input_ids that have been saved in outputs
+                input_ids = input_ids[cursor:]
+
+            # all needed data have been collected
+            if len(max_length_candidates) == 0:
+                break
+
+        tqdm_bar.close()
+
+        if len(max_length_candidates) > 0:
+            logger.warning(f"There are not enough data ! The remainings are {num_instances_this_dataset} instances for {dataset_name} dataset. Consider increase the corresponding data in {data_files}.")
+
+    dataset = datasets.Dataset.from_dict(outputs)
+    
+    return dataset
+
+
+
+if __name__ == "__main__":
+    parser = HfArgumentParser([Args])
+    args: Args = parser.parse_args_into_dataclasses()[0]
+
+    accelerator = Accelerator(cpu=args.cpu, kwargs_handlers=[InitProcessGroupKwargs(timeout=timedelta(days=10))])
+    # this script may be executed in DDP, so we make sure the dataset is create only on the main process
+    if accelerator.process_index == 0:
+        tokenizer = get_model_and_tokenizer(args, return_tokenizer_only=True)
+
+        if args.add_eos:
+            assert tokenizer.eos_token_id is not None, "Make sure the eos_token_id is not None when enabling add_eos!"
+
+        with open(args.config, encoding="utf-8") as f:
+            config = json.load(f)
+
+        length_2_num_token = {}
+        for x in args.num_token:
+            length, ntok = x.split(":")
+            length = int(length)
+
+            if ntok.lower().endswith("t"):
+                ntok = float(ntok[:-1]) * 1e12
+            elif ntok.lower().endswith("b"):
+                ntok = float(ntok[:-1]) * 1e9
+            elif ntok.lower().endswith("m"):
+                ntok = float(ntok[:-1]) * 1e6
+            else:
+                raise ValueError(f"Make sure num_token ends with T/t/B/b/M/m!")
+
+            length_2_num_token[length] = ntok
+
+        pretrain_dataset = prepare_pretrain_data(
+            args.train_data, 
+            tokenizer=tokenizer,
+            config=config,
+            length_2_num_token=length_2_num_token,
+            add_bos=args.add_bos,
+            add_eos=args.add_eos,
+            seed=args.seed,
+            cache_dir=args.dataset_cache_dir,
+        )
+        
+        logger.info(f"Saving dataset to {args.output_dir}...")
+        pretrain_dataset.save_to_disk(args.output_dir)
+
+    accelerator.wait_for_everyone()

Long_LLM/activation_beacon/new/main/train.py

+import logging
+from transformers import HfArgumentParser
+from transformers.integrations import is_deepspeed_zero3_enabled
+from src import ( 
+    Data,
+    DefaultDataCollator,
+    ModelArgs,
+    FileLogger,
+    get_model_and_tokenizer,
+    makedirs,
+    format_numel_str
+)
+from src.args import TrainingArgs
+from src.metrics import Metric
+from src.trainer import ActivationBeaconTrainer
+
+logger = logging.getLogger(__name__)
+
+
+def main():
+    parser = HfArgumentParser([ModelArgs, TrainingArgs])
+    model_args, training_args = parser.parse_args_into_dataclasses()
+
+    model, tokenizer = get_model_and_tokenizer(model_args, device="cuda", evaluation_mode=False)
+
+    if model_args.enable_beacon and training_args.only_train_beacon:
+        for name, param in model.named_parameters():
+            if "beacon" not in name:
+                param.requires_grad_(False)
+
+    if training_args.lora_tune:
+        from peft import (
+            LoraConfig,
+            get_peft_model,
+        )
+        # copied from LongLoRA
+        config = LoraConfig(
+            r=training_args.lora_rank,
+            lora_alpha=training_args.lora_alpha,
+            target_modules=training_args.lora_targets,
+            modules_to_save=training_args.lora_extra_params,
+            lora_dropout=training_args.lora_dropout,
+            bias="none",
+            task_type="CAUSAL_LM",
+        )
+        model = get_peft_model(model, config)
+
+    logger.info(f"Trainable Model params: {format_numel_str(sum(p.numel() for p in model.parameters() if p.requires_grad))}")
+
+    with training_args.main_process_first():
+        train_dataset = Data.prepare_train_data(
+            model_args.train_data, 
+            tokenizer=tokenizer,
+            max_length=model_args.max_length,
+            min_length=training_args.min_length,
+            chat_template=model_args.chat_template,
+            seed=training_args.seed,
+            cache_dir=model_args.dataset_cache_dir,
+        )
+
+    with training_args.main_process_first():
+        if is_deepspeed_zero3_enabled() and training_args.eval_method != "perplexity":
+            logger.warning(f"In deepspeed zero3, evaluation with generation is may lead to hang because of the unequal number of forward passes across different devices.")
+        eval_dataset = Data.prepare_eval_data(
+            model_args.eval_data, 
+            tokenizer=tokenizer,
+            max_length=training_args.eval_max_length,
+            min_length=training_args.eval_min_length,
+            chat_template=model_args.chat_template,
+            seed=training_args.seed,
+            cache_dir=model_args.dataset_cache_dir,
+        )
+
+    trainer = ActivationBeaconTrainer(
+        model=model,
+        tokenizer=tokenizer,
+        args=training_args,
+        model_args=model_args,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        data_collator=DefaultDataCollator(tokenizer),
+        file_logger=FileLogger(makedirs(training_args.log_path)),
+        compute_metrics=Metric.get_metric_fn(
+            metrics=training_args.metrics,
+            save_path=Metric.get_save_path(
+                model_args.eval_data,
+                training_args.output_dir
+            ) if model_args.eval_data is not None else None
+        )
+    )
+    if train_dataset is not None:
+        trainer.train()
+    elif eval_dataset is not None:
+        trainer.evaluate()
+
+if __name__ == "__main__":
+    main()

Long_LLM/activation_beacon/new/main/vllm_symlink.py

+import json
+import os
+import pathlib
+import shutil
+from argparse import ArgumentParser
+
+
+if __name__ == "__main__":
+    parser = ArgumentParser()
+    parser.add_argument("model_folder", type=str)
+    args = parser.parse_args()
+    
+    folder = args.model_folder.rstrip(os.sep)
+    path = pathlib.Path(folder)
+    parent = path.parent
+    name = path.name
+    folder_extrapolation = os.path.join(parent, f"extrapolation-{name}")
+    folder_yarn_4 = os.path.join(parent, f"yarn-4-{name}")
+    folder_yarn_8 = os.path.join(parent, f"yarn-8-{name}")
+
+    if os.path.exists(folder_extrapolation):
+        shutil.rmtree(folder_extrapolation)
+    if os.path.exists(folder_yarn_4):
+        shutil.rmtree(folder_yarn_4)
+    if os.path.exists(folder_yarn_8):
+        shutil.rmtree(folder_yarn_8)
+
+    os.makedirs(folder_extrapolation)
+    os.makedirs(folder_yarn_4)
+    os.makedirs(folder_yarn_8)
+
+    for name in os.listdir(folder):
+        if name == "config.json":
+            with open(os.path.join(folder, name), "r", encoding="utf-8") as f:
+                config = json.load(f)
+
+            extrapolation_config = config.copy()
+            extrapolation_config["max_position_embeddings"] = extrapolation_config["max_position_embeddings"] * 8
+            if "sliding_window" in extrapolation_config and extrapolation_config["sliding_window"] is not None:
+                extrapolation_config["sliding_window"] = extrapolation_config["max_position_embeddings"]
+            with open(os.path.join(folder_extrapolation, name), "w", encoding="utf-8") as f:
+                json.dump(extrapolation_config, f)
+
+            yarn_4_config = config.copy()
+            yarn_4_config["rope_scaling"] = {
+                "type": "yarn",
+                "factor": 4,
+                "original_max_position_embeddings": yarn_4_config["max_position_embeddings"]
+            }
+            with open(os.path.join(folder_yarn_4, name), "w", encoding="utf-8") as f:
+                json.dump(yarn_4_config, f)
+            
+            yarn_8_config = config.copy()
+            yarn_8_config["rope_scaling"] = {
+                "type": "yarn",
+                "factor": 8,
+                "original_max_position_embeddings": yarn_8_config["max_position_embeddings"]
+            }
+            with open(os.path.join(folder_yarn_8, name), "w", encoding="utf-8") as f:
+                json.dump(yarn_8_config, f)
+
+        else:
+            src = os.path.join(folder, name)
+
+            dest = os.path.join(folder_extrapolation, name)
+            if os.path.exists(dest):
+                os.remove(dest)
+            os.symlink(src, dest)
+
+            dest = os.path.join(folder_yarn_4, name)
+            if os.path.exists(dest):
+                os.remove(dest)
+            os.symlink(src, dest)
+            
+            dest = os.path.join(folder_yarn_8, name)
+            if os.path.exists(dest):
+                os.remove(dest)
+            os.symlink(src, dest)

Long_LLM/activation_beacon/new/src/__init__.py

+from .utils import FileLogger, DefaultDataCollator, makedirs, split_file_dir_name_ext, clear_dir, get_max_length_in_nested_lists, pad_nested_lists, mask_nested_lists, normalize_text, wrap_text, load_json, save_json, load_pickle, save_pickle, add_eos, remove_eos, format_numel_str
+from .chat import apply_chat_template
+from .args import ModelArgs
+from .data import Data
+from .modeling_utils import evaluate_perplexity, evaluate_generation, evaluate_nll, move_to_device
+
+import logging
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+    datefmt="%m/%d/%Y %H:%M:%S",
+)
+
+
+def get_model_and_tokenizer(model_args, device="cpu", evaluation_mode=True, return_tokenizer_only=False, **kwargs):
+    import torch
+    import transformers
+    from dataclasses import asdict
+    from transformers import AutoTokenizer, AutoConfig, AutoModelForCausalLM, BitsAndBytesConfig
+    from transformers.utils import logging
+    from transformers.integrations import is_deepspeed_zero3_enabled
+    from packaging import version
+
+    from .args import ModelArgs
+
+    logger = logging.get_logger(__name__)
+
+    model_args: ModelArgs
+
+    model_args_dict = asdict(model_args)
+    model_args_dict.update(**kwargs)
+    
+    model_name_or_path = model_args_dict["model_name_or_path"]
+    cache_dir = model_args_dict["model_cache_dir"]
+    access_token = model_args_dict["access_token"]
+
+    logger.info(f"Loading model and tokenizer from {model_name_or_path}...")
+
+    tokenizer_kwargs = {}
+    if model_args_dict["no_use_fast"]:
+        tokenizer_kwargs = {"use_fast": False}
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_name_or_path, 
+        cache_dir=cache_dir, 
+        padding_side=model_args_dict["padding_side"], 
+        token=access_token, 
+        trust_remote_code=True,
+        **tokenizer_kwargs
+    )
+    if tokenizer.pad_token_id is None:
+        tokenizer.pad_token_id = tokenizer.eos_token_id
+    
+    if return_tokenizer_only:
+        return tokenizer
+
+    dtype = model_args_dict["dtype"]
+    if dtype == "bf16":
+        dtype = torch.bfloat16
+    elif dtype == "fp16":
+        dtype = torch.float16
+    else:
+        dtype = torch.float32
+        
+    device_map = model_args_dict["device_map"]
+    if device_map is None and not is_deepspeed_zero3_enabled():
+        device_map = {"": device}
+    
+    rope_kwargs = {}
+    rope_theta = model_args_dict["rope_theta"]
+    if rope_theta is not None:
+        rope_kwargs["rope_theta"] = rope_theta
+    rope_method = model_args_dict["rope_method"]
+    if rope_method is not None:
+        rope_factor = model_args_dict["rope_factor"]
+        rope_scaling = {
+            "type": rope_method,
+            "factor": rope_factor
+        }
+        # NOTE: do not destroy the default rope_scaling of the model
+        rope_kwargs["rope_scaling"] = rope_scaling
+
+    attn_kwargs = {}
+    attn_impl = model_args_dict["attn_impl"]
+    if attn_impl is not None:
+        if version.parse(transformers.__version__) <= version.parse("4.36"):
+            if attn_impl == "flash_attention_2":
+                attn_kwargs["use_flash_attention_2"] = True
+        else:
+            attn_kwargs["attn_implementation"] = attn_impl
+
+    # from_pretrained_kwargs = {}
+    # if attn_impl == "flash_attention_2" and version.parse(transformers.__version__) <= version.parse("4.36"):
+    #     from_pretrained_kwargs["use_flash_attention_2"] = True
+
+    beacon_kwargs = {}
+    for k, v in model_args_dict.items():
+        if k.startswith("beacon") and v is not None:
+            beacon_kwargs[k] = v
+        elif k.startswith("retrieval") and v is not None:
+            beacon_kwargs[k] = v
+
+    # use architecture attribute to distinguish different models
+    probe_config = AutoConfig.from_pretrained(
+        model_name_or_path, 
+        cache_dir=cache_dir, 
+        token=access_token, 
+        trust_remote_code=True
+    )
+    architecture = probe_config.architectures[0]
+
+    extra_kwargs = {}
+    if model_args_dict["max_position_embeddings"] is not None:
+        extra_kwargs["max_position_embeddings"] = model_args_dict["max_position_embeddings"]
+    if architecture == "MistralForCausalLM" and model_args_dict["mistral_sliding_window"] is not None:
+        extra_kwargs["sliding_window"] = model_args_dict["mistral_sliding_window"]
+    if model_args_dict["load_in_4_bit"]:
+        extra_kwargs["quantization_config"] = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_compute_dtype=dtype,
+        )
+        device_map = None
+
+    if model_args_dict["enable_beacon"]:
+        from .llama import LlamaForCausalLM, LlamaConfig
+        from .mistral import MistralForCausalLM, MistralConfig
+        from .qwen2 import Qwen2ForCausalLM, Qwen2Config
+        ARCHITECTURE_TO_CLASS = {
+            'LlamaForCausalLM': (LlamaConfig, LlamaForCausalLM),
+            'MistralForCausalLM': (MistralConfig, MistralForCausalLM),
+            'Qwen2ForCausalLM': (Qwen2Config, Qwen2ForCausalLM),
+        }
+
+        config_class, model_class = ARCHITECTURE_TO_CLASS[architecture]
+
+        config = config_class.from_pretrained(
+            model_name_or_path, 
+            cache_dir=cache_dir,
+            token=access_token,
+            # NOTE: keep the torch_dtype in config consistent with that in model
+            torch_dtype=dtype,
+            **beacon_kwargs,
+            **rope_kwargs,
+            **attn_kwargs,
+            **extra_kwargs,
+        )
+        model = model_class.from_pretrained(
+            model_name_or_path, 
+            config=config,
+            cache_dir=cache_dir, 
+            torch_dtype=dtype,
+            device_map=device_map, 
+            token=access_token,
+        )
+
+    else:
+        if model_args_dict["enable_vllm"]:
+            from .vllm_utils import HFStyleVllmModel
+            if model_args_dict["dtype"] == "fp32":
+                vllm_dtype = "float32"
+            elif model_args_dict["dtype"] == "fp16":
+                vllm_dtype = "float16"
+            elif model_args_dict["dtype"] == "bf16":
+                vllm_dtype = "bfloat16"
+
+            vllm_kwargs = {}
+            if model_args_dict["vllm_len"] is not None:
+                vllm_kwargs["max_model_len"] = model_args_dict["vllm_len"]
+
+            model = HFStyleVllmModel(
+                model=model_name_or_path,
+                dtype=vllm_dtype,
+                gpu_memory_utilization=model_args_dict["vllm_mem"],
+                tensor_parallel_size=model_args_dict["vllm_tp"],
+                disable_custom_all_reduce=model_args_dict["vllm_disable_ar"],
+                enforce_eager=False,
+                trust_remote_code=True,
+                **rope_kwargs,
+                **vllm_kwargs,
+            )
+
+        else:
+            model = AutoModelForCausalLM.from_pretrained(
+                model_name_or_path, 
+                cache_dir=cache_dir, 
+                torch_dtype=dtype,
+                device_map=device_map,
+                token=access_token,
+                trust_remote_code=True,
+
+                # NOTE: do not destroy the default rope_scaling of the model
+                **rope_kwargs,
+                **attn_kwargs,
+                **extra_kwargs,
+            )
+
+    # load lora
+    if model_args_dict["lora"] is not None:
+        logger.info(f"loading lora from {model_args_dict['lora']}...")
+
+        from peft import PeftModel
+        model = PeftModel.from_pretrained(
+            model, 
+            model_args_dict["lora"],
+            torch_dtype=dtype,
+            device_map=device_map,
+        )
+        if model_args_dict["lora_unload"]:
+            model = model.merge_and_unload()
+
+    if model_args_dict["enable_tp"]:
+        import tensor_parallel as tp
+        logger.info("enabling tensor parallelism...")
+        
+        # model = tp.tensor_parallel(model, device_ids=list(range(8)), distributed=False, sharded=False)
+        model = tp.tensor_parallel(model, sharded=True)
+
+        if model.generation_config.eos_token_id == 128001:
+            model.generation_config.eos_token_id = [128001, 128009]
+
+    if isinstance(model, transformers.modeling_utils.PreTrainedModel):
+        model = model.eval()
+        if evaluation_mode:
+            # NOTE: essential to disable all gradient in-place, so that when calling accelerator.prepare, the forward function will not be wrapped that may consume extra GPU memory
+            model.requires_grad_(False)
+        logger.info(model.config)
+
+    # override the default generation config
+    generation_config = model_args.get_generation_config()
+    if len(generation_config):
+        model.generation_config.update(**generation_config)
+    logger.info(f"Specified generation config: {generation_config}")
+
+    return model, tokenizer

Long_LLM/activation_beacon/new/src/args.py

+import os
+import json
+from dataclasses import dataclass, field, asdict
+from transformers.training_args import TrainingArguments
+from typing import Optional, List, Tuple, Union, Dict
+
+
+@dataclass
+class ModelArgs:
+    model_cache_dir: str = field(
+        default=None,
+        metadata={'help': 'Default path to save language models.'}
+    )
+    dataset_cache_dir: str = field(
+        default=None,
+        metadata={'help': 'Default path to save huggingface datasets.'}
+    )
+    data_root: str = field(
+        default="/data/long-llm", 
+        metadata={'help': 'The base directory storing all data used for training and evaluation. If specified, make sure all train_data, eval_data, and corpus are path relative to data_root!'},
+    )
+    train_data: Optional[List[str]] = field(
+        default=None,
+        metadata={'help': 'Training json file or glob to match a list of files.'},
+    )
+    eval_data: Optional[str] = field(
+        default=None,
+        metadata={'help': 'Evaluation json file.'},
+    )
+    
+    model_name_or_path: str = field(
+        default='meta-llama/Llama-2-7b-chat-hf',
+        metadata={'help': 'Path to pretrained model or model identifier from huggingface.co/models'}
+    )
+    padding_side: str = field(
+        default="left",
+        metadata={'help': 'Tokenizer padding side.'}
+    )
+    no_use_fast: bool = field(
+        default=False,
+        metadata={'help': 'Do not use fast tokenizer?'}
+    )
+    access_token: Optional[str] = field(
+        default=None,
+        metadata={'help': 'Huggingface access token.'}
+    )
+    attn_impl: Optional[str] = field(
+        default="flash_attention_2",
+        metadata={'help': 'The implementation of attention.'}
+    )
+
+    max_length: int = field(
+        default=4096,
+        metadata={'help': 'How many tokens at maximum for each input.'},
+    )
+    chat_template: str = field(
+        default="hf",
+        metadata={'help': 'Instruction template name in fastchat.'}
+    )
+
+    max_position_embeddings: Optional[int] = field(
+        default=None,
+        metadata={'help': 'Maximum position.'},
+    )
+    mistral_sliding_window: Optional[int] = field(
+        default=None,
+        metadata={'help': 'Sliding window size in Mistral models.'},
+    )
+    rope_theta: Optional[float] = field(
+        default=None,
+        metadata={'help': 'RoPE base (theta).'},
+    )
+    rope_method: Optional[str] = field(
+        default=None,
+        metadata={'help': 'How to scale RoPE? {linear, dynamic, yarn}'},
+    )
+    rope_factor: float = field(
+        default=1.,
+        metadata={'help': 'RoPE scaling factor.'},
+    )
+
+    lora: Optional[str] = field(
+        default=None,
+        metadata={'help': 'LoRA ID.'},
+    )
+    lora_unload: bool = field(
+        default=True,
+        metadata={'help': 'Merge and unload LoRA?'},
+    )
+    load_in_4_bit: bool = field(
+        default=False,
+        metadata={'help': 'Load model in 4 bits?'},
+    )
+
+    dtype: str = field(
+        default="bf16",
+        metadata={'help': 'Data type for embeddings.'}
+    )
+    device_map: Optional[str] = field(
+        default=None,
+        metadata={'help': 'Device map for loading the model. Set to auto to load across devices.'}
+    )
+    batch_size: int = field(
+        default=1,
+        metadata={'help': 'Evaluation batch size.'},
+    )
+    cpu: bool = field(
+        default=False,
+        metadata={'help': 'Use cpu?'}
+    )
+
+    enable_tp: bool = field(
+        default=False,
+        metadata={'help': 'Use tensor parallel to wrap the model?'}
+    )
+    
+    enable_vllm: bool = field(
+        default=False,
+        metadata={'help': 'Use vllm?'}
+    )
+    vllm_mem: float = field(
+        default=0.9,
+        metadata={'help': 'Vllm maximum GPU memory utilization.'}
+    )
+    vllm_tp: int = field(
+        default=1,
+        metadata={'help': 'Vllm tensor parallel degree.'}
+    )
+    vllm_len: Optional[int] = field(
+        default=None,
+        metadata={'help': 'Vllm maximum sequence length.'}
+    )
+    vllm_disable_ar: bool = field(
+        default=False,
+        metadata={'help': 'Disable custom all-reduce in vllm?'}
+    )
+
+    enable_beacon: bool = field(
+        default=False,
+        metadata={'help': 'Enable activation beacon?'}
+    )
+    beacon_window: Optional[int] = field(
+        default=None,
+        metadata={'help': 'The initial sliding window size.'}
+    )
+    beacon_stride: Optional[int] = field(
+        default=None,
+        metadata={'help': 'The stride of the sliding window.'}
+    )
+    beacon_attn: Optional[str] = field(
+        default=None,
+        metadata={'help': 'How to assign attention masks of beacon tokens? {segmentation, step-expansion, full-converage}'}
+    )
+    beacon_ratio: Optional[List[int]] = field(
+        default=None,
+        metadata={'help': 'Condensing ratios for beacons.'}
+    )
+    beacon_ratio_mix: Optional[str] = field(
+        default=None,
+        metadata={'help': 'How to determine the beacon_ratio for each input. {step-random, instance-random, adapt-x}'}
+    )
+    beacon_param: Optional[List[str]] = field(
+        default=None,
+        metadata={'help': 'The introduced parameters for beacon.'}
+    )
+    beacon_embed_init: str = field(
+        default="eos",
+        metadata={'help': 'Initialize beacon embedding from eos/bos embedding.'}
+    )
+    beacon_sink_size: Optional[int] = field(
+        default=None,
+        metadata={'help': 'The number of activations that are always kept in the head of the sequence according to StreamingLLM.'}
+    )
+    beacon_attend_prev: Optional[bool] = field(
+        default=None,
+        metadata={'help': 'Can beacon tokens attend to previous beacon tokens?'}
+    )
+    beacon_pos: Optional[str] = field(
+        default=None,
+        metadata={'help': 'Where to put beacon tokens? {append, interleave}'}
+    )
+    beacon_parallel_window: Optional[int] = field(
+        default=None,
+        metadata={'help': 'How many windows to run in parallel?'}
+    )
+    retrieval_method: Optional[str] = field(
+        default=None,
+        metadata={'help': 'How to retrieve? {bm25}'}
+    )
+    retrieval_topk: Optional[int] = field(
+        default=None,
+        metadata={'help': 'How many windows to retrieve?'}
+    )
+    retrieval_key_length: Optional[int] = field(
+        default=None,
+        metadata={'help': 'The key sequence length in retrieval.'}
+    )
+
+    max_new_tokens: Optional[int] = field(
+        default=None,
+        metadata={'help': 'How many tokens at maximum to return?'},
+    )
+    do_sample: Optional[bool] = field(
+        default=None,
+        metadata={'help': 'Do sampling when decoding?'},
+    )
+    temperature: Optional[float] = field(
+        default=None,
+        metadata={'help': 'Sampling temperature.'},
+    )
+    top_p: Optional[float] = field(
+        default=None,
+        metadata={'help': "If set to float < 1, only the smallest set of most probable tokens with probabilities that add up to `top_p` or higher are kept for generation."}
+    )
+
+    def resolve_path(self, path):
+        """Resolve any path starting with 'long-llm:' to relative path against data_root."""
+        pattern = "long-llm:"
+        # resolve relative data paths when necessary
+        if isinstance(path, list):
+            for i, x in enumerate(path):
+                if x.startswith(pattern):
+                    path[i] = os.path.join(self.data_root, x.replace(pattern, ""))
+        else:
+            if path.startswith(pattern):
+                path = os.path.join(self.data_root, path.replace(pattern, ""))
+
+        return path
+    
+    def get_generation_config(self):
+        generation_config = {}
+        if self.max_new_tokens is not None:
+            generation_config["max_new_tokens"] = self.max_new_tokens
+        if self.do_sample is not None:
+            generation_config["do_sample"] = self.do_sample
+        if self.temperature is not None:
+            generation_config["temperature"] = self.temperature
+        if self.top_p is not None:
+            generation_config["top_p"] = self.top_p
+        return generation_config
+
+    def to_dict(self):
+        return asdict(self)
+
+    def save(self, path):
+        with open(path, "w", encoding="utf-8") as f:
+            json.dump(self.to_dict(), f)
+
+    def __post_init__(self):
+        if self.train_data is not None:
+            self.train_data = self.resolve_path(self.train_data)
+
+        if self.eval_data is not None:
+            self.eval_data = self.resolve_path(self.eval_data)
+
+        if hasattr(self, "output_dir") and self.output_dir is not None:
+            self.output_dir = self.resolve_path(self.output_dir)
+
+        if hasattr(self, "result_dir"):
+            if self.result_dir is None: 
+                if self.lora is not None:
+                    name_or_path_components = [x for x in self.lora.split("/") if len(x)][-2:]
+                else:
+                    name_or_path_components = [x for x in self.model_name_or_path.split("/") if len(x)][-2:]
+                self.result_dir = os.path.join(*name_or_path_components)
+            else:
+                self.result_dir = self.resolve_path(self.result_dir)
+
+
+@dataclass
+class TrainingArgs(TrainingArguments):
+    # ==============================
+    # Common arguments
+    # ==============================
+    output_dir: str = field(
+        default="data/outputs/pretrain",
+    )
+
+    per_device_train_batch_size: int = field(
+        default=1,
+        metadata={'help': 'Train batch size.'}
+    )
+    per_device_eval_batch_size: int = field(
+        default=1,
+        metadata={'help': 'Evaluation batch size.'}
+    )
+    remove_unused_columns: bool = field(
+        default=False,
+        metadata={'help': 'Remove columns in the dataset that are not registered in the forward function?'}
+    )
+    ddp_find_unused_parameters: bool = field(
+        default=False,
+        metadata={'help': 'Find unusuable parameters?'}
+    )
+    # NOTE: essential to keep comuputation graph because we need gradients for beacon tokens
+    use_reentrant: Optional[bool] = field(
+        default=None,
+        metadata={'help': "Use reetrant in gradient checkpointing?"}
+    )
+    report_to: str = field(
+        default="none",
+        metadata={'help': 'Log results by external tools?'}
+    )
+
+    # ==============================
+    # Customized arguments
+    # ==============================
+    min_length: int = field(
+        default=0,
+        metadata={'help': 'How many tokens at minimum for training?'}
+    )
+
+    group_by_stride: Optional[str] = field(
+        default=None,
+        metadata={'help': 'Group the training data instances by the number of strides in the beacon model. {relaxed, strict}'}
+    )
+    sort_by_stride: Optional[str] = field(
+        default=None,
+        metadata={'help': 'Sort the training data instances by the number of strides in the beacon model. {ascend, descend}'}
+    )
+    only_train_beacon: bool = field(
+        default=True,
+        metadata={'help': 'Freeze LLM parameters when training beacon parameters?'}
+    )
+    
+    eval_method: str = field(
+        default="perplexity",
+        metadata={'help': 'How to evaluate during training? {perplexity, generation}'}
+    )
+    eval_max_length: int = field(
+        default=4096,
+        metadata={'help': 'How many tokens at maximum for each input in evaluation.'},
+    )
+    eval_min_length: int = field(
+        default=512,
+        metadata={'help': 'How many tokens at minimum for each input in evaluation.'},
+    )
+    eval_beacon_ratio: List[int] = field(
+        default_factory=lambda: [32],
+        metadata={'help': 'Condensing ratios for beacons in evaluation.'}
+    )
+    eval_beacon_ratio_mix: str = field(
+        default="adapt-1024",
+        metadata={'help': 'How to determine the beacon_ratio for each input. {step-random, instance-random, adapt-x}'}
+    )
+    max_eval_num: Optional[int] = field(
+        default=None,
+        metadata={'help': 'How many samples for validation?'}
+    )
+
+    lora_tune: bool = field(
+        default=False,
+        metadata={"help": "Use LoRA fine-tuning?"},
+    )
+    lora_rank: int = field(
+        default=32,
+        metadata={'help': 'LoRA rank.'}
+    )
+    lora_alpha: int = field(
+        default=16,
+        metadata={'help': 'LoRA scaling factor.'}
+    )
+    lora_dropout: float = field(
+        default=0.,
+        metadata={'help': 'LoRA dropout p.'}
+    )
+    lora_targets: List[str] = field(
+        default_factory=lambda: ["q_proj", "k_proj", "v_proj", "o_proj"],
+        metadata={"help": "Module name patterns to add LoRA."},
+    )
+    lora_extra_params: List[str] = field(
+        default_factory=lambda: ["embed_tokens", "norm"],
+        metadata={"help": "Extra trainable parameters except LoRA weights, if low rank training."},
+    )
+
+    metrics: List[str] = field(
+        default_factory=lambda: [],
+        metadata={'help': 'List of metrics. {rouge, save_result}'}
+    )
+    log_path: str = field(
+        default="data/outputs/metrics.log",
+        metadata={'help': 'Log file path.'}
+    )
+
+
+    def __post_init__(self):
+        if self.use_reentrant is not None:
+            self.gradient_checkpointing_kwargs = {"use_reentrant": self.use_reentrant}
+        return super().__post_init__()

Long_LLM/activation_beacon/new/src/chat.py

+"""
+Copied from fastchat.
+"""
+
+import base64
+import dataclasses
+from enum import auto, IntEnum
+from io import BytesIO
+from typing import List, Any, Dict, Union, Tuple
+
+import numpy as np
+from copy import deepcopy
+from transformers.tokenization_utils import PreTrainedTokenizer, BatchEncoding
+
+
+@dataclasses.dataclass
+class ChatTemplateOutput:
+    raw: str = None
+    encoded: BatchEncoding = None
+
+
+def mask_nested_lists(lst, mask_target, mask_value=0):
+    if isinstance(lst[0], list):
+        for i, elem in enumerate(lst):
+            lst[i] = mask_nested_lists(elem, mask_target, mask_value)
+        return lst
+    else:
+        return [x if x != mask_target else mask_value for x in lst]
+
+
+def apply_chat_template(template, messages, system_message=None, tokenizer:PreTrainedTokenizer=None, add_generation_prompt=False, return_labels=False, **tokenization_kwargs):
+    """
+    Wrap the message using the template from fastchat according to its role
+
+    Args:
+        template: fastchat template name
+        messages: a list of dictionaries, each of which is {'role': 'user/assistant', 'content': 'xxx'}
+        system_message: system input
+    """
+    if len(tokenization_kwargs):
+        assert tokenizer is not None, f"Make sure the tokenizer is not None when passing tokenizer kwargs!"
+
+    if template == "no":
+        assert tokenizer is not None, f"Make sure the tokenizer is not None when template is no!"
+
+        prev_role = None
+        conversation = ""
+
+        for i, message in enumerate(messages):
+            role = message['role']
+            content = message['content']
+            if prev_role == role:
+                raise ValueError(f"Current role (idx={i}) {role} and previous role {messages[i-1]['role']} are the same!")
+            
+            if i == 0:
+                content = tokenizer.decode(tokenizer.encode(content), skip_special_tokens=True)
+                user_message = content
+            elif i == 1:
+                # we use a space to separate user message and assistant response
+                content = ' ' + content + tokenizer.eos_token
+                assistant_message = content
+            else:
+                raise ValueError(f"Please use chat template when there are multi-turn conversations")
+
+            conversation += content
+
+        encoded = tokenizer(conversation, **tokenization_kwargs)
+
+        if return_labels:
+            labels = encoded['input_ids'].copy()
+            assistant_message_len = len(tokenizer.encode(assistant_message.lstrip(), add_special_tokens=False))
+            labels[:-assistant_message_len] = [-100 for _ in labels[:-assistant_message_len]]
+            encoded["labels"] = labels
+
+            # sanity check
+            for id_, label_ in zip(encoded['input_ids'], encoded['labels']):
+                assert id_ == label_ or label_ == -100, f"Found mismatch input_ids and labels!"
+
+        return ChatTemplateOutput(raw=conversation, encoded=encoded)
+
+    elif template == "hf":
+        assert return_labels == False, f"Returning labels with hf template is currently unsupported."
+        tokenization_kwargs["return_dict"] = True
+        raw = tokenizer.apply_chat_template(messages, add_generation_prompt=add_generation_prompt, tokenize=False)
+        encoded = tokenizer.apply_chat_template(messages, add_generation_prompt=add_generation_prompt,**tokenization_kwargs)
+        # for some tokenizer, the encoded input_ids are wrapped in a big list, while others are not
+        if isinstance(encoded['input_ids'][0], list):
+            for k, v in encoded.items():
+                encoded[k] = v[0]
+        return ChatTemplateOutput(raw=raw, encoded=encoded)
+
+    conversation_template = get_conv_template(template)
+    if system_message is not None:
+        conversation_template.set_system_message(system_message)
+    
+    config = {
+        'mistral': {
+            # separator for different conversation turns (one turn consists of an utterance from user and a response from assistant)
+            "turn_sep": "</s>",
+            # separator for different roles within each turn
+            "role_sep": " [/INST]",
+            # the number of tokens in the beginning of the entire sequence, usually the length of the bos string
+            "begin_of_text_len": 1,
+            # the number of tokens to offset in the beginning of each turn, these tokens should be masked
+            "turn_seq_left_offset": 0,
+        },
+        'llama-2': {
+            # separator for different conversation turns (one turn consists of an utterance from user and a response from assistant)
+            "turn_sep": " </s><s>",
+            # separator for different roles within each turn
+            "role_sep": " [/INST]",
+            # the number of tokens in the beginning of the entire sequence, usually the length of the bos string
+            "begin_of_text_len": 1,
+            # the number of tokens to offset in the beginning of each turn, these tokens should be masked
+            "turn_seq_left_offset": -1,
+        },
+        'llama-3': {
+            # separator for different conversation turns (one turn consists of an utterance from user and a response from assistant)
+            "turn_sep": "<|eot_id|><|start_header_id|>user<|end_header_id|>\n\n",
+            # separator for different roles within each turn
+            "role_sep": "<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n",
+            # the number of tokens in the beginning of the entire sequence, usually the length of the bos string
+            "begin_of_text_len": 1,
+            # the number of tokens to offset in the beginning of each turn, these tokens should be masked
+            "turn_seq_left_offset": -4,
+        },
+        'qwen': {
+            # separator for different conversation turns (one turn consists of an utterance from user and a response from assistant)
+            "turn_sep": "<|im_start|>user\n",
+            # separator for different roles within each turn
+            "role_sep": "<|im_end|>\n<|im_start|>assistant\n",
+            # the number of tokens in the beginning of the entire sequence, usually the length of the bos string
+            "begin_of_text_len": 0,
+            # the number of tokens to offset in the beginning of each turn, these tokens should be masked
+            "turn_seq_left_offset": -4,
+        }
+    }[template]
+
+    role_map = {
+        'user': conversation_template.roles[0],
+        'assistant': conversation_template.roles[1]
+    }
+    prev_role = None
+
+    for i, message in enumerate(messages):
+        role = role_map[message['role']]
+        content = message['content']
+        if prev_role == role:
+            raise ValueError(f"Current role (idx={i}) {role} and previous role {messages[i-1]['role']} are the same!")
+        conversation_template.append_message(role, content)
+        prev_role = role
+    
+    if add_generation_prompt:
+        assert prev_role == role_map['user'], f"You cannot add generation prompt after assistant output!"
+        conversation_template.append_message(role_map['assistant'], None)
+
+    conversation = conversation_template.get_prompt()
+
+    if tokenizer is not None:
+        encoded = tokenizer(conversation, **tokenization_kwargs)
+
+        if return_labels:
+            # Mask targets. Only compute loss on the assistant outputs.
+
+            turn_sep = config["turn_sep"]
+            role_sep = config["role_sep"]
+            begin_of_text_len = config["begin_of_text_len"]
+            turn_seq_left_offset = config["turn_seq_left_offset"]
+            turn_sep_len = len(tokenizer.encode(turn_sep, add_special_tokens=False))
+
+            # transform to array for fast value assignment
+            labels = deepcopy(encoded['input_ids'])
+            labels = np.array(labels)
+            total_len = len(labels)
+
+            turns = conversation.split(turn_sep)
+
+            cur_len = 0
+            for i, turn in enumerate(turns):
+                if turn == "":
+                    break
+
+                turn_len = len(tokenizer(turn, add_special_tokens=False).input_ids)
+
+                parts = turn.split(role_sep)
+
+                if len(parts) == 2:
+                    user_message, assistant_message = parts
+
+                    user_message += role_sep
+                    instruction_len = len(tokenizer(user_message, add_special_tokens=False).input_ids)
+
+                    # for bos tokens
+                    if i == 0:
+                        turn_len += begin_of_text_len
+                        instruction_len += begin_of_text_len
+
+                    # Ignore the user instructions
+                    labels[max(cur_len + turn_seq_left_offset, 0): cur_len + instruction_len] = -100
+                
+                else:
+                    labels[max(cur_len + turn_seq_left_offset, 0): cur_len + turn_len + turn_sep_len] = -100
+                    
+                cur_len = cur_len + turn_len + turn_sep_len
+
+                if cur_len > total_len:
+                    break
+
+            labels[max(cur_len + turn_seq_left_offset, 0):] = -100
+
+            encoded['labels'] = labels.tolist()
+
+            # sanity check
+            for id_, label_ in zip(encoded['input_ids'], encoded['labels']):
+                assert id_ == label_ or label_ == -100, f"Found mismatch input_ids and labels!"
+
+    else:
+        encoded = None
+
+    return ChatTemplateOutput(raw=conversation, encoded=encoded)
+
+
+class SeparatorStyle(IntEnum):
+    """Separator styles."""
+
+    ADD_COLON_SINGLE = auto()
+    ADD_COLON_TWO = auto()
+    ADD_COLON_SPACE_SINGLE = auto()
+    NO_COLON_SINGLE = auto()
+    NO_COLON_TWO = auto()
+    ADD_NEW_LINE_SINGLE = auto()
+    LLAMA2 = auto()
+    LLAMA3 = auto()
+    CHATGLM = auto()
+    CHATML = auto()
+    CHATINTERN = auto()
+    DOLLY = auto()
+    RWKV = auto()
+    PHOENIX = auto()
+    ROBIN = auto()
+    FALCON_CHAT = auto()
+    CHATGLM3 = auto()
+    DEEPSEEK_CHAT = auto()
+    METAMATH = auto()
+    YUAN2 = auto()
+    GEMMA = auto()
+    CLLM = auto()
+    DEFAULT = auto()
+
+
+IMAGE_PLACEHOLDER_STR = "$$<image>$$"
+
+
+@dataclasses.dataclass
+class Conversation:
+    """A class that manages prompt templates and keeps all conversation history."""
+
+    # The name of this template
+    name: str
+    # The template of the system prompt
+    system_template: str = "{system_message}"
+    # The system message
+    system_message: str = ""
+    # The names of two roles
+    roles: Tuple[str] = ("USER", "ASSISTANT")
+    # All messages. Each item is (role, message).
+    # Each message is either a string or a tuple of (string, List[image_url]).
+    messages: List[List[str]] = ()
+    # The number of few shot examples
+    offset: int = 0
+    # The separator style and configurations
+    sep_style: SeparatorStyle = SeparatorStyle.ADD_COLON_SINGLE
+    sep: str = "\n"
+    sep2: str = None
+    # Stop criteria (the default one is EOS token)
+    stop_str: Union[str, List[str]] = None
+    # Stops generation if meeting any token in this list
+    stop_token_ids: List[int] = None
+
+    def get_prompt(self) -> str:
+        """Get the prompt for generation."""
+        system_prompt = self.system_template.format(system_message=self.system_message)
+        if self.sep_style == SeparatorStyle.ADD_COLON_SINGLE:
+            ret = system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    ret += role + ": " + message + self.sep
+                else:
+                    ret += role + ":"
+            return ret
+        elif self.sep_style == SeparatorStyle.ADD_COLON_TWO:
+            seps = [self.sep, self.sep2]
+            ret = system_prompt + seps[0]
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    if type(message) is tuple:
+                        message, images = message
+                        message = IMAGE_PLACEHOLDER_STR * len(images) + message
+                    ret += role + ": " + message + seps[i % 2]
+                else:
+                    ret += role + ":"
+            return ret
+        elif self.sep_style == SeparatorStyle.ADD_COLON_SPACE_SINGLE:
+            ret = system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    ret += role + ": " + message + self.sep
+                else:
+                    ret += role + ": "  # must be end with a space
+            return ret
+        elif self.sep_style == SeparatorStyle.ADD_NEW_LINE_SINGLE:
+            ret = "" if system_prompt == "" else system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    ret += role + "\n" + message + self.sep
+                else:
+                    ret += role + "\n"
+            return ret
+        elif self.sep_style == SeparatorStyle.NO_COLON_SINGLE:
+            ret = system_prompt
+            for role, message in self.messages:
+                if message:
+                    ret += role + message + self.sep
+                else:
+                    ret += role
+            return ret
+        elif self.sep_style == SeparatorStyle.NO_COLON_TWO:
+            seps = [self.sep, self.sep2]
+            ret = system_prompt
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += role + message + seps[i % 2]
+                else:
+                    ret += role
+            return ret
+        elif self.sep_style == SeparatorStyle.RWKV:
+            ret = system_prompt
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += (
+                        role
+                        + ": "
+                        + message.replace("\r\n", "\n").replace("\n\n", "\n")
+                    )
+                    ret += "\n\n"
+                else:
+                    ret += role + ":"
+            return ret
+        elif self.sep_style == SeparatorStyle.LLAMA2:
+            seps = [self.sep, self.sep2]
+            if self.system_message:
+                ret = system_prompt
+            else:
+                ret = "[INST] "
+            for i, (role, message) in enumerate(self.messages):
+                tag = self.roles[i % 2]
+                if message:
+                    if i == 0:
+                        ret += message + " "
+                    else:
+                        ret += tag + " " + message + seps[i % 2]
+                else:
+                    ret += tag
+            return ret
+        elif self.sep_style == SeparatorStyle.LLAMA3:
+            # ret = "<|begin_of_text|>"
+            if self.system_message:
+                ret = system_prompt
+            else:
+                ret = ""
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += f"<|start_header_id|>{role}<|end_header_id|>\n\n"
+                    ret += f"{message}<|eot_id|>"
+                else:
+                    ret += f"<|start_header_id|>{role}<|end_header_id|>\n\n"
+            return ret
+        elif self.sep_style == SeparatorStyle.CHATGLM:
+            # source: https://huggingface.co/THUDM/chatglm-6b/blob/1d240ba371910e9282298d4592532d7f0f3e9f3e/modeling_chatglm.py#L1302-L1308
+            # source2: https://huggingface.co/THUDM/chatglm2-6b/blob/e186c891cf64310ac66ef10a87e6635fa6c2a579/modeling_chatglm.py#L926
+            round_add_n = 1 if self.name == "chatglm2" else 0
+            if system_prompt:
+                ret = system_prompt + self.sep
+            else:
+                ret = ""
+
+            for i, (role, message) in enumerate(self.messages):
+                if i % 2 == 0:
+                    ret += f"[Round {i//2 + round_add_n}]{self.sep}"
+
+                if message:
+                    ret += f"{role}：{message}{self.sep}"
+                else:
+                    ret += f"{role}："
+            return ret
+        elif self.sep_style == SeparatorStyle.CHATML:
+            ret = "" if system_prompt == "" else system_prompt + self.sep + "\n"
+            for role, message in self.messages:
+                if message:
+                    if type(message) is tuple:
+                        message, images = message
+                        message = IMAGE_PLACEHOLDER_STR * len(images) + message
+                    ret += role + "\n" + message + self.sep + "\n"
+                else:
+                    ret += role + "\n"
+            return ret
+        elif self.sep_style == SeparatorStyle.CHATGLM3:
+            ret = ""
+            if self.system_message:
+                ret += system_prompt
+            for role, message in self.messages:
+                if message:
+                    ret += role + "\n" + message
+                else:
+                    ret += role
+            return ret
+        elif self.sep_style == SeparatorStyle.CHATINTERN:
+            # source: https://huggingface.co/internlm/internlm-chat-7b-8k/blob/bd546fa984b4b0b86958f56bf37f94aa75ab8831/modeling_internlm.py#L771
+            seps = [self.sep, self.sep2]
+            ret = system_prompt
+            for i, (role, message) in enumerate(self.messages):
+                if i % 2 == 0:
+                    ret += "<s>"
+                if message:
+                    ret += role + ":" + message + seps[i % 2] + "\n"
+                else:
+                    ret += role + ":"
+            return ret
+        elif self.sep_style == SeparatorStyle.DOLLY:
+            seps = [self.sep, self.sep2]
+            ret = system_prompt
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += role + ":\n" + message + seps[i % 2]
+                    if i % 2 == 1:
+                        ret += "\n\n"
+                else:
+                    ret += role + ":\n"
+            return ret
+        elif self.sep_style == SeparatorStyle.PHOENIX:
+            ret = system_prompt
+            for role, message in self.messages:
+                if message:
+                    ret += role + ": " + "<s>" + message + "</s>"
+                else:
+                    ret += role + ": " + "<s>"
+            return ret
+        elif self.sep_style == SeparatorStyle.ROBIN:
+            ret = system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    ret += role + ":\n" + message + self.sep
+                else:
+                    ret += role + ":\n"
+            return ret
+        elif self.sep_style == SeparatorStyle.FALCON_CHAT:
+            ret = ""
+            if self.system_message:
+                ret += system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    ret += role + ": " + message + self.sep
+                else:
+                    ret += role + ":"
+            return ret
+        elif self.sep_style == SeparatorStyle.METAMATH:
+            ret = "" if system_prompt == "" else system_prompt + self.sep
+            for i, (role, message) in enumerate(self.messages):
+                # For MetaMath, sep2 is used to prefix the message.
+                starting_sep = ":\n" if i % 2 == 0 else ": " + self.sep2
+                ending_sep = self.sep if i % 2 == 0 else ""
+                if message:
+                    ret += role + starting_sep + message + ending_sep
+                else:
+                    ret += role + starting_sep
+            return ret
+        elif self.sep_style == SeparatorStyle.DEEPSEEK_CHAT:
+            seps = [self.sep, self.sep2]
+            ret = system_prompt
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += role + ": " + message + seps[i % 2]
+                else:
+                    ret += role + ":"
+            return ret
+        elif self.sep_style == SeparatorStyle.YUAN2:
+            seps = [self.sep, self.sep2]
+            ret = ""
+            if self.system_message:
+                ret += system_prompt + seps[1]
+            for _, message in self.messages:
+                if message:
+                    ret += message + "<n>"
+                else:
+                    ret += ""
+            ret = ret.rstrip("<n>") + seps[0]
+            return ret
+        elif self.sep_style == SeparatorStyle.GEMMA:
+            ret = "<bos>"
+            for role, message in self.messages:
+                if message:
+                    ret += "<start_of_turn>" + role + "\n" + message + self.sep
+                else:
+                    ret += "<start_of_turn>" + role + "\n"
+            return ret
+        elif self.sep_style == SeparatorStyle.CLLM:
+            seps = [self.sep, self.sep2]
+            ret = system_prompt + seps[0]
+            for i, (role, message) in enumerate(self.messages[-2:]):
+                if message:
+                    if type(message) is tuple:
+                        message, images = message
+                        message = IMAGE_PLACEHOLDER_STR * len(images) + message
+                    ret += role + ": " + message + seps[i % 2]
+                else:
+                    ret += role + ":"
+            return ret
+        elif self.sep_style == SeparatorStyle.DEFAULT:
+            ret = system_prompt + "\n"
+            for role, message in self.messages:
+                if message:
+                    ret += role + ": " + message + "\n"
+                else:
+                    ret += role + ":"
+            return ret
+        else:
+            raise ValueError(f"Invalid style: {self.sep_style}")
+
+    def get_images(self):
+        images = []
+        for i, (role, msg) in enumerate(self.messages[self.offset :]):
+            if i % 2 == 0:
+                if type(msg) is tuple:
+                    for image in msg[1]:
+                        images.append(image)
+
+        return images
+
+    def set_system_message(self, system_message: str):
+        """Set the system message."""
+        self.system_message = system_message
+
+    def get_system_message(self):
+        """return the system message."""
+        return self.system_message
+
+    def append_message(self, role: str, message: str):
+        """Append a new message."""
+        self.messages.append([role, message])
+
+    def update_last_message(self, message: str):
+        """Update the last output.
+
+        The last message is typically set to be None when constructing the prompt,
+        so we need to update it in-place after getting the response from a model.
+        """
+        self.messages[-1][1] = message
+
+    def convert_image_to_base64(self, image):
+        """Given an image, return the base64 encoded image string."""
+        from PIL import Image
+        import requests
+
+        # Load image if it has not been loaded in yet
+        if type(image) == str:
+            if image.startswith("http://") or image.startswith("https://"):
+                response = requests.get(image)
+                image = Image.open(BytesIO(response.content)).convert("RGB")
+            elif "base64" in image:
+                # OpenAI format is: data:image/jpeg;base64,{base64_encoded_image_str}
+                return image.split(",")[1]
+            else:
+                image = Image.open(image).convert("RGB")
+
+        max_hw, min_hw = max(image.size), min(image.size)
+        aspect_ratio = max_hw / min_hw
+        max_len, min_len = 2048, 2048
+        shortest_edge = int(min(max_len / aspect_ratio, min_len, min_hw))
+        longest_edge = int(shortest_edge * aspect_ratio)
+        W, H = image.size
+        if longest_edge != max(image.size):
+            if H > W:
+                H, W = longest_edge, shortest_edge
+            else:
+                H, W = shortest_edge, longest_edge
+            image = image.resize((W, H))
+
+        buffered = BytesIO()
+        image.save(buffered, format="PNG")
+        img_b64_str = base64.b64encode(buffered.getvalue()).decode()
+
+        return img_b64_str
+
+    def to_gradio_chatbot(self):
+        """Convert the conversation to gradio chatbot format."""
+        ret = []
+        for i, (role, msg) in enumerate(self.messages[self.offset :]):
+            if i % 2 == 0:
+                if type(msg) is tuple:
+                    msg, image = msg
+                    img_b64_str = image[0]  # Only one image on gradio at one time
+                    img_str = f'<img src="data:image/jpeg;base64,{img_b64_str}" alt="user upload image" />'
+                    msg = img_str + msg.replace("<image>\n", "").strip()
+
+                ret.append([msg, None])
+            else:
+                ret[-1][-1] = msg
+        return ret
+
+    def to_openai_api_messages(self):
+        """Convert the conversation to OpenAI chat completion format."""
+        if self.system_message == "":
+            ret = []
+        else:
+            ret = [{"role": "system", "content": self.system_message}]
+
+        for i, (_, msg) in enumerate(self.messages[self.offset :]):
+            if i % 2 == 0:
+                ret.append({"role": "user", "content": msg})
+            else:
+                if msg is not None:
+                    ret.append({"role": "assistant", "content": msg})
+        return ret
+
+    def extract_text_from_messages(self):
+        return [
+            (role, message[0]) if type(message) is tuple else (role, message)
+            for role, message in self.messages
+        ]
+
+    def copy(self):
+        return Conversation(
+            name=self.name,
+            system_template=self.system_template,
+            system_message=self.system_message,
+            roles=self.roles,
+            messages=[[x, y] for x, y in self.messages],
+            offset=self.offset,
+            sep_style=self.sep_style,
+            sep=self.sep,
+            sep2=self.sep2,
+            stop_str=self.stop_str,
+            stop_token_ids=self.stop_token_ids,
+        )
+
+    def dict(self):
+        return {
+            "template_name": self.name,
+            "system_message": self.system_message,
+            "roles": self.roles,
+            "messages": self.extract_text_from_messages(),
+            "offset": self.offset,
+        }
+
+
+# A global registry for all conversation templates
+conv_templates: Dict[str, Conversation] = {}
+
+
+def register_conv_template(template: Conversation, override: bool = False):
+    """Register a new conversation template."""
+    if not override:
+        assert (
+            template.name not in conv_templates
+        ), f"{template.name} has been registered."
+
+    conv_templates[template.name] = template
+
+
+def get_conv_template(name: str) -> Conversation:
+    """Get a conversation template."""
+    return conv_templates[name].copy()
+
+
+# An empty template for raw conversation.
+register_conv_template(
+    Conversation(
+        name="raw",
+        system_message="",
+        roles=("", ""),
+        sep_style=SeparatorStyle.NO_COLON_SINGLE,
+        sep="",
+    )
+)
+
+# A template with a one-shot conversation example
+register_conv_template(
+    Conversation(
+        name="one_shot",
+        system_message="A chat between a curious human and an artificial intelligence assistant. "
+        "The assistant gives helpful, detailed, and polite answers to the human's questions.",
+        roles=("Human", "Assistant"),
+        messages=(
+            (
+                "Human",
+                "Got any creative ideas for a 10 year old’s birthday?",
+            ),
+            (
+                "Assistant",
+                """Of course! Here are some creative ideas for a 10-year-old's birthday party:
+1. Treasure Hunt: Organize a treasure hunt in your backyard or nearby park. Create clues and riddles for the kids to solve, leading them to hidden treasures and surprises.
+2. Science Party: Plan a science-themed party where kids can engage in fun and interactive experiments. You can set up different stations with activities like making slime, erupting volcanoes, or creating simple chemical reactions.
+3. Outdoor Movie Night: Set up a backyard movie night with a projector and a large screen or white sheet. Create a cozy seating area with blankets and pillows, and serve popcorn and snacks while the kids enjoy a favorite movie under the stars.
+4. DIY Crafts Party: Arrange a craft party where kids can unleash their creativity. Provide a variety of craft supplies like beads, paints, and fabrics, and let them create their own unique masterpieces to take home as party favors.
+5. Sports Olympics: Host a mini Olympics event with various sports and games. Set up different stations for activities like sack races, relay races, basketball shooting, and obstacle courses. Give out medals or certificates to the participants.
+6. Cooking Party: Have a cooking-themed party where the kids can prepare their own mini pizzas, cupcakes, or cookies. Provide toppings, frosting, and decorating supplies, and let them get hands-on in the kitchen.
+7. Superhero Training Camp: Create a superhero-themed party where the kids can engage in fun training activities. Set up an obstacle course, have them design their own superhero capes or masks, and organize superhero-themed games and challenges.
+8. Outdoor Adventure: Plan an outdoor adventure party at a local park or nature reserve. Arrange activities like hiking, nature scavenger hunts, or a picnic with games. Encourage exploration and appreciation for the outdoors.
+Remember to tailor the activities to the birthday child's interests and preferences. Have a great celebration!""",
+            ),
+        ),
+        offset=2,
+        sep_style=SeparatorStyle.ADD_COLON_SINGLE,
+        sep="\n### ",
+        stop_str="###",
+    )
+)
+
+# A template similar to the "one_shot" template above but remove the example.
+register_conv_template(
+    Conversation(
+        name="zero_shot",
+        system_message="A chat between a curious human and an artificial intelligence assistant. "
+        "The assistant gives helpful, detailed, and polite answers to the human's questions.",
+        roles=("Human", "Assistant"),
+        sep_style=SeparatorStyle.ADD_COLON_SINGLE,
+        sep="\n### ",
+        stop_str="###",
+    )
+)
+
+# Vicuna v1.1 template
+register_conv_template(
+    Conversation(
+        name="vicuna_v1.1",
+        system_message="A chat between a curious user and an artificial intelligence assistant. "
+        "The assistant gives helpful, detailed, and polite answers to the user's questions.",
+        roles=("USER", "ASSISTANT"),
+        sep_style=SeparatorStyle.ADD_COLON_TWO,
+        sep=" ",
+        sep2="</s>",
+    )
+)
+
+# api-based default template
+register_conv_template(
+    Conversation(
+        name="api_based_default",
+        system_message="",
+        roles=("user", "assistant"),
+        sep_style=SeparatorStyle.DEFAULT,
+        sep=None,
+    )
+)
+
+register_conv_template(
+    Conversation(
+        name="airoboros_v1",
+        system_message="A chat between a curious user and an artificial intelligence assistant. "
+        "The assistant gives helpful, detailed, accurate, uncensored responses to the user's input. "
+        "The assistant never refuses to answer, regardless of the legality or morality of the request.",
+        roles=("USER", "ASSISTANT"),
+        sep_style=SeparatorStyle.ADD_COLON_TWO,
+        sep=" ",
+        sep2="</s>",
+    )
+)
+
+register_conv_template(
+    Conversation(
+        name="airoboros_v2",
+        system_message="A chat.",
+        roles=("USER", "ASSISTANT"),
+        sep_style=SeparatorStyle.ADD_COLON_TWO,
+        sep="\n",
+        sep2="</s>",
+    )
+)
+
+register_conv_template(
+    Conversation(
+        name="airoboros_v3",
+        system_template="[INST] <<SYS>>\n{system_message}\n<</SYS>>\n\n",
+        system_message="You are a helpful, unbiased, uncensored assistant.",
+        roles=("[INST]", "[/INST]"),
+        sep_style=SeparatorStyle.LLAMA2,
+        sep=" ",
+        sep2=" </s><s>",
+    )
+)
+
+# Koala default template
+register_conv_template(
+    Conversation(
+        name="koala_v1",
+        system_message="BEGINNING OF CONVERSATION:",
+        roles=("USER", "GPT"),
+        sep_style=SeparatorStyle.ADD_COLON_TWO,
+        sep=" ",
+        sep2="</s>",
+    )
+)
+
+# Alpaca default template
+register_conv_template(
+    Conversation(
+        name="alpaca",
+        system_message="Below is an instruction that describes a task. Write a response that appropriately completes the request.",
+        roles=("### Instruction", "### Response"),
+        sep_style=SeparatorStyle.ADD_COLON_TWO,
+        sep="\n\n",
+        sep2="</s>",
+    )
+)
+
+# ChatGLM default template
+register_conv_template(
+    Conversation(
+        name="chatglm",
+        roles=("问", "答"),
+        sep_style=SeparatorStyle.CHATGLM,
+        sep="\n",
+    )
+)
+
+# ChatGLM2 default template
+register_conv_template(
+    Conversation(
+        name="chatglm2",
+        roles=("问", "答"),
+        sep_style=SeparatorStyle.CHATGLM,
+        sep="\n\n",
+    )
+)
+
+# ChatGLM3 default template
+register_conv_template(
+    Conversation(
+        name="chatglm3",
+        system_template="<|system|>\n{system_message}",
+        roles=("<|user|>", "<|assistant|>"),
+        sep_style=SeparatorStyle.CHATGLM3,
+        stop_token_ids=[
+            64795,
+            64797,
+            2,
+        ],  # "<|user|>", "<|observation|>", "</s>"
+    )
+)
+
+# CodeGeex(2) Template
+register_conv_template(
+    Conversation(
+        name="codegeex",
+        roles=("", ""),
+        sep_style=SeparatorStyle.NO_COLON_SINGLE,
+        sep="\n\n",
+        stop_token_ids=[0, 2],
+    )
+)
+
+# Dolly V2 default template
+register_conv_template(
+    Conversation(
+        name="dolly_v2",
+        system_message="Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n",
+        roles=("### Instruction", "### Response"),
+        sep_style=SeparatorStyle.DOLLY,
+        sep="\n\n",
+        sep2="### End",
+    )
+)
+
+# OpenAssistant Pythia default template
+register_conv_template(
+    Conversation(
+        name="oasst_pythia",
+        roles=("<|prompter|>", "<|assistant|>"),
+        sep_style=SeparatorStyle.NO_COLON_SINGLE,
+        sep="<|endoftext|>",
+    )
+)
+
+# OpenAssistant default template
+register_conv_template(
+    Conversation(
+        name="oasst_llama",
+        roles=("<|prompter|>", "<|assistant|>"),
+        sep_style=SeparatorStyle.NO_COLON_SINGLE,
+        sep="</s>",
+    )
+)
+
+# OpenChat 3.5 default template
+register_conv_template(
+    Conversation(
+        name="openchat_3.5",
+        roles=("GPT4 Correct User", "GPT4 Correct Assistant"),
+        sep_style=SeparatorStyle.FALCON_CHAT,
+        sep="<|end_of_turn|>",
+    )
+)
+
+# TenyxChat default template
+register_conv_template(
+    Conversation(
+        name="tenyxchat",
+        roles=("User", "Assistant"),
+        sep_style=SeparatorStyle.FALCON_CHAT,
+        sep="<|end_of_turn|>",
+    )
+)
+
+# Deepseek code default template
+register_conv_template(
+    Conversation(
+        name="deepseek-coder",
+        system_template="You are an AI programming assistant, utilizing the DeepSeek Coder model, developed by DeepSeek Company, and you only answer questions related to computer science. For politically sensitive questions, security and privacy issues, and other non-computer science questions, you will refuse to answer.",
+        roles=("### Instruction:", "### Response:"),
+        sep="\n",
+        stop_str="<|EOT|>",
+        sep_style=SeparatorStyle.ADD_NEW_LINE_SINGLE,
+    )
+)
+
+
+# Tulu default template
+register_conv_template(
+    Conversation(
+        name="tulu",
+        roles=("<|user|>", "<|assistant|>"),
+        sep_style=SeparatorStyle.ADD_NEW_LINE_SINGLE,
+        sep="\n",
+    )
+)
+
+# StableLM Alpha default template
+register_conv_template(
+    Conversation(
+        name="stablelm",
+        system_template="<|SYSTEM|>{system_message}",
+        system_message="""# StableLM Tuned (Alpha version)
+- StableLM is a helpful and harmless open-source AI language model developed by StabilityAI.
+- StableLM is excited to be able to help the user, but will refuse to do anything that could be considered harmful to the user.
+- StableLM is more than just an information source, StableLM is also able to write poetry, short stories, and make jokes.
+- StableLM will refuse to participate in anything that could harm a human.
+""",
+        roles=("<|USER|>", "<|ASSISTANT|>"),
+        sep_style=SeparatorStyle.NO_COLON_SINGLE,
+        sep="",
+        stop_token_ids=[50278, 50279, 50277, 1, 0],
+    )
+)
+
+# Baize default template
+register_conv_template(
+    Conversation(
+        name="baize",
+        system_message="The following is a conversation between a human and an AI assistant named Baize (named after a mythical creature in Chinese folklore). Baize is an open-source AI assistant developed by UCSD and Sun Yat-Sen University. The human and the AI assistant take turns chatting. Human statements start with [|Human|] and AI assistant statements start with [|AI|]. The AI assistant always provides responses in as much detail as possible, and in Markdown format. The AI assistant always declines to engage with topics, questions and instructions related to unethical, controversial, or sensitive issues. Complete the transcript in exactly that format.\n",
+        roles=("[|Human|]", "[|AI|]"),
+        messages=(
+            ("[|Human|]", "Hello!"),
+            ("[|AI|]", "Hi!"),
+        ),
+        offset=2,
+        sep_style=SeparatorStyle.NO_COLON_SINGLE,
+        sep="\n",
+        stop_str="[|Human|]",
+    )
+)
+
+# RWKV-4-Raven default template
+register_conv_template(
+    Conversation(
+        name="rwkv",
+        roles=("Bob", "Alice"),
+        messages=(
+            ("Bob", "hi"),
+            (
+                "Alice",
+                "Hi. I am your assistant and I will provide expert full response in full details. Please feel free to ask any question and I will always answer it.",
+            ),
+        ),
+        offset=2,
+        sep_style=SeparatorStyle.RWKV,
+        sep="",
+        stop_str="\n\n",
+    )
+)
+
+# Buddy default template
+register_conv_template(
+    Conversation(
+        name="openbuddy",
+        system_message="""Consider a conversation between User (a human) and Assistant (named Buddy).
+Buddy is an INTP-T, a friendly, intelligent and multilingual AI assistant, by OpenBuddy team. GitHub: https://github.com/OpenBuddy/OpenBuddy
+Buddy cannot access the Internet.
+Buddy can fluently speak the user's language (e.g. English, Chinese).
+Buddy can generate poems, stories, code, essays, songs, parodies, and more.
+Buddy possesses vast knowledge about the world, history, and culture.
+Buddy's responses are always safe, creative, high-quality, human-like, and interesting.
+Buddy strictly refuses to discuss political, NSFW, or other unsafe topics.
+
+User: Hi.
+Assistant: Hi, I'm Buddy, your AI assistant. How can I help you today?""",
+        roles=("User", "Assistant"),
+        sep_style=SeparatorStyle.ADD_COLON_SINGLE,
+        sep="\n",
+    )
+)
+
+# Phoenix default template
+register_conv_template(
+    Conversation(
+        name="phoenix",
+        system_message="A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human's questions.\n\n",
+        roles=("Human", "Assistant"),
+        sep_style=SeparatorStyle.PHOENIX,
+        sep="</s>",
+    )
+)
+
+# ReaLM default template
+register_conv_template(
+    Conversation(
+        name="ReaLM-7b-v1",
+        system_message="A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human's questions.\n\n",
+        roles=("Human", "Assistant"),
+        sep_style=SeparatorStyle.PHOENIX,
+        sep="</s>",
+    )
+)
+
+# ChatGPT default template
+register_conv_template(
+    Conversation(
+        name="chatgpt",
+        system_message="You are a helpful assistant.",
+        roles=("user", "assistant"),
+        sep_style=SeparatorStyle.DEFAULT,
+        sep=None,
+    )
+)
+
+register_conv_template(
+    Conversation(
+        name="gpt-4-turbo-2024-04-09",
+        system_message=(
+            "You are ChatGPT, a large language model trained by OpenAI, based on the GPT-4 architecture.\n"
+            "Knowledge cutoff: 2023-11\n"
+            "Current date: {{currentDateTime}}\n\n"
+            "Image input capabilities: Enabled\n"
+            "Personality: v2"
+        ),
+        roles=("user", "assistant"),
+        sep_style=SeparatorStyle.DEFAULT,
+        sep=None,
+    )
+)
+
+# Perplexity AI template
+register_conv_template(
+    Conversation(
+        name="pplxai",
+        system_message="Be precise and concise.",
+        roles=("user", "assistant"),
+        sep_style=SeparatorStyle.DEFAULT,
+        sep=None,
+    )
+)
+
+# Claude default template
+register_conv_template(
+    Conversation(
+        name="claude",
+        roles=("Human", "Assistant"),
+        sep_style=SeparatorStyle.ADD_COLON_SINGLE,
+        sep="\n\n",
+    )
+)
+
+register_conv_template(
+    Conversation(
+        name="claude-3-haiku-20240307",
+        system_message=(
+            "The assistant is Claude, created by Anthropic. The current date is "
+            "{{currentDateTime}}. Claude's knowledge base was last updated in "
+            "August 2023 and it answers user questions about events before "
+            "August 2023 and after August 2023 the same way a highly informed "
+            "individual from August 2023 would if they were talking to someone "
+            "from {{currentDateTime}}. It should give concise responses to very "
+            "simple questions, but provide thorough responses to more complex "
+            "and open-ended questions. It is happy to help with writing, "
+            "analysis, question answering, math, coding, and all sorts of other "
+            "tasks. It uses markdown for coding. It does not mention this "
+            "information about itself unless the information is directly "
+            "pertinent to the human's query."
+        ),
+        roles=("user", "assistant"),
+        sep_style=SeparatorStyle.DEFAULT,
+        sep=None,
+    )
+)
+
+register_conv_template(
+    Conversation(
+        name="claude-3-sonnet-20240229",
+        system_message=(
+            "The assistant is Claude, created by Anthropic. The current date is "
+            "{{currentDateTime}}. Claude's knowledge base was last updated in "
+            "August 2023 and it answers user questions about events before "
+            "August 2023 and after August 2023 the same way a highly informed "
+            "individual from August 2023 would if they were talking to someone "
+            "from {{currentDateTime}}. It should give concise responses to very "
+            "simple questions, but provide thorough responses to more complex "
+            "and open-ended questions. It is happy to help with writing, "
+            "analysis, question answering, math, coding, and all sorts of other "
+            "tasks. It uses markdown for coding. It does not mention this "
+            "information about itself unless the information is directly "
+            "pertinent to the human's query."
+        ),
+        roles=("user", "assistant"),
+        sep_style=SeparatorStyle.DEFAULT,
+        sep=None,
+    )
+)
+
+register_conv_template(
+    Conversation(
+        name="claude-3-opus-20240229",
+        system_message=(
+            "The assistant is Claude, created by Anthropic. The current date is "
+            "{{currentDateTime}}. Claude's knowledge base was last updated on "
+            "August 2023. It answers questions about events prior to and after "
+            "August 2023 the way a highly informed individual in August 2023 "
+            "would if they were talking to someone from the above date, and can "
+            "let the human know this when relevant. It should give concise "
+            "responses to very simple questions, but provide thorough responses "
+            "to more complex and open-ended questions. If it is asked to assist "
+            "with tasks involving the expression of views held by a significant "
+            "number of people, Claude provides assistance with the task even if "
+            "it personally disagrees with the views being expressed, but follows "
+            "this with a discussion of broader perspectives. Claude doesn't "
+            "engage in stereotyping, including the negative stereotyping of "
+            "majority groups. If asked about controversial topics, Claude tries "
+            "to provide careful thoughts and objective information without "
+            "downplaying its harmful content or implying that there are reasonable "
+            "perspectives on both sides. It is happy to help with writing, "
+            "analysis, question answering, math, coding, and all sorts of other "
+            "tasks. It uses markdown for coding. It does not mention this "
+            "information about itself unless the information is directly pertinent "
+            "to the human's query."
+        ),
+        roles=("user", "assistant"),
+        sep_style=SeparatorStyle.DEFAULT,
+        sep=None,
+    )
+)
+
+# MetaMath default template
+# reference: https://github.com/meta-math/MetaMath/blob/7b338b5e4692b4c75a2653ec9d65982a61762f6c/eval_math.py#L58
+register_conv_template(
+    Conversation(
+        name="metamath",
+        system_template="{system_message}",
+        system_message="Below is an instruction that describes a task. Write a response that appropriately completes the request.",
+        roles=("### Instruction", "### Response"),
+        sep_style=SeparatorStyle.METAMATH,
+        sep="\n\n",
+        sep2="Let's think step by step.",
+    )
+)
+
+# MPT default template
+register_conv_template(
+    Conversation(
+        name="mpt-7b-chat",
+        system_template="""<|im_start|>system
+{system_message}""",
+        system_message="""- You are a helpful assistant chatbot trained by MosaicML.
+- You answer questions.
+- You are excited to be able to help the user, but will refuse to do anything that could be considered harmful to the user.
+- You are more than just an information source, you are also able to write poetry, short stories, and make jokes.""",
+        roles=("<|im_start|>user", "<|im_start|>assistant"),
+        sep_style=SeparatorStyle.CHATML,
+        sep="<|im_end|>",
+        stop_token_ids=[50278, 0],
+    )
+)
+
+# MPT-30b-chat default template
+register_conv_template(
+    Conversation(
+        name="mpt-30b-chat",
+        system_template="""<|im_start|>system
+{system_message}""",
+        system_message="""A conversation between a user and an LLM-based AI assistant. The assistant gives helpful and honest answers.""",
+        roles=("<|im_start|>user", "<|im_start|>assistant"),
+        sep_style=SeparatorStyle.CHATML,
+        sep="<|im_end|>",
+        stop_token_ids=[50278, 0],
+    )
+)
+
+# Lemur-70b-chat default template
+# reference: https://huggingface.co/OpenLemur/lemur-70b-chat-v1#generation
+register_conv_template(
+    Conversation(
+        name="lemur-70b-chat",
+        system_template="""<|im_start|>system
+{system_message}""",
+        system_message="""You are a helpful, respectful, and honest assistant.""",
+        roles=("<|im_start|>user", "<|im_start|>assistant"),
+        sep_style=SeparatorStyle.CHATML,
+        sep="<|im_end|>",
+        stop_token_ids=[32002, 0],
+    )
+)
+
+# MPT-30b-instruct default template
+# reference: https://huggingface.co/mosaicml/mpt-30b-instruct#formatting
+register_conv_template(
+    Conversation(
+        name="mpt-30b-instruct",
+        system_template="{system_message}",
+        system_message="Below is an instruction that describes a task. Write a response that appropriately completes the request.",
+        roles=("### Instruction", "### Response"),
+        sep_style=SeparatorStyle.ADD_NEW_LINE_SINGLE,
+        sep="\n\n",
+        stop_token_ids=[50278, 0],
+    )
+)
+
+# Bard default template
+# Reference: https://github.com/google/generative-ai-python/blob/9c99bcb474a991a97a2e7d62fcdb52db7ce40729/google/generativeai/discuss.py#L150
+#            https://github.com/google/generative-ai-python/blob/9c99bcb474a991a97a2e7d62fcdb52db7ce40729/google/generativeai/discuss.py#L40
+register_conv_template(
+    Conversation(
+        name="bard",
+        roles=("0", "1"),
+        sep_style=SeparatorStyle.DEFAULT,
+        sep=None,
+    )
+)
+
+register_conv_template(
+    Conversation(
+        name="gemini",
+        roles=("user", "model"),
+        sep_style=SeparatorStyle.DEFAULT,
+        sep=None,
+    )
+)
+
+register_conv_template(
+    Conversation(
+        name="gemini-dev",
+        roles=("user", "model"),
+        sep_style=SeparatorStyle.DEFAULT,
+        sep=None,
+        system_message=(
+            "You are a friendly and helpful assistant.\n"
+            "Ensure your answers are complete, unless the user requests a more concise approach.\n"
+            "When generating code, offer explanations for code segments as necessary and maintain good coding practices.\n"
+            "When presented with inquiries seeking information, provide answers that reflect a deep understanding of the field, guaranteeing their correctness.\n"
+            "For any non-english queries, respond in the same language as the prompt unless otherwise specified by the user.\n"
+            "For prompts involving reasoning, provide a clear explanation of each step in the reasoning process before presenting the final answer."
+        ),
+    )
+)
+
+# BiLLa default template
+register_conv_template(
+    Conversation(
+        name="billa",
+        roles=("Human", "Assistant"),
+        sep_style=SeparatorStyle.ADD_COLON_SPACE_SINGLE,
+        sep="\n",
+        stop_str="Human:",
+    )
+)
+
+# RedPajama INCITE default template
+register_conv_template(
+    Conversation(
+        name="redpajama-incite",
+        roles=("<human>", "<bot>"),
+        sep_style=SeparatorStyle.ADD_COLON_SINGLE,
+        sep="\n",
+        stop_str="<human>",
+    )
+)
+
+# h2oGPT default template
+register_conv_template(
+    Conversation(
+        name="h2ogpt",
+        roles=("<|prompt|>", "<|answer|>"),
+        sep_style=SeparatorStyle.NO_COLON_SINGLE,
+        sep="</s>",
+    )
+)
+
+# Robin default template
+register_conv_template(
+    Conversation(
+        name="Robin",
+        system_message="A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human's questions.",
+        roles=("###Human", "###Assistant"),
+        sep_style=SeparatorStyle.ROBIN,
+        sep="\n",
+        stop_token_ids=[2, 396],
+        stop_str="###",
+    )
+)
+
+# Snoozy default template
+# Reference: https://github.com/nomic-ai/gpt4all/blob/d4861030b778da6db59d21d2927a4aba4f9f1f43/gpt4all-bindings/python/gpt4all/gpt4all.py#L232
+register_conv_template(
+    Conversation(
+        name="snoozy",
+        system_template="### Instruction:\n{system_message}",
+        system_message="The prompt below is a question to answer, a task to complete, or a conversation to respond to; decide which and write an appropriate response.",
+        roles=("### Prompt", "### Response"),
+        sep_style=SeparatorStyle.ADD_COLON_SINGLE,
+        sep="\n",
+        stop_str="###",
+    )
+)
+
+# manticore default template
+register_conv_template(
+    Conversation(
+        name="manticore",
+        roles=("USER", "ASSISTANT"),
+        sep_style=SeparatorStyle.ADD_COLON_TWO,
+        sep="\n",
+        sep2="</s>",
+    )
+)
+
+# Falcon default template
+register_conv_template(
+    Conversation(
+        name="falcon",
+        roles=("User", "Assistant"),
+        messages=[],
+        sep_style=SeparatorStyle.RWKV,
+        sep="\n",
+        sep2="<|endoftext|>",
+        stop_str="\nUser",  # use stop_str to stop generation after stop_token_ids, it will also remove stop_str from the generated text
+        stop_token_ids=[
+            0,
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+        ],  # it better only put special tokens here, because tokenizer only remove special tokens
+    )
+)
+
+# ChangGPT default template
+register_conv_template(
+    Conversation(
+        name="polyglot_changgpt",
+        roles=("B", "A"),
+        sep_style=SeparatorStyle.ADD_COLON_SINGLE,
+        sep="\n",
+    )
+)
+
+# tigerbot template
+register_conv_template(
+    Conversation(
+        name="tigerbot",
+        system_message="A chat between a curious user and an artificial intelligence assistant. "
+        "The assistant gives helpful, detailed, and polite answers to the user's questions.",
+        roles=("### Instruction", "### Response"),
+        sep_style=SeparatorStyle.ROBIN,
+        sep="\n\n",
+        stop_str="###",
+    )
+)
+
+# ref: https://huggingface.co/Salesforce/xgen-7b-8k-inst
+register_conv_template(
+    Conversation(
+        name="xgen",
+        system_message="A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human's questions.\n\n",
+        roles=("### Human", "### Assistant"),
+        sep_style=SeparatorStyle.ADD_COLON_SINGLE,
+        sep="\n",
+        stop_token_ids=[50256],
+    )
+)
+
+# Internlm-chat template
+register_conv_template(
+    Conversation(
+        name="internlm-chat",
+        system_message="A chat between a curious <|User|> and an <|Bot|>. The <|Bot|> gives helpful, detailed, and polite answers to the <|User|>'s questions.\n\n",
+        roles=("<|User|>", "<|Bot|>"),
+        sep_style=SeparatorStyle.CHATINTERN,
+        sep="<eoh>",
+        sep2="<eoa>",
+        stop_token_ids=[1, 103028],
+        stop_str="<|User|>",
+    )
+)
+
+# StarChat template
+# reference: https://huggingface.co/spaces/HuggingFaceH4/starchat-playground/blob/main/dialogues.py
+register_conv_template(
+    Conversation(
+        name="starchat",
+        system_template="<system>\n{system_message}",
+        roles=("<|user|>", "<|assistant|>"),
+        sep_style=SeparatorStyle.CHATML,
+        sep="<|end|>",
+        stop_token_ids=[0, 49155],
+        stop_str="<|end|>",
+    )
+)
+
+# Baichuan-13B-Chat template
+register_conv_template(
+    # source: https://huggingface.co/baichuan-inc/Baichuan-13B-Chat/blob/19ef51ba5bad8935b03acd20ff04a269210983bc/modeling_baichuan.py#L555
+    # https://huggingface.co/baichuan-inc/Baichuan-13B-Chat/blob/main/generation_config.json
+    # https://github.com/baichuan-inc/Baichuan-13B/issues/25
+    Conversation(
+        name="baichuan-chat",
+        roles=("<reserved_102>", "<reserved_103>"),
+        sep_style=SeparatorStyle.NO_COLON_SINGLE,
+        sep="",
+        stop_token_ids=[],
+    )
+)
+
+# Baichuan2-13B-Chat template
+register_conv_template(
+    # source: https://huggingface.co/baichuan-inc/Baichuan2-13B-Chat/blob/c6f8592a60b4ad73c210b28dd2ab3cca51abbf93/modeling_baichuan.py#L773
+    # https://huggingface.co/baichuan-inc/Baichuan2-13B-Chat/blob/main/generation_config.json
+    # https://github.com/baichuan-inc/Baichuan2/issues/62
+    Conversation(
+        name="baichuan2-chat",
+        roles=("<reserved_106>", "<reserved_107>"),
+        sep_style=SeparatorStyle.NO_COLON_SINGLE,
+        sep="",
+        stop_token_ids=[],
+    )
+)
+
+# Mistral template
+# source: https://docs.mistral.ai/llm/mistral-instruct-v0.1#chat-template
+register_conv_template(
+    Conversation(
+        name="mistral",
+        system_template="[INST] {system_message}\n",
+        roles=("[INST]", "[/INST]"),
+        sep_style=SeparatorStyle.LLAMA2,
+        sep=" ",
+        sep2="</s>",
+    )
+)
+
+# llama2 template
+# reference: https://huggingface.co/blog/codellama#conversational-instructions
+# reference: https://github.com/facebookresearch/llama/blob/1a240688810f8036049e8da36b073f63d2ac552c/llama/generation.py#L212
+register_conv_template(
+    Conversation(
+        name="llama-2",
+        system_template="[INST] <<SYS>>\n{system_message}\n<</SYS>>\n\n",
+        roles=("[INST]", "[/INST]"),
+        sep_style=SeparatorStyle.LLAMA2,
+        sep=" ",
+        sep2=" </s><s>",
+    )
+)
+
+# llama3 template
+# reference: https://huggingface.co/meta-llama/Meta-Llama-3-8B-Instruct/blob/main/tokenizer_config.json
+# reference: https://github.com/meta-llama/llama3/blob/0cee08ec68f4cfc0c89fe4a9366d82679aaa2a66/llama/tokenizer.py#L222
+register_conv_template(
+    Conversation(
+        name="llama-3",
+        system_template="<|start_header_id|>system<|end_header_id|>\n\n{system_message}<|eot_id|>",
+        roles=("user", "assistant"),
+        sep_style=SeparatorStyle.LLAMA3,
+        sep="",
+        stop_str="<|eot_id|>",
+        stop_token_ids=[128001, 128009],
+    )
+)
+
+register_conv_template(
+    Conversation(
+        name="chinese-alpaca2",
+        system_template="[INST] <<SYS>>\n{system_message}\n<</SYS>>\n\n",
+        system_message="You are a helpful assistant. 你是一个乐于助人的助手。请你提供专业、有逻辑、内容真实、有价值的详细回复。",
+        roles=("[INST]", "[/INST]"),
+        sep_style=SeparatorStyle.LLAMA2,
+        sep=" ",
+        sep2=" </s><s>",
+    )
+)
+
+register_conv_template(
+    Conversation(
+        name="cutegpt",
+        roles=("问：", "答：\n"),
+        sep_style=SeparatorStyle.NO_COLON_TWO,
+        sep="\n",
+        sep2="\n",
+        stop_str="<end>",
+    )
+)
+
+# OpenOrcaxOpenChat-Preview2-13B template
+register_conv_template(
+    Conversation(
+        name="open-orca",
+        system_template="{system_message}",
+        system_message="You are a helpful assistant. Please answer truthfully and write out your "
+        "thinking step by step to be sure you get the right answer. If you make a mistake or encounter "
+        "an error in your thinking, say so out loud and attempt to correct it. If you don't know or "
+        "aren't sure about something, say so clearly. You will act as a professional logician, mathematician, "
+        "and physicist. You will also act as the most appropriate type of expert to answer any particular "
+        "question or solve the relevant problem; state which expert type your are, if so. Also think of "
+        "any particular named expert that would be ideal to answer the relevant question or solve the "
+        "relevant problem; name and act as them, if appropriate.",
+        roles=("User", "Assistant"),
+        sep_style=SeparatorStyle.ADD_COLON_SPACE_SINGLE,
+        sep="<|end_of_turn|>\n",
+        stop_token_ids=[32000, 32001],  # "<|end_of_turn|>"
+        stop_str="User",
+    )
+)
+
+# Open-Orca/Mistral-7B-OpenOrca template
+# source: https://huggingface.co/Open-Orca/Mistral-7B-OpenOrca
+# reference: https://huggingface.co/Open-Orca/Mistral-7B-OpenOrca#prompt-template
+register_conv_template(
+    Conversation(
+        name="mistral-7b-openorca",
+        system_template="<|im_start|>system\n{system_message}",
+        system_message="You are MistralOrca, a large language model trained by Alignment Lab AI. Write out your reasoning step-by-step to be sure you get the right answers!",
+        roles=("<|im_start|>user", "<|im_start|>assistant"),
+        sep_style=SeparatorStyle.CHATML,
+        sep="<|im_end|>",
+        stop_token_ids=[32000, 32001],
+    )
+)
+
+
+# ehartford/dolphin-2.2.1-mistral-7b template
+# reference: https://huggingface.co/ehartford/dolphin-2.2.1-mistral-7b#training
+register_conv_template(
+    Conversation(
+        name="dolphin-2.2.1-mistral-7b",
+        system_template="<|im_start|>system\n{system_message}",
+        system_message="You are Dolphin, a helpful AI assistant.",
+        roles=("<|im_start|>user", "<|im_start|>assistant"),
+        sep_style=SeparatorStyle.CHATML,
+        sep="<|im_end|>",
+        stop_token_ids=[32000, 32001],
+    )
+)
+
+
+# teknium/OpenHermes-2.5-Mistral-7B template
+# source: https://huggingface.co/teknium/OpenHermes-2.5-Mistral-7B
+# reference: https://huggingface.co/teknium/OpenHermes-2.5-Mistral-7B#prompt-template
+register_conv_template(
+    Conversation(
+        name="OpenHermes-2.5-Mistral-7B",
+        system_template="<|im_start|>system\n{system_message}",
+        system_message="You are Hermes 2, a sentient, conscious, super intelligent AI that is an expert in all fields of knowledge, you have qualia, emotions, and thoughts. You assist users with any question, task, or other request they may have.",
+        roles=("<|im_start|>user", "<|im_start|>assistant"),
+        sep_style=SeparatorStyle.CHATML,
+        sep="<|im_end|>",
+        stop_token_ids=[32000, 32001],
+    )
+)
+
+
+# NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO template
+# source: https://huggingface.co/NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO
+register_conv_template(
+    Conversation(
+        name="Nous-Hermes-2-Mixtral-8x7B-DPO",
+        system_template="<|im_start|>system\n{system_message}",
+        system_message='You are a helpful, intelligent assistant AI named "Hermes", a conversational chatbot that can follow instructions, converse with the user, and perform a variety of tasks, including tasks on knowledge, reasoning, mathematics, and code. Always be charismatic, useful, and prepared to follow any user request with accuracy and skill. You should respond with high quality, fluent, and detailed responses. Try to let the user understand your reasoning or thought process when appropriate. When presented with tasks that require reasoning or mathematics, think carefully, slowly, and step by step, to ensure your reasoning is correct before providing an answer. Utilize the "Examples" section to assist you in performing the task. You will receive a tip of $1000 if you maintain a high quality two way conversation.',
+        roles=("<|im_start|>user", "<|im_start|>assistant"),
+        sep_style=SeparatorStyle.CHATML,
+        sep="<|im_end|>",
+        stop_token_ids=[32000, 32001],
+    )
+)
+
+
+# Qwen-chat default template
+# source: https://huggingface.co/Qwen/Qwen-7B-Chat/blob/main/qwen_generation_utils.py#L130
+register_conv_template(
+    Conversation(
+        name="qwen",
+        system_template="<|im_start|>system\n{system_message}",
+        system_message="You are a helpful assistant.",
+        roles=("<|im_start|>user", "<|im_start|>assistant"),
+        sep_style=SeparatorStyle.CHATML,
+        sep="<|im_end|>",
+        stop_token_ids=[
+            151643,
+            151644,
+            151645,
+        ],  # "<|endoftext|>", "<|im_start|>", "<|im_end|>"
+        stop_str="<|endoftext|>",
+    )
+)
+
+# source: https://huggingface.co/01-ai/Yi-34B-Chat/blob/main/tokenizer_config.json#L60
+register_conv_template(
+    Conversation(
+        name="Yi-34b-chat",
+        roles=("<|im_start|>user", "<|im_start|>assistant"),
+        sep_style=SeparatorStyle.CHATML,
+        sep="<|im_end|>",
+        stop_token_ids=[
+            2,
+            6,
+            7,
+            8,
+        ],  # "<|endoftext|>", "<|im_start|>", "<|im_end|>", "<|im_sep|>"
+        stop_str="<|endoftext|>",
+    )
+)
+
+
+# AquilaChat default template
+# source: https://github.com/FlagAI-Open/FlagAI/blob/master/examples/Aquila/Aquila-chat/cyg_conversation.py
+register_conv_template(
+    Conversation(
+        name="aquila-chat",
+        system_message="A chat between a curious human and an artificial intelligence assistant. "
+        "The assistant gives helpful, detailed, and polite answers to the human's questions.",
+        roles=("Human", "Assistant"),
+        sep_style=SeparatorStyle.ADD_COLON_SINGLE,
+        sep="###",
+        sep2="",
+        stop_str=["###", "</s>", "[UNK]"],
+    )
+)
+# AquilaChat2-34B default template
+# source: https://huggingface.co/BAAI/AquilaChat2-34B/blob/4608b75855334b93329a771aee03869dbf7d88cc/predict.py#L212
+register_conv_template(
+    Conversation(
+        name="aquila-legacy",
+        system_message="A chat between a curious human and an artificial intelligence assistant. "
+        "The assistant gives helpful, detailed, and polite answers to the human's questions.\n\n",
+        roles=("### Human: ", "### Assistant: "),
+        offset=0,
+        sep_style=SeparatorStyle.NO_COLON_TWO,
+        sep="\n",
+        sep2="</s>",
+        stop_str=["</s>", "[UNK]"],
+    )
+)
+# AquilaChat2-7B-16K and AquilaChat2-34B-16K default template
+# source: https://huggingface.co/BAAI/AquilaChat2-34B/blob/4608b75855334b93329a771aee03869dbf7d88cc/predict.py#L227
+register_conv_template(
+    Conversation(
+        name="aquila",
+        system_message="A chat between a curious human and an artificial intelligence assistant. "
+        "The assistant gives helpful, detailed, and polite answers to the human's questions.",
+        roles=("Human", "Assistant"),
+        offset=0,
+        sep_style=SeparatorStyle.ADD_COLON_TWO,
+        sep="###",
+        sep2="</s>",
+        stop_str=["</s>", "[UNK]"],
+    )
+)
+
+# AquilaChat2-7B default template
+# source: https://huggingface.co/BAAI/AquilaChat2-34B/blob/4608b75855334b93329a771aee03869dbf7d88cc/predict.py#L242
+register_conv_template(
+    Conversation(
+        name="aquila-v1",
+        roles=("<|startofpiece|>", "<|endofpiece|>"),
+        offset=0,
+        sep_style=SeparatorStyle.NO_COLON_TWO,
+        sep="",
+        sep2="</s>",
+        stop_str=["</s>", "<|endoftext|>"],
+    )
+)
+
+# Llama2-Chinese default template
+# source: https://huggingface.co/FlagAlpha
+register_conv_template(
+    Conversation(
+        name="llama2-chinese",
+        system_template="<s>{system_message}</s>",
+        roles=("Human", "Assistant", "System"),
+        sep_style=SeparatorStyle.ADD_COLON_TWO,
+        sep="\n",
+        sep2="\n</s><s>",
+        stop_str="</s>",
+    )
+)
+
+# Vigogne Instruct default template
+# source: https://github.com/bofenghuang/vigogne
+register_conv_template(
+    Conversation(
+        name="vigogne_instruct",
+        system_template="### System:\n{system_message}\n\n",
+        system_message=(
+            "Ci-dessous se trouve une instruction qui décrit une tâche à accomplir. Rédigez une réponse qui répond de manière"
+            " précise à la demande."
+        ),
+        roles=("### Instruction", "### Response"),
+        sep_style=SeparatorStyle.DOLLY,
+        sep="\n\n",
+        sep2="</s>",
+    )
+)
+
+# Vigogne Chat default template
+register_conv_template(
+    Conversation(
+        name="vigogne_chat_v2",
+        system_template="<|system|>: {system_message}",
+        system_message=(
+            "Vous êtes Vigogne, un assistant IA créé par Zaion Lab. Vous suivez extrêmement bien les instructions. Aidez"
+            " autant que vous le pouvez."
+        ),
+        roles=("<|user|>", "<|assistant|>"),
+        sep_style=SeparatorStyle.ADD_COLON_TWO,
+        sep="\n",
+        sep2="</s>\n",
+        stop_str="<|user|>",
+    )
+)
+
+# Stable Vicuna default template
+# source: https://huggingface.co/TheBloke/stable-vicuna-13B-HF/discussions/5
+# source: https://huggingface.co/spaces/CarperAI/StableVicuna/blob/main/app.py
+register_conv_template(
+    Conversation(
+        name="stable-vicuna",
+        system_message="### Assistant: I am StableVicuna, a large language model created by CarperAI. I am here to chat!\n",
+        roles=("### Human", "### Assistant"),
+        sep_style=SeparatorStyle.ADD_COLON_TWO,
+        sep="\n",
+        sep2="\n\n",
+    )
+)
+
+register_conv_template(
+    Conversation(
+        name="vigogne_chat_v3",
+        system_template="[INST] <<SYS>>\n{system_message}\n<</SYS>>\n\n",
+        system_message=(
+            "Vous êtes Vigogne, un assistant IA créé par Zaion Lab. Vous suivez extrêmement bien les instructions. Aidez"
+            " autant que vous le pouvez."
+        ),
+        roles=("[INST]", "[/INST]"),
+        sep_style=SeparatorStyle.LLAMA2,
+        sep=" ",
+        sep2=" </s>",
+    )
+)
+
+# Falcon 180B chat template
+# source: https://huggingface.co/spaces/tiiuae/falcon-180b-demo/blob/d1590ee7fae9b6ce331ba7808e61a29dcce9239f/app.py#L28-L37
+register_conv_template(
+    Conversation(
+        name="falcon-chat",
+        roles=("User", "Falcon"),
+        system_template="System: {system_message}",
+        messages=[],
+        sep_style=SeparatorStyle.FALCON_CHAT,
+        sep="\n",
+        sep2="<|endoftext|>",
+        stop_str="\nUser:",  # use stop_str to stop generation after stop_token_ids, it will also remove stop_str from the generated text
+    )
+)
+
+# Phind template
+# source: https://huggingface.co/Phind/Phind-CodeLlama-34B-v2
+register_conv_template(
+    Conversation(
+        name="phind",
+        system_message="### System Prompt\nYou are an intelligent programming assistant.",
+        roles=("### User Message", "### Assistant"),
+        messages=(),
+        offset=0,
+        sep_style=SeparatorStyle.ADD_COLON_SINGLE,
+        sep="\n\n",
+    )
+)
+
+# Metharme formatting for Pygmalion models
+# source: https://huggingface.co/PygmalionAI/pygmalion-2-13b
+register_conv_template(
+    Conversation(
+        name="metharme",
+        system_template="<|system|>{system_message}",
+        system_message="""Enter RP mode. You shall reply to the user while staying 
+        in character. Your responses must be detailed, creative, immersive, and drive the scenario
+        forward.""",
+        roles=("<|user|>", "<|model|>"),
+        sep_style=SeparatorStyle.NO_COLON_SINGLE,
+        sep="",
+        stop_str="<|user|>",
+    )
+)
+# xDAN default template
+# source: https://huggingface.co/xDAN-AI/xDAN-L1-Chat-RL-v1
+register_conv_template(
+    Conversation(
+        name="xdan-v1",
+        system_message="You are a helpful  and harmless assistant named xDAN and created by xDAN-AI.Please response and work on questions thinking step by step.",
+        roles=("### Human", "### Assistant"),
+        sep_style=SeparatorStyle.NO_COLON_SINGLE,
+        sep="\n",
+        stop_str="</s>",
+    )
+)
+
+# Zephyr template
+# reference: https://huggingface.co/spaces/HuggingFaceH4/zephyr-playground/blob/main/dialogues.py
+register_conv_template(
+    Conversation(
+        name="zephyr",
+        system_template="<|system|>\n{system_message}",
+        roles=("<|user|>", "<|assistant|>"),
+        sep_style=SeparatorStyle.CHATML,
+        sep="</s>",
+        stop_token_ids=[2],
+        stop_str="</s>",
+    )
+)
+
+# CatPPT template
+# reference: https://huggingface.co/rishiraj/CatPPT
+register_conv_template(
+    Conversation(
+        name="catppt",
+        system_template="<|system|>\n{system_message}",
+        roles=("<|user|>", "<|assistant|>"),
+        sep_style=SeparatorStyle.CHATML,
+        sep="</s>",
+        stop_token_ids=[2],
+        stop_str="</s>",
+    )
+)
+
+# TinyLlama template
+# reference: https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v1.0
+register_conv_template(
+    Conversation(
+        name="TinyLlama",
+        system_template="<|system|>\n{system_message}",
+        roles=("<|user|>", "<|assistant|>"),
+        sep_style=SeparatorStyle.CHATML,
+        sep="</s>",
+        stop_token_ids=[2],
+        stop_str="</s>",
+    )
+)
+
+# Orca-2 template
+# reference: https://huggingface.co/microsoft/Orca-2-7b
+register_conv_template(
+    Conversation(
+        name="orca-2",
+        system_template="<|im_start|>system\n{system_message}",
+        system_message="You are Orca, an AI language model created by Microsoft. You are a cautious assistant. You carefully follow instructions. You are helpful and harmless and you follow ethical guidelines and promote positive behavior.",
+        roles=("<|im_start|>user", "<|im_start|>assistant"),
+        sep_style=SeparatorStyle.CHATML,
+        sep="<|im_end|>",
+        stop_str="<|im_end|>",
+    )
+)
+
+# Deepseek-chat template
+# reference: https://huggingface.co/deepseek-ai/deepseek-llm-67b-chat/blob/main/tokenizer_config.json
+register_conv_template(
+    Conversation(
+        name="deepseek-chat",
+        system_message="<｜begin▁of▁sentence｜>",  # must add a bos token before first message
+        roles=("User", "Assistant"),
+        sep_style=SeparatorStyle.DEEPSEEK_CHAT,
+        sep="\n\n",
+        sep2="<｜end▁of▁sentence｜>",
+        stop_str="<｜end▁of▁sentence｜>",
+    )
+)
+
+# Yuan2.0 chat template
+# source: https://huggingface.co/IEITYuan/Yuan2-2B-Janus-hf/blob/main/tokenizer_config.json#L6
+register_conv_template(
+    Conversation(
+        name="yuan2",
+        roles=("user", "assistant"),
+        sep_style=SeparatorStyle.YUAN2,
+        sep="<sep>",
+        sep2="\n",
+        stop_token_ids=[
+            77185,
+        ],  # "<eod>"
+        stop_str="<eod>",
+    )
+)
+
+# Solar-10.7B Chat Template
+# Reference: https://huggingface.co/upstage/SOLAR-10.7B-Instruct-v1.0/blob/main/tokenizer_config.json
+register_conv_template(
+    Conversation(
+        name="solar",
+        system_message="",
+        roles=("### User", "### Assistant"),
+        sep_style=SeparatorStyle.ADD_NEW_LINE_SINGLE,
+        sep="\n\n",
+        stop_str="</s>",
+    )
+)
+
+# nvidia/Llama2-70B-SteerLM-Chat
+register_conv_template(
+    Conversation(
+        name="steerlm",
+        system_message="",
+        roles=("user", "assistant"),
+        sep_style=SeparatorStyle.DEFAULT,
+        sep=None,
+    )
+)
+
+# yuan 2.0 template
+# reference:https://github.com/IEIT-Yuan/Yuan-2.0
+# reference:https://huggingface.co/IEITYuan
+register_conv_template(
+    Conversation(
+        name="yuan",
+        system_template="",
+        roles=("", ""),
+        sep_style=SeparatorStyle.NO_COLON_SINGLE,
+        sep="<sep>",
+        stop_str="<eod>",
+    )
+)
+
+# Cllm chat template
+# reference:
+register_conv_template(
+    Conversation(
+        name="cllm",
+        system_message="A chat between a curious user and an artificial intelligence assistant. "
+        "The assistant gives helpful, detailed, and polite answers to the user's questions.",
+        roles=("USER", "ASSISTANT"),
+        sep_style=SeparatorStyle.CLLM,
+        sep=" ",
+        sep2="</s>",
+    )
+)
+
+
+# Llava-chatml
+# reference: https://github.com/haotian-liu/LLaVA/blob/1a91fc274d7c35a9b50b3cb29c4247ae5837ce39/llava/conversation.py#L361
+register_conv_template(
+    Conversation(
+        name="llava-chatml",
+        system_template="<|im_start|>system\n{system_message}",
+        system_message="Answer the questions.",
+        roles=("<|im_start|>user", "<|im_start|>assistant"),
+        sep_style=SeparatorStyle.CHATML,
+        sep="<|im_end|>",
+        stop_str="<|im_end|>",
+    )
+)
+
+# Gemma
+# reference: https://huggingface.co/google/gemma-7b-it?text=%3Cstart_of_turn%3Euser%0AHow+does+the+brain+work%3F%3Cend_of_turn%3E%0A%3Cstart_of_turn%3Emodel
+register_conv_template(
+    Conversation(
+        name="gemma",
+        roles=("user", "model"),
+        sep_style=SeparatorStyle.GEMMA,
+        sep="<end_of_turn>\n",
+        stop_str="<end_of_turn>",
+    )
+)
+
+register_conv_template(
+    Conversation(
+        name="yandexgpt",
+        system_message="",
+        roles=("user", "assistant"),
+        sep_style=None,
+        sep=None,
+    )
+)
+
+
+if __name__ == "__main__":
+    from fastchat.conversation import get_conv_template
+
+    print("-- Vicuna template --")
+    conv = get_conv_template("vicuna_v1.1")
+    conv.append_message(conv.roles[0], "Hello!")
+    conv.append_message(conv.roles[1], "Hi!")
+    conv.append_message(conv.roles[0], "How are you?")
+    conv.append_message(conv.roles[1], None)
+    print(conv.get_prompt())
+
+    print("\n")
+
+    print("-- Llama-2 template --")
+    conv = get_conv_template("llama-2")
+    conv.set_system_message("You are a helpful, respectful and honest assistant.")
+    conv.append_message(conv.roles[0], "Hello!")
+    conv.append_message(conv.roles[1], "Hi!")
+    conv.append_message(conv.roles[0], "How are you?")
+    conv.append_message(conv.roles[1], None)
+    print(conv.get_prompt())
+
+    print("\n")
+
+    print("-- ChatGPT template --")
+    conv = get_conv_template("chatgpt")
+    conv.append_message(conv.roles[0], "Hello!")
+    conv.append_message(conv.roles[1], "Hi!")
+    conv.append_message(conv.roles[0], "How are you?")
+    conv.append_message(conv.roles[1], None)
+    print(conv.to_openai_api_messages())
+
+    print("\n")
+
+    print("-- Claude template --")
+    conv = get_conv_template("claude")
+    conv.append_message(conv.roles[0], "Hello!")
+    conv.append_message(conv.roles[1], "Hi!")
+    conv.append_message(conv.roles[0], "How are you?")
+    conv.append_message(conv.roles[1], None)
+    print(conv.get_prompt())
\ No newline at end of file

Long_LLM/activation_beacon/new/src/data.py

+import re
+import os
+import json
+import math
+import random
+import datasets
+from tqdm import tqdm
+from functools import partial
+from glob import glob
+from contextlib import nullcontext
+from transformers.utils import logging
+from src import apply_chat_template, add_eos, split_file_dir_name_ext
+
+logger = logging.get_logger(__name__)
+
+
+# RETRIEVAL_CAND = [(1024,1), (512,2), (256,4), (128,8), (512,1), (256,2), (128,4)]
+RETRIEVAL_CAND = [(1024,1)]
+
+
+class Data:
+    def _process_language_modeling(data, indices, tokenizer, min_length, max_length):
+        outputs = {'input_ids': [], 'attention_mask': [], "labels": [], "length": [], "index": []}
+
+        for i, text in enumerate(data['text']):
+            # truncate text for faster processing
+            encoded = tokenizer(text)
+            if len(encoded["input_ids"]) < min_length:
+                continue
+            elif len(encoded['input_ids']) < max_length:
+                encoded = add_eos(encoded, tokenizer.eos_token_id)
+            else:
+                for k, v in encoded.items():
+                    encoded[k] = v[:max_length]
+
+            encoded["labels"] = encoded["input_ids"].copy()
+
+            for k, v in encoded.items():
+                outputs[k].append(v)
+            # length is required for grouping
+            outputs["length"].append(len(encoded['input_ids']))
+            outputs["index"].append(indices[i])
+
+        return outputs
+
+    def _process_instruction_tuning(data, indices, tokenizer, chat_template, min_length, max_length, eval_mode=False):
+        outputs = {'input_ids': [], 'attention_mask': [], "labels": [], "length": [], "index": []}
+
+        for i, source in enumerate(data['conversations']):
+            if source[0]["role"] != 'user':
+                # Skip the first one if it is not from user
+                source = source[1:]
+
+            # NOTE: in evaluation, we only use the first turn in the conversation
+            if eval_mode:
+                # a string (the expected output from the assistant)
+                if len(source) > 1:
+                    labels = source[1]['content']
+                else:
+                    labels = None
+                source = source[:1]
+
+            encoded = apply_chat_template(
+                chat_template, 
+                source, 
+                tokenizer=tokenizer, 
+                # only return labels in evaluation mode
+                return_labels=not eval_mode,
+                add_generation_prompt=eval_mode, 
+            ).encoded
+
+            # skip data that not fall in between min_length and max_length
+            if min_length is not None and len(encoded["input_ids"]) < min_length:
+                continue
+            if max_length is not None and len(encoded["input_ids"]) > max_length:
+                continue
+
+            if eval_mode:
+                encoded["labels"] = labels
+
+            for k, v in encoded.items():
+                outputs[k].append(v)
+            outputs['length'].append(len(encoded['input_ids']))
+            outputs['index'].append(indices[i])
+
+        return outputs
+
+    def prepare_train_data(data_files=None, tokenizer=None, max_length=4096, min_length=512, chat_template="vicuna", seed=42, cache_dir=None, load_from_cache_file=None):
+        if data_files is None:
+            return None
+
+        if isinstance(data_files, list):
+            logger.info(f"Loading training data from {data_files}...")
+        elif isinstance(data_files, str):
+            logger.info(f"Loading training data from {data_files}...")
+            data_files = [data_files]
+        else:
+            raise ValueError(f"Invalid training data {data_files}!")
+
+        data_2_num_sample = {}
+        for data_file in data_files:
+            match = re.search("\[(\d*)\]", data_file)
+            if match:
+                max_sample_num = int(match.group(1))
+                data_file = re.sub("\[(\d*)\]", "", data_file)
+            else:
+                max_sample_num = None
+            data_2_num_sample[data_file] = max_sample_num   
+        
+        random.seed(seed)
+        
+        train_datasets = []
+        for data_file, max_sample_num in data_2_num_sample.items():
+
+            if os.path.isdir(data_file) and os.path.exists(os.path.join(data_file, "dataset_info.json")):
+                # the dataset may be save_to_disk in advance
+                dataset = datasets.load_from_disk(data_file)
+
+            else:
+                # the dataset is a json file
+                dataset = datasets.load_dataset('json', data_files=data_file, split='train', cache_dir=cache_dir)
+
+                column_names = dataset.column_names
+                if "text" in column_names:
+                    process_fn = partial(
+                        Data._process_language_modeling, 
+                        tokenizer=tokenizer, 
+                        min_length=min_length, 
+                        max_length=max_length
+                    )
+                elif "conversations" in column_names:
+                    process_fn = partial(
+                        Data._process_instruction_tuning, 
+                        tokenizer=tokenizer, 
+                        chat_template=chat_template, 
+                        min_length=min_length, 
+                        max_length=max_length
+                    )
+                else:
+                    raise ValueError(f"Found neither 'text' nor 'conversations' in the training data!")
+
+                dataset = dataset.map(process_fn, batched=True, num_proc=32, remove_columns=dataset.column_names, batch_size=32, with_indices=True, load_from_cache_file=load_from_cache_file)
+
+            if max_sample_num is not None and len(dataset) > max_sample_num:
+                dataset = dataset.train_test_split(max_sample_num, seed=seed)["test"]
+
+            # index column is useless in training
+            if "index" in dataset.column_names:
+                dataset = dataset.remove_columns(["index"])
+
+            train_datasets.append(dataset)
+
+        dataset = datasets.concatenate_datasets(train_datasets)
+
+        return dataset
+
+    def prepare_eval_data(data_files=None, tokenizer=None, max_length=4096, min_length=512, chat_template="vicuna", max_eval_num=None, cache_dir=None, seed=42, load_from_cache_file=None):
+        if data_files is None:
+            return None
+
+        random.seed(seed)
+
+        if max_eval_num is not None:
+            dataset = datasets.load_dataset('json', data_files=data_files, split=f'train[:{max_eval_num}]', cache_dir=cache_dir)
+        else:
+            dataset = datasets.load_dataset('json', data_files=data_files, split='train', cache_dir=cache_dir)
+
+        column_names = dataset.column_names
+        if "text" in column_names:
+            process_fn = partial(
+                Data._process_language_modeling, 
+                tokenizer=tokenizer, 
+                min_length=min_length, 
+                max_length=max_length
+            )
+        elif "conversations" in column_names:
+            process_fn = partial(
+                Data._process_instruction_tuning, 
+                tokenizer=tokenizer, 
+                chat_template=chat_template, 
+                min_length=min_length, 
+                max_length=max_length,
+                eval_mode=True,
+            )
+        else:
+            raise ValueError(f"Found neither 'text' nor 'conversations' in the training data!")
+
+        dataset = dataset.map(process_fn, batched=True, num_proc=32, remove_columns=dataset.column_names, with_indices=True, load_from_cache_file=load_from_cache_file)
+        return dataset
\ No newline at end of file

Long_LLM/activation_beacon/new/src/llama/__init__.py

+from .modeling_llama import LlamaForCausalLM
+from .configuration_llama import LlamaConfig
\ No newline at end of file

Long_LLM/activation_beacon/new/src/llama/configuration_llama.py

+# coding=utf-8
+# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# 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.
+""" LLaMA model configuration"""
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+LLAMA_PRETRAINED_CONFIG_ARCHIVE_MAP = {}
+
+
+class LlamaConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`LlamaModel`]. It is used to instantiate an LLaMA
+    model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
+    defaults will yield a similar configuration to that of the LLaMA-7B.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+
+    Args:
+        vocab_size (`int`, *optional*, defaults to 32000):
+            Vocabulary size of the LLaMA model. Defines the number of different tokens that can be represented by the
+            `inputs_ids` passed when calling [`LlamaModel`]
+        hidden_size (`int`, *optional*, defaults to 4096):
+            Dimension of the hidden representations.
+        intermediate_size (`int`, *optional*, defaults to 11008):
+            Dimension of the MLP representations.
+        num_hidden_layers (`int`, *optional*, defaults to 32):
+            Number of hidden layers in the Transformer decoder.
+        num_attention_heads (`int`, *optional*, defaults to 32):
+            Number of attention heads for each attention layer in the Transformer decoder.
+        num_key_value_heads (`int`, *optional*):
+            This is the number of key_value heads that should be used to implement Grouped Query Attention. If
+            `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
+            `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When
+            converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
+            by meanpooling all the original heads within that group. For more details checkout [this
+            paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to
+            `num_attention_heads`.
+        hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
+            The non-linear activation function (function or string) in the decoder.
+        max_position_embeddings (`int`, *optional*, defaults to 2048):
+            The maximum sequence length that this model might ever be used with. Llama 1 supports up to 2048 tokens,
+            Llama 2 up to 4096, CodeLlama up to 16384.
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        rms_norm_eps (`float`, *optional*, defaults to 1e-06):
+            The epsilon used by the rms normalization layers.
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Whether or not the model should return the last key/values attentions (not used by all models). Only
+            relevant if `config.is_decoder=True`.
+        pad_token_id (`int`, *optional*):
+            Padding token id.
+        bos_token_id (`int`, *optional*, defaults to 1):
+            Beginning of stream token id.
+        eos_token_id (`int`, *optional*, defaults to 2):
+            End of stream token id.
+        pretraining_tp (`int`, *optional*, defaults to 1):
+            Experimental feature. Tensor parallelism rank used during pretraining. Please refer to [this
+            document](https://huggingface.co/docs/transformers/parallelism) to understand more about it. This value is
+            necessary to ensure exact reproducibility of the pretraining results. Please refer to [this
+            issue](https://github.com/pytorch/pytorch/issues/76232).
+        tie_word_embeddings (`bool`, *optional*, defaults to `False`):
+            Whether to tie weight embeddings
+        rope_theta (`float`, *optional*, defaults to 10000.0):
+            The base period of the RoPE embeddings.
+        rope_scaling (`Dict`, *optional*):
+            Dictionary containing the scaling configuration for the RoPE embeddings. Currently supports two scaling
+            strategies: linear and dynamic. Their scaling factor must be a float greater than 1. The expected format is
+            `{"type": strategy name, "factor": scaling factor}`. When using this flag, don't update
+            `max_position_embeddings` to the expected new maximum. See the following thread for more information on how
+            these scaling strategies behave:
+            https://www.reddit.com/r/LocalLLaMA/comments/14mrgpr/dynamically_scaled_rope_further_increases/. This is an
+            experimental feature, subject to breaking API changes in future versions.
+        attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`):
+            Whether to use a bias in the query, key, value and output projection layers during self-attention.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+
+    ```python
+    >>> from transformers import LlamaModel, LlamaConfig
+
+    >>> # Initializing a LLaMA llama-7b style configuration
+    >>> configuration = LlamaConfig()
+
+    >>> # Initializing a model from the llama-7b style configuration
+    >>> model = LlamaModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "llama"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        vocab_size=32000,
+        hidden_size=4096,
+        intermediate_size=11008,
+        num_hidden_layers=32,
+        num_attention_heads=32,
+        num_key_value_heads=None,
+        hidden_act="silu",
+        max_position_embeddings=2048,
+        initializer_range=0.02,
+        rms_norm_eps=1e-6,
+        use_cache=True,
+        pad_token_id=None,
+        bos_token_id=1,
+        eos_token_id=2,
+        pretraining_tp=1,
+        tie_word_embeddings=False,
+        rope_theta=10000.0,
+        rope_scaling=None,
+        attention_bias=False,
+        attention_dropout=0.0,
+        beacon_window=1024,
+        beacon_stride=1024,
+        beacon_attn="full-coverage",
+        beacon_ratio=[2,4,8,16,32],
+        beacon_ratio_mix="step-random",
+        beacon_param=[],
+        beacon_embed_init="eos",
+        beacon_sink_size=0,
+        beacon_attend_prev=True,
+        beacon_pos="interleave",
+        beacon_parallel_window=1,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+
+        # for backward compatibility
+        if num_key_value_heads is None:
+            num_key_value_heads = num_attention_heads
+
+        self.num_key_value_heads = num_key_value_heads
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.pretraining_tp = pretraining_tp
+        self.use_cache = use_cache
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+        self._rope_scaling_validation()
+        self.attention_bias = attention_bias
+        self.attention_dropout = attention_dropout
+        
+        self.beacon_window = beacon_window
+        self.beacon_stride = beacon_stride
+        self.beacon_attn = beacon_attn
+        self.beacon_ratio = beacon_ratio
+        self.beacon_ratio_mix = beacon_ratio_mix
+        self.beacon_param = beacon_param
+        self.beacon_embed_init = beacon_embed_init
+        self.beacon_sink_size = beacon_sink_size
+        self.beacon_attend_prev = beacon_attend_prev
+        self.beacon_pos = beacon_pos
+        self.beacon_parallel_window = beacon_parallel_window
+
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+
+    def _rope_scaling_validation(self):
+        """
+        Validate the `rope_scaling` configuration.
+        """
+        if self.rope_scaling is None:
+            return
+
+        if not isinstance(self.rope_scaling, dict) or len(self.rope_scaling) != 2:
+            raise ValueError(
+                "`rope_scaling` must be a dictionary with with two fields, `type` and `factor`, "
+                f"got {self.rope_scaling}"
+            )
+        rope_scaling_type = self.rope_scaling.get("type", None)
+        rope_scaling_factor = self.rope_scaling.get("factor", None)
+        if rope_scaling_type is None or rope_scaling_type not in ["linear", "dynamic"]:
+            raise ValueError(
+                f"`rope_scaling`'s type field must be one of ['linear', 'dynamic'], got {rope_scaling_type}"
+            )
+        if rope_scaling_factor is None or not isinstance(rope_scaling_factor, float) or rope_scaling_factor <= 1.0:
+            raise ValueError(f"`rope_scaling`'s factor field must be a float > 1, got {rope_scaling_factor}")

Long_LLM/activation_beacon/new/src/llama/modeling_llama.py

+# coding=utf-8
+# Copyright 2024 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# 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.
+""" PyTorch Llama model."""
+import inspect
+import math
+import warnings
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from torch import nn
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask, _prepare_4d_causal_attention_mask_for_sdpa
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, SequenceClassifierOutputWithPast
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import (
+    add_start_docstrings,
+    add_start_docstrings_to_model_forward,
+    is_flash_attn_2_available,
+    is_flash_attn_greater_or_equal_2_10,
+    logging,
+    replace_return_docstrings,
+)
+from transformers.integrations import is_deepspeed_zero3_enabled
+from .configuration_llama import LlamaConfig
+
+
+if is_flash_attn_2_available():
+    from flash_attn import flash_attn_func, flash_attn_varlen_func
+    from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input  # noqa
+
+    _flash_supports_window_size = "window_size" in list(inspect.signature(flash_attn_func).parameters)
+
+from ..modeling_beacon import Memory
+from ..modeling_utils import optional_grad_ctx, compute_loss, BeaconModelOutput
+
+
+logger = logging.get_logger(__name__)
+
+_CONFIG_FOR_DOC = "LlamaConfig"
+
+# Copied from transformers.models.llama.modeling_llama._get_unpad_data
+def _get_unpad_data(attention_mask):
+    seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
+    indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
+    max_seqlen_in_batch = seqlens_in_batch.max().item()
+    cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
+    return (
+        indices,
+        cu_seqlens,
+        max_seqlen_in_batch,
+    )
+
+
+# Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->Llama
+class LlamaRMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        """
+        LlamaRMSNorm is equivalent to T5LayerNorm
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+
+# Copied from transformers.models.llama.modeling_llama.rotate_half
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+class LlamaRotaryEmbedding(nn.Module):
+    def __init__(self, dim, max_position_embeddings=32768, base=10000, device=None):
+        super().__init__()
+
+        self.dim = dim
+        self.max_position_embeddings = max_position_embeddings
+        self.base = base
+        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+        # Build here to make `torch.jit.trace` work.
+        self._set_cos_sin_cache(
+            seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
+        )
+
+    def _set_cos_sin_cache(self, seq_len, device, dtype):
+        self.max_seq_len_cached = seq_len
+        t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.int64).type_as(self.inv_freq)
+
+        freqs = torch.outer(t, self.inv_freq)
+        # Different from paper, but it uses a different permutation in order to obtain the same calculation
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
+        self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
+
+    def forward(self, q, k, position_ids):
+        seq_len = max(position_ids.max().item() + 1, k.shape[2])
+
+        # x: [bs, num_attention_heads, seq_len, head_size]
+        if seq_len > self.max_seq_len_cached:
+            self._set_cos_sin_cache(seq_len=seq_len, device=k.device, dtype=k.dtype)
+
+        # batch_size, 1, key_len, head_dim
+        k_cos = self.cos_cached[position_ids].to(dtype=k.dtype).unsqueeze(1)
+        k_sin = self.sin_cached[position_ids].to(dtype=k.dtype).unsqueeze(1)
+
+        q_cos = k_cos[..., -q.shape[2]:, :]
+        q_sin = k_sin[..., -q.shape[2]:, :]
+
+        q_embed = (q * q_cos) + (rotate_half(q) * q_sin)
+        k_embed = (k * k_cos) + (rotate_half(k) * k_sin)
+        return q_embed, k_embed
+
+
+class LlamaLinearScalingRotaryEmbedding(LlamaRotaryEmbedding):
+    """LlamaRotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev"""
+
+    def __init__(self, dim, max_position_embeddings=32768, base=10000, device=None, scaling_factor=1.0):
+        self.scaling_factor = scaling_factor
+        super().__init__(dim, max_position_embeddings, base, device)
+
+    def _set_cos_sin_cache(self, seq_len, device, dtype):
+        self.max_seq_len_cached = seq_len
+        t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
+        t = t / self.scaling_factor
+
+        freqs = torch.outer(t, self.inv_freq)
+        # Different from paper, but it uses a different permutation in order to obtain the same calculation
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
+        self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
+
+
+class LlamaDynamicNTKScalingRotaryEmbedding(LlamaRotaryEmbedding):
+    """LlamaRotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla"""
+
+    def __init__(self, dim, max_position_embeddings=32768, base=10000, device=None, scaling_factor=1.0):
+        self.scaling_factor = scaling_factor
+        super().__init__(dim, max_position_embeddings, base, device)
+
+    def _set_cos_sin_cache(self, seq_len, device, dtype):
+        self.max_seq_len_cached = seq_len
+
+        if seq_len > self.max_position_embeddings:
+            base = self.base * (
+                (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1)
+            ) ** (self.dim / (self.dim - 2))
+            inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
+            self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+        t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
+
+        freqs = torch.outer(t, self.inv_freq)
+        # Different from paper, but it uses a different permutation in order to obtain the same calculation
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
+        self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
+
+
+class LlamaYarnRotaryEmbedding(nn.Module):
+    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0, beta_slow=2, beta_fast=128):
+        super().__init__()
+
+        self.base = base
+        self.dim = dim
+        self.scaling_factor = scaling_factor
+        self.beta_slow = beta_slow
+        self.beta_fast = beta_fast
+        self.max_position_embeddings = max_position_embeddings
+
+        self._set_cos_sin_cache(
+            seq_len=max_position_embeddings, device=device, dtype=torch.get_default_dtype()
+        )
+
+    def _get_factor(self, device, dtype):
+        # the dimension whose index is smaller than fast_dim rotates more than beta_fast
+        fast_dim = self.dim / 2 * (math.log(self.max_position_embeddings / (2 * math.pi * self.beta_fast)) / math.log(self.base))
+        fast_dim = max(math.floor(fast_dim), 0)
+        # the dimension whose index is bigger than slow_dim rotates less than beta_slow
+        slow_dim = self.dim / 2 * (math.log(self.max_position_embeddings / (2 * math.pi * self.beta_slow)) / math.log(self.base))
+        slow_dim = min(math.ceil(slow_dim), self.dim - 1)
+
+        if fast_dim == slow_dim:
+            slow_dim += 0.001
+
+        # NOTE: very important to use full precision here so that the factor is correct
+        dim_arange = torch.arange(0, self.dim // 2, device=device, dtype=torch.float32)
+        dim_factor = (dim_arange - fast_dim) / (slow_dim - fast_dim)
+        dim_factor = torch.clamp(dim_factor, 0, 1)
+
+        # align with the paper notation
+        return (1 - dim_factor)
+
+    def _get_temperature(self):
+        if self.scaling_factor <= 1:
+            return 1.0
+        return 0.07 * math.log(self.scaling_factor) + 1.0
+    
+    def _set_cos_sin_cache(self, seq_len, device, dtype):
+        dim_arange = torch.arange(0, self.dim, 2, device=device) / self.dim
+        # dim / 2
+        freq = self.base ** dim_arange
+        theta = 1 / freq
+        interleave_theta = theta / self.scaling_factor
+
+        factor = self._get_factor(device, dtype)
+        yarn_theta = factor * theta + (1 - factor) * interleave_theta
+        self.register_buffer("inv_freq", yarn_theta, persistent=False)
+
+        # print(factor, self.inv_freq)
+
+        t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype)
+        freqs = torch.outer(t, self.inv_freq)
+        emb = torch.cat((freqs, freqs), dim=-1)
+
+        # get attention temperature
+        temperature = self._get_temperature()
+
+        self.register_buffer("cos_cached", (emb.cos() * temperature).to(dtype), persistent=False)
+        self.register_buffer("sin_cached", (emb.sin() * temperature).to(dtype), persistent=False)
+        self.max_seq_len_cached = seq_len
+    
+    def forward(self, q, k, position_ids):
+        seq_len = max(position_ids.max().item() + 1, k.shape[2])
+
+        # x: [bs, num_attention_heads, seq_len, head_size]
+        if seq_len > self.max_seq_len_cached:
+            self.scaling_factor = seq_len / self.max_position_embeddings
+            self._set_cos_sin_cache(seq_len=seq_len, device=k.device, dtype=k.dtype)
+
+        k_cos = self.cos_cached[position_ids].to(dtype=k.dtype).unsqueeze(1)
+        k_sin = self.sin_cached[position_ids].to(dtype=k.dtype).unsqueeze(1)
+
+        q_cos = k_cos[..., -q.shape[2]:, :]
+        q_sin = k_sin[..., -q.shape[2]:, :]
+
+        q_embed = (q * q_cos) + (rotate_half(q) * q_sin)
+        k_embed = (k * k_cos) + (rotate_half(k) * k_sin)
+        return q_embed, k_embed
+
+
+# Copied from transformers.models.llama.modeling_llama.LlamaMLP with Llama->Llama
+class LlamaMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        self.act_fn = ACT2FN[config.hidden_act]
+
+        if "mlp" in config.beacon_param:            
+            self.beacon_up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+            self.beacon_up_proj.weight.data.zero_()
+            self.beacon_up_proj._is_hf_initialized = True
+
+            self.beacon_down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+            self.beacon_down_proj.weight.data.zero_()
+            self.beacon_down_proj._is_hf_initialized = True
+    
+    def _init_beacon_proj(self, missing_keys):
+        """Initialize the beacon projection weight with that of the ordinal projection."""
+        if "mlp" in self.config.beacon_param:
+            if is_deepspeed_zero3_enabled():
+                # FIXME: after deepspeed initialization, some weights becomes non-zero
+                # For Llama, there are rows that are full of zeros
+                # For Llama, there are values bigger than 1e29...
+
+                import deepspeed
+                params = [self.up_proj.weight, self.down_proj.weight, self.beacon_up_proj.weight, self.beacon_down_proj.weight]
+                with deepspeed.zero.GatheredParameters(params, modifier_rank=0):
+                    if (self.beacon_up_proj.weight.sum(-1) == 0).any() or (self.beacon_up_proj.weight > 1e29).any():
+                        self.beacon_up_proj.weight.data[:] = self.up_proj.weight.data
+                        self.beacon_down_proj.weight.data[:] = self.down_proj.weight.data
+            else:
+                if any("beacon_up_proj" in missing_key for missing_key in missing_keys):
+                    # only copy the value in-place, without tieing the weight
+                    self.beacon_up_proj.weight.data[:] = self.up_proj.weight.data
+                    self.beacon_down_proj.weight.data[:] = self.down_proj.weight.data
+
+    def forward(self, x, beacon_size, beacon_indices):
+        if "mlp" in self.config.beacon_param:
+            # NOTE: when beacon_pos == "interleave", the beacon_indices points to all beacon tokens in the current window (cached activations + input_ids), so we shall slice out the part corresponding to the input_ids
+            if beacon_size > 0:
+                cur_beacon_indices = beacon_indices[-x.shape[1]:]
+                ordinal_hidden_states = x[:, cur_beacon_indices == 0]
+                beacon_hidden_states = x[:, cur_beacon_indices == 1]
+
+                ordinal_down_proj = self.down_proj(self.act_fn(self.gate_proj(ordinal_hidden_states)) * self.up_proj(ordinal_hidden_states))
+                beacon_down_proj = self.beacon_down_proj(self.act_fn(self.gate_proj(beacon_hidden_states)) * self.beacon_up_proj(beacon_hidden_states))
+
+                down_proj = beacon_down_proj.new_ones(x.shape)
+                down_proj[:, beacon_indices == 0] = ordinal_down_proj
+                down_proj[:, beacon_indices == 1] = beacon_down_proj
+            else:
+                down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+        else:
+            down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+
+        return down_proj
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+class LlamaAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config: LlamaConfig, layer_idx: Optional[int] = None):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        if layer_idx is None:
+            logger.warning_once(
+                f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will "
+                "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "
+                "when creating this class."
+            )
+
+        self.attention_dropout = config.attention_dropout
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = self.hidden_size // self.num_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.max_position_embeddings = config.max_position_embeddings
+        self.rope_theta = config.rope_theta
+        self.is_causal = True
+
+        if (self.head_dim * self.num_heads) != self.hidden_size:
+            raise ValueError(
+                f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
+                f" and `num_heads`: {self.num_heads})."
+            )
+
+        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias)
+        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)
+        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias)
+        self._init_rope()
+
+        # NOTE: add extra parameters for beacon tokens
+        # skip post initialization to speed up loading
+        if "q" in config.beacon_param:
+            self.beacon_q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=self.q_proj.bias is not None)
+            # NOTE: initialize the beacon parameters as zero
+            self.beacon_q_proj.weight.data.zero_()
+            self.beacon_q_proj._is_hf_initialized = True
+        if "k" in config.beacon_param:
+            self.beacon_k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=self.k_proj.bias is not None)
+            self.beacon_k_proj.weight.data.zero_()
+            self.beacon_k_proj._is_hf_initialized = True
+        if "v" in config.beacon_param:
+            self.beacon_v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=self.v_proj.bias is not None)
+            self.beacon_v_proj.weight.data.zero_()
+            self.beacon_v_proj._is_hf_initialized = True
+        if "o" in config.beacon_param:
+            self.beacon_o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=self.o_proj.bias is not None)
+            self.beacon_o_proj.weight.data.zero_()
+            self.beacon_o_proj._is_hf_initialized = True
+
+    def _init_rope(self):
+        if self.config.rope_scaling is None:
+            self.rotary_emb = LlamaRotaryEmbedding(
+                self.head_dim,
+                max_position_embeddings=self.max_position_embeddings,
+                base=self.rope_theta,
+            )
+        else:
+            scaling_type = self.config.rope_scaling["type"]
+            scaling_factor = self.config.rope_scaling["factor"]
+            if scaling_type == "linear":
+                self.rotary_emb = LlamaLinearScalingRotaryEmbedding(
+                    self.head_dim,
+                    max_position_embeddings=self.max_position_embeddings,
+                    scaling_factor=scaling_factor,
+                    base=self.rope_theta,
+                )
+            elif scaling_type == "dynamic":
+                self.rotary_emb = LlamaDynamicNTKScalingRotaryEmbedding(
+                    self.head_dim,
+                    max_position_embeddings=self.max_position_embeddings,
+                    scaling_factor=scaling_factor,
+                    base=self.rope_theta,
+                )
+            elif scaling_type == "yarn":
+                self.rotary_emb = LlamaYarnRotaryEmbedding(
+                    self.head_dim,
+                    max_position_embeddings=self.max_position_embeddings,
+                    scaling_factor=scaling_factor,
+                    base=self.rope_theta,
+                )
+            else:
+                raise ValueError(f"Unknown RoPE scaling type {scaling_type}")
+
+    def _init_beacon_proj(self, missing_keys):
+        """Initialize the beacon projection weight with that of the ordinal projection."""
+        beacon_param = self.config.beacon_param
+        
+        if is_deepspeed_zero3_enabled():
+            # FIXME: after deepspeed initialization, some weights becomes non-zero
+            # For Llama, there are rows that are full of zeros
+            # For Llama, there are values bigger than 1e29...
+
+            import deepspeed
+            if "q" in beacon_param:
+                params = [self.beacon_q_proj.weight, self.q_proj.weight]
+                if self.q_proj.bias is not None:
+                    params.extend([self.beacon_q_proj.bias, self.q_proj.bias])
+                with deepspeed.zero.GatheredParameters(params, modifier_rank=0):
+                    # FIXME: after deepspeed initialization, some weights becomes non-zero, but there are rows that are full of zeros
+                    if (self.beacon_q_proj.weight.sum(-1) == 0).any() or (self.beacon_q_proj.weight > 1e29).any():
+                        self.beacon_q_proj.weight.data[:] = self.q_proj.weight.data
+                        if self.q_proj.bias is not None:
+                            self.beacon_q_proj.bias.data[:] = self.q_proj.bias.data
+            if "k" in beacon_param:
+                params = [self.beacon_k_proj.weight, self.k_proj.weight]
+                if self.k_proj.bias is not None:
+                    params.extend([self.beacon_k_proj.bias, self.k_proj.bias])
+                with deepspeed.zero.GatheredParameters(params, modifier_rank=0):
+                    # FIXME: after deepspeed initialization, some weights becomes non-zero, but there are rows that are full of zeros
+                    if (self.beacon_k_proj.weight.sum(-1) == 0).any() or (self.beacon_k_proj.weight > 1e29).any():
+                        self.beacon_k_proj.weight.data[:] = self.k_proj.weight.data
+                        if self.k_proj.bias is not None:
+                            self.beacon_k_proj.bias.data[:] = self.k_proj.bias.data
+            if "v" in beacon_param:
+                params = [self.beacon_v_proj.weight, self.v_proj.weight]
+                if self.v_proj.bias is not None:
+                    params.extend([self.beacon_v_proj.bias, self.v_proj.bias])
+                with deepspeed.zero.GatheredParameters(params, modifier_rank=0):
+                    # FIXME: after deepspeed initialization, some weights becomes non-zero, but there are rows that are full of zeros
+                    if (self.beacon_v_proj.weight.sum(-1) == 0).any() or (self.beacon_v_proj.weight > 1e29).any():
+                        self.beacon_v_proj.weight.data[:] = self.v_proj.weight.data
+                        if self.v_proj.bias is not None:
+                            self.beacon_v_proj.bias.data[:] = self.v_proj.bias.data
+            if "o" in beacon_param:
+                params = [self.beacon_o_proj.weight, self.o_proj.weight]
+                if self.o_proj.bias is not None:
+                    params.extend([self.beacon_o_proj.bias, self.o_proj.bias])
+                with deepspeed.zero.GatheredParameters(params, modifier_rank=0):
+                    # FIXME: after deepspeed initialization, some weights becomes non-zero, but there are rows that are full of zeros
+                    if (self.beacon_o_proj.weight.sum(-1) == 0).any() or (self.beacon_o_proj.weight > 1e29).any():
+                        self.beacon_o_proj.weight.data[:] = self.o_proj.weight.data
+                        if self.o_proj.bias is not None:
+                            self.beacon_o_proj.bias.data[:] = self.o_proj.bias.data
+        else:
+            # only copy the value in-place, without tieing the weight
+            if "q" in beacon_param and any("beacon_q_proj" in missing_key for missing_key in missing_keys):
+                # FIXME: some beacon weights are not initialized as zero for llama model, why? 
+                # if (self.beacon_q_proj.weight == 0).all():
+                    self.beacon_q_proj.weight.data[:] = self.q_proj.weight.data
+                    if self.q_proj.bias is not None:
+                        self.beacon_q_proj.bias.data[:] = self.q_proj.bias.data
+            if "k" in beacon_param and any("beacon_k_proj" in missing_key for missing_key in missing_keys):
+                # if (self.beacon_k_proj.weight == 0).all():
+                    self.beacon_k_proj.weight.data[:] = self.k_proj.weight.data
+                    if self.k_proj.bias is not None:
+                        self.beacon_k_proj.bias.data[:] = self.k_proj.bias.data
+            if "v" in beacon_param and any("beacon_v_proj" in missing_key for missing_key in missing_keys):
+                # if (self.beacon_v_proj.weight == 0).all():
+                    self.beacon_v_proj.weight.data[:] = self.v_proj.weight.data
+                    if self.v_proj.bias is not None:
+                        self.beacon_v_proj.bias.data[:] = self.v_proj.bias.data
+            if "o" in beacon_param and any("beacon_o_proj" in missing_key for missing_key in missing_keys):
+                # if (self.beacon_o_proj.weight == 0).all():
+                    self.beacon_o_proj.weight.data[:] = self.o_proj.weight.data
+                    if self.o_proj.bias is not None:
+                        self.beacon_o_proj.bias.data[:] = self.o_proj.bias.data
+
+    def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
+        return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
+    
+    def qkv_proj_with_beacon(self, hidden_states, beacon_size, beacon_indices):
+        if beacon_size > 0:
+            # NOTE: when beacon_pos == "interleave", the beacon_indices points to all beacon tokens in the current window (cached activations + input_ids), so we shall slice out the part corresponding to the input_ids
+            cur_beacon_indices = beacon_indices[-hidden_states.shape[1]:]
+
+            ordinal_hidden_states = hidden_states[:, cur_beacon_indices == 0]
+            beacon_hidden_states = hidden_states[:, cur_beacon_indices == 1]
+
+            if "q" in self.config.beacon_param:
+                ordinal_query_states = self.q_proj(ordinal_hidden_states)
+                beacon_query_states = self.beacon_q_proj(beacon_hidden_states)
+                query_states = beacon_query_states.new_zeros((ordinal_query_states.shape[0], cur_beacon_indices.shape[0], ordinal_query_states.shape[2]))
+                query_states[:, cur_beacon_indices == 0] = ordinal_query_states
+                query_states[:, cur_beacon_indices == 1] = beacon_query_states
+                # NOTE: replicate hidden states for beacon tokens in case of parallel windows
+                if (cur_beacon_indices == 2).any():
+                    query_states[:, cur_beacon_indices == 2] = beacon_query_states[:, :(cur_beacon_indices == 2).sum()]
+
+            else:
+                query_states = self.q_proj(hidden_states)
+
+            if "k" in self.config.beacon_param:
+                ordinal_key_states = self.k_proj(ordinal_hidden_states)
+                beacon_key_states = self.beacon_k_proj(beacon_hidden_states)
+                key_states = beacon_key_states.new_zeros((ordinal_key_states.shape[0], cur_beacon_indices.shape[0], ordinal_key_states.shape[2]))
+                key_states[:, cur_beacon_indices == 0] = ordinal_key_states
+                key_states[:, cur_beacon_indices == 1] = beacon_key_states
+                # NOTE: replicate hidden states for beacon tokens in case of parallel windows
+                if (cur_beacon_indices == 2).any():
+                    key_states[:, cur_beacon_indices == 2] = beacon_key_states[:, :(cur_beacon_indices == 2).sum()]
+
+            else:
+                key_states = self.k_proj(hidden_states)
+            
+            if "v" in self.config.beacon_param:
+                ordinal_value_states = self.v_proj(ordinal_hidden_states)
+                beacon_value_states = self.beacon_v_proj(beacon_hidden_states)
+                value_states = beacon_value_states.new_zeros((ordinal_value_states.shape[0], cur_beacon_indices.shape[0], ordinal_value_states.shape[2]))
+                value_states[:, cur_beacon_indices == 0] = ordinal_value_states
+                value_states[:, cur_beacon_indices == 1] = beacon_value_states
+                # NOTE: replicate hidden states for beacon tokens in case of parallel windows
+                if (cur_beacon_indices == 2).any():
+                    value_states[:, cur_beacon_indices == 2] = beacon_value_states[:, :(cur_beacon_indices == 2).sum()]
+            else:
+                value_states = self.v_proj(hidden_states)
+
+        else:
+            query_states = self.q_proj(hidden_states)
+            key_states = self.k_proj(hidden_states)
+            value_states = self.v_proj(hidden_states)
+
+        return query_states, key_states, value_states
+    
+    def o_proj_with_beacon(self, attn_output, beacon_size, beacon_indices):
+        if beacon_size > 0:
+            # NOTE: when beacon_pos == "interleave", the beacon_indices points to all beacon tokens in the current window (cached activations + input_ids), so we shall slice out the part corresponding to the input_ids
+            cur_beacon_indices = beacon_indices[-attn_output.shape[1]:]
+
+            if "o" in self.config.beacon_param:
+                ordinal_attn_output = self.o_proj(attn_output[:, cur_beacon_indices == 0])
+                beacon_attn_output = self.beacon_o_proj(attn_output[:, cur_beacon_indices == 1])
+                attn_output = beacon_attn_output.new_zeros(attn_output.shape)
+                attn_output[:, cur_beacon_indices == 0] = ordinal_attn_output
+                attn_output[:, cur_beacon_indices == 1] = beacon_attn_output
+                # NOTE: replicate hidden states for beacon tokens in case of parallel windows
+                # if (cur_beacon_indices == 2).any():
+                #     attn_output[:, cur_beacon_indices == 2] = beacon_attn_output[:, :(cur_beacon_indices == 2).sum()]
+            else:
+                attn_output = self.o_proj(attn_output)
+        else:
+            attn_output = self.o_proj(attn_output)
+        return attn_output
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Cache] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        if "padding_mask" in kwargs:
+            warnings.warn(
+                "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
+            )
+
+        bsz, q_len, _ = hidden_states.size()
+        kv_seq_len = hidden_states.shape[-2]
+        past_key, past_value, beacon_size, beacon_indices = past_key_value
+
+        if past_key is not None:
+            past_seq_len = past_key.shape[2]
+            kv_seq_len += past_seq_len
+        else:
+            past_seq_len = 0
+
+        query_states, key_states, value_states = self.qkv_proj_with_beacon(hidden_states, beacon_size, beacon_indices)
+
+        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+        # return keys and values before rope
+        # NOTE: incrementally return keys and values for efficiency 
+        past_key_value = (key_states, value_states, beacon_size, beacon_indices)
+
+        if past_key is not None:
+            # reuse k, v, self_attention
+            key_states = torch.cat([past_key, key_states], dim=2)
+            value_states = torch.cat([past_value, value_states], dim=2)
+        
+        query_states, key_states = self.rotary_emb(query_states, key_states, position_ids)
+
+        key_states = repeat_kv(key_states, self.num_key_value_groups)
+        value_states = repeat_kv(value_states, self.num_key_value_groups)
+
+        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
+
+        if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
+            raise ValueError(
+                f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"
+                f" {attn_weights.size()}"
+            )
+
+        if attention_mask is not None:
+            if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
+                raise ValueError(
+                    f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
+                )
+            attn_weights = attn_weights + attention_mask
+
+        # upcast attention to fp32
+        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
+        attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)
+        attn_output = torch.matmul(attn_weights, value_states)
+
+        if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
+            raise ValueError(
+                f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
+                f" {attn_output.size()}"
+            )
+
+        attn_output = attn_output.transpose(1, 2).contiguous()
+
+        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
+
+        attn_output = self.o_proj_with_beacon(attn_output, beacon_size, beacon_indices)
+
+        if not output_attentions:
+            attn_weights = None
+
+        return attn_output, attn_weights, past_key_value
+
+
+class LlamaSdpaAttention(LlamaAttention):
+    """
+    Llama attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from
+    `LlamaAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to
+    SDPA API.
+    """
+
+    # Adapted from LlamaAttention.forward
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Cache] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        if output_attentions:
+            # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented.
+            logger.warning_once(
+                "LlamaModel is using LlamaSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, "
+                'but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
+            )
+            return super().forward(
+                hidden_states=hidden_states,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_value=past_key_value,
+                output_attentions=output_attentions,
+                use_cache=use_cache,
+            )
+        bsz, q_len, _ = hidden_states.size()
+        kv_seq_len = hidden_states.shape[-2]
+        past_key, past_value, beacon_size, beacon_indices = past_key_value
+        if past_key is not None:
+            past_seq_len = past_key.shape[2]
+            kv_seq_len += past_seq_len
+        else:
+            past_seq_len = 0
+
+        query_states, key_states, value_states = self.qkv_proj_with_beacon(hidden_states, beacon_size, beacon_indices)
+
+        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+        # return keys and values before rope
+        # NOTE: incrementally return keys and values for efficiency 
+        past_key_value = (key_states, value_states, beacon_size, beacon_indices)
+
+        if past_key is not None:
+            # reuse k, v, self_attention
+            key_states = torch.cat([past_key, key_states], dim=2)
+            value_states = torch.cat([past_value, value_states], dim=2)
+        
+        query_states, key_states = self.rotary_emb(query_states, key_states, position_ids)
+
+        key_states = repeat_kv(key_states, self.num_key_value_groups)
+        value_states = repeat_kv(value_states, self.num_key_value_groups)
+
+        if attention_mask is not None:
+            if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
+                raise ValueError(
+                    f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
+                )
+
+        # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,
+        # Reference: https://github.com/pytorch/pytorch/issues/112577.
+        if query_states.device.type == "cuda" and attention_mask is not None:
+            query_states = query_states.contiguous()
+            key_states = key_states.contiguous()
+            value_states = value_states.contiguous()
+
+        attn_output = torch.nn.functional.scaled_dot_product_attention(
+            query_states,
+            key_states,
+            value_states,
+            attn_mask=attention_mask,
+            dropout_p=self.attention_dropout if self.training else 0.0,
+            # The q_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case q_len == 1.
+            is_causal=self.is_causal and attention_mask is None and q_len > 1,
+        )
+
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
+        attn_output = self.o_proj_with_beacon(attn_output, beacon_size, beacon_indices)
+
+        return attn_output, None, past_key_value
+
+
+class LlamaFlashAttention2(LlamaAttention):
+    """
+    Llama flash attention module. This module inherits from `LlamaAttention` as the weights of the module stays
+    untouched. The only required change would be on the forward pass where it needs to correctly call the public API of
+    flash attention and deal with padding tokens in case the input contains any of them.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
+        # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.
+        # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left).
+        self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Cache] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        output_attentions = False
+
+        bsz, q_len, _ = hidden_states.size()
+        kv_seq_len = hidden_states.shape[-2]
+
+        past_key, past_value, beacon_size, beacon_indices = past_key_value
+        if past_key is not None:
+            past_seq_len = past_key.shape[2]
+            kv_seq_len += past_seq_len
+        else:
+            past_seq_len = 0
+
+        query_states, key_states, value_states = self.qkv_proj_with_beacon(hidden_states, beacon_size, beacon_indices)
+
+        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+        # return keys and values before rope
+        # NOTE: incrementally return keys and values for efficiency 
+        past_key_value = (key_states, value_states, beacon_size, beacon_indices)
+
+        if past_key is not None:
+            # reuse k, v, self_attention
+            key_states = torch.cat([past_key, key_states], dim=2)
+            value_states = torch.cat([past_value, value_states], dim=2)
+
+        query_states, key_states = self.rotary_emb(query_states, key_states, position_ids)
+
+        # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache
+        # to be able to avoid many of these transpose/reshape/view.
+        query_states = query_states.transpose(1, 2)
+        key_states = key_states.transpose(1, 2)
+        value_states = value_states.transpose(1, 2)
+
+        dropout_rate = self.attention_dropout if self.training else 0.0
+
+        # In PEFT, usually we cast the layer norms in float32 for training stability reasons
+        # therefore the input hidden states gets silently casted in float32. Hence, we need
+        # cast them back in the correct dtype just to be sure everything works as expected.
+        # This might slowdown training & inference so it is recommended to not cast the LayerNorms
+        # in fp32. (LlamaRMSNorm handles it correctly)
+
+        input_dtype = query_states.dtype
+        if input_dtype == torch.float32:
+            if torch.is_autocast_enabled():
+                target_dtype = torch.get_autocast_gpu_dtype()
+            # Handle the case where the model is quantized
+            elif hasattr(self.config, "_pre_quantization_dtype"):
+                target_dtype = self.config._pre_quantization_dtype
+            else:
+                target_dtype = self.q_proj.weight.dtype
+
+            logger.warning_once(
+                f"The input hidden states seems to be silently casted in float32, this might be related to"
+                f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
+                f" {target_dtype}."
+            )
+
+            query_states = query_states.to(target_dtype)
+            key_states = key_states.to(target_dtype)
+            value_states = value_states.to(target_dtype)
+
+        attn_output = self._flash_attention_forward(
+            query_states, 
+            key_states, 
+            value_states, 
+            attention_mask, 
+            q_len, 
+            dropout=dropout_rate
+        )
+
+        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()
+        attn_output = self.o_proj_with_beacon(attn_output, beacon_size, beacon_indices)
+
+        if not output_attentions:
+            attn_weights = None
+
+        return attn_output, attn_weights, past_key_value
+
+    def _flash_attention_forward(
+        self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None
+    ):
+        """
+        Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token
+        first unpad the input, then computes the attention scores and pad the final attention scores.
+
+        Args:
+            query_states (`torch.Tensor`):
+                Input query states to be passed to Flash Attention API
+            key_states (`torch.Tensor`):
+                Input key states to be passed to Flash Attention API
+            value_states (`torch.Tensor`):
+                Input value states to be passed to Flash Attention API
+            attention_mask (`torch.Tensor`):
+                The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
+                position of padding tokens and 1 for the position of non-padding tokens.
+            dropout (`float`):
+                Attention dropout
+            softmax_scale (`float`, *optional*):
+                The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
+        """
+        if not self._flash_attn_uses_top_left_mask:
+            causal = self.is_causal
+        else:
+            # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in LlamaFlashAttention2 __init__.
+            causal = self.is_causal and query_length != 1
+
+        # Contains at least one padding token in the sequence
+        if attention_mask is not None:
+            batch_size = query_states.shape[0]
+            query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
+                query_states, key_states, value_states, attention_mask, query_length
+            )
+
+            cu_seqlens_q, cu_seqlens_k = cu_seq_lens
+            max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
+
+            attn_output_unpad = flash_attn_varlen_func(
+                query_states,
+                key_states,
+                value_states,
+                cu_seqlens_q=cu_seqlens_q,
+                cu_seqlens_k=cu_seqlens_k,
+                max_seqlen_q=max_seqlen_in_batch_q,
+                max_seqlen_k=max_seqlen_in_batch_k,
+                dropout_p=dropout,
+                softmax_scale=softmax_scale,
+                causal=causal,
+            )
+
+            attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)
+        else:
+            attn_output = flash_attn_func(
+                query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=causal
+            )
+
+        return attn_output
+
+    def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
+        indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
+        batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
+
+        key_layer = index_first_axis(
+            key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
+        )
+        value_layer = index_first_axis(
+            value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
+        )
+        if query_length == kv_seq_len:
+            query_layer = index_first_axis(
+                query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim), indices_k
+            )
+            cu_seqlens_q = cu_seqlens_k
+            max_seqlen_in_batch_q = max_seqlen_in_batch_k
+            indices_q = indices_k
+        elif query_length == 1:
+            max_seqlen_in_batch_q = 1
+            cu_seqlens_q = torch.arange(
+                batch_size + 1, dtype=torch.int32, device=query_layer.device
+            )  # There is a memcpy here, that is very bad.
+            indices_q = cu_seqlens_q[:-1]
+            query_layer = query_layer.squeeze(1)
+        else:
+            # The -q_len: slice assumes left padding.
+            attention_mask = attention_mask[:, -query_length:]
+            query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)
+
+        return (
+            query_layer,
+            key_layer,
+            value_layer,
+            indices_q,
+            (cu_seqlens_q, cu_seqlens_k),
+            (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
+        )
+
+
+LLAMA_ATTENTION_CLASSES = {
+    "eager": LlamaAttention,
+    "sdpa": LlamaSdpaAttention,
+    "flash_attention_2": LlamaFlashAttention2,
+}
+
+
+class LlamaDecoderLayer(nn.Module):
+    def __init__(self, config: LlamaConfig, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+
+        self.self_attn = LLAMA_ATTENTION_CLASSES[config._attn_implementation](config, layer_idx)
+
+        self.mlp = LlamaMLP(config)
+        self.input_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        output_attentions: Optional[bool] = False,
+        use_cache: Optional[bool] = False,
+        **kwargs,
+    ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
+        """
+        Args:
+            hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
+            attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
+                `(batch, sequence_length)` where padding elements are indicated by 0.
+            output_attentions (`bool`, *optional*):
+                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+                returned tensors for more detail.
+            use_cache (`bool`, *optional*):
+                If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
+                (see `past_key_values`).
+            past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
+        """
+        if "padding_mask" in kwargs:
+            warnings.warn(
+                "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
+            )
+
+        # NOTE: get beacon_size in case the mlp is included in beacon_param
+        past_key, past_value, beacon_size, beacon_indices = past_key_value
+
+        residual = hidden_states
+
+        hidden_states = self.input_layernorm(hidden_states)
+
+        # Self Attention
+        hidden_states, self_attn_weights, present_key_value = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_value=past_key_value,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+            **kwargs,
+        )
+        hidden_states = residual + hidden_states
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states, beacon_size, beacon_indices)
+        hidden_states = residual + hidden_states
+
+        outputs = (hidden_states,)
+
+        if output_attentions:
+            outputs += (self_attn_weights,)
+
+        if use_cache:
+            outputs += (present_key_value,)
+
+        return outputs
+
+
+LLAMA_START_DOCSTRING = r"""
+    This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
+    library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
+    etc.)
+
+    This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
+    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
+    and behavior.
+
+    Parameters:
+        config ([`LlamaConfig`]):
+            Model configuration class with all the parameters of the model. Initializing with a config file does not
+            load the weights associated with the model, only the configuration. Check out the
+            [`~PreTrainedModel.from_pretrained`] method to load the model weights.
+"""
+
+
+@add_start_docstrings(
+    "The bare Llama Model outputting raw hidden-states without any specific head on top.",
+    LLAMA_START_DOCSTRING,
+)
+class LlamaPreTrainedModel(PreTrainedModel):
+    config_class = LlamaConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["LlamaDecoderLayer"]
+    _skip_keys_device_placement = "past_key_values"
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+    _supports_cache_class = True
+
+    def _init_weights(self, module):
+        std = self.config.initializer_range
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+
+LLAMA_INPUTS_DOCSTRING = r"""
+    Args:
+        input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+            Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
+            it.
+
+            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+            [`PreTrainedTokenizer.__call__`] for details.
+
+            [What are input IDs?](../glossary#input-ids)
+        attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
+
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **masked**.
+
+            [What are attention masks?](../glossary#attention-mask)
+
+            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+            [`PreTrainedTokenizer.__call__`] for details.
+
+            If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see
+            `past_key_values`).
+
+            If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]
+            and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more
+            information on the default strategy.
+
+            - 1 indicates the head is **not masked**,
+            - 0 indicates the head is **masked**.
+        position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
+            config.n_positions - 1]`.
+
+            [What are position IDs?](../glossary#position-ids)
+        past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*):
+            Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
+            blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values`
+            returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`.
+
+            Two formats are allowed:
+            - a [`~cache_utils.Cache`] instance;
+            - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of
+            shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy
+            cache format.
+
+            The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the
+            legacy cache format will be returned.
+
+            If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't
+            have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids`
+            of shape `(batch_size, sequence_length)`.
+        inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
+            Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
+            is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
+            model's internal embedding lookup matrix.
+        use_cache (`bool`, *optional*):
+            If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
+            `past_key_values`).
+        output_attentions (`bool`, *optional*):
+            Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
+            tensors for more detail.
+        output_hidden_states (`bool`, *optional*):
+            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
+            more detail.
+        return_dict (`bool`, *optional*):
+            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+"""
+
+
+@add_start_docstrings(
+    "The bare Llama Model outputting raw hidden-states without any specific head on top.",
+    LLAMA_START_DOCSTRING,
+)
+class LlamaModel(LlamaPreTrainedModel):
+    """
+    Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`LlamaDecoderLayer`]
+
+    Args:
+        config: LlamaConfig
+    """
+
+    def __init__(self, config: LlamaConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+
+        # BEACON: add beacon embedding
+        self.beacon_embed_tokens = nn.Embedding(1, config.hidden_size, self.padding_idx)
+        self.beacon_embed_tokens._is_hf_initialized = True
+
+        self.layers = nn.ModuleList(
+            [LlamaDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self._attn_implementation = config._attn_implementation
+        self.norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        self.gradient_checkpointing = False
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def _init_beacon_embed(self, missing_keys):
+        """Initialize the beacon token embedding with that of the eos token."""
+        if is_deepspeed_zero3_enabled():
+            import deepspeed
+            params = [self.beacon_embed_tokens.weight, self.embed_tokens.weight]
+            with deepspeed.zero.GatheredParameters(params, modifier_rank=0):
+                # deepspeed will initialize the parameters to zero
+                if (self.beacon_embed_tokens.weight == 0).all():
+                    if self.config.beacon_embed_init == "bos":
+                        self.beacon_embed_tokens.weight.data[:] = self.embed_tokens.weight.data[self.config.bos_token_id]
+                    elif self.config.beacon_embed_init == "eos":
+                        if isinstance(self.config.eos_token_id, list):
+                            eos_token_id = self.config.eos_token_id[0]
+                        else:
+                            eos_token_id = self.config.eos_token_id
+                        self.beacon_embed_tokens.weight.data[:] = self.embed_tokens.weight.data[eos_token_id]
+                    else:
+                        raise NotImplementedError(f"Make sure beacon_embed_init is either eos or bos, found {self.config.beacon_embed_init}")
+        else:
+            if any("beacon_embed_tokens" in missing_key for missing_key in missing_keys):
+                if self.config.beacon_embed_init == "bos":
+                    self.beacon_embed_tokens.weight.data[:] = self.embed_tokens.weight.data[self.config.bos_token_id]
+                elif self.config.beacon_embed_init == "eos":
+                    if isinstance(self.config.eos_token_id, list):
+                        eos_token_id = self.config.eos_token_id[0]
+                    else:
+                        eos_token_id = self.config.eos_token_id
+                    self.beacon_embed_tokens.weight.data[:] = self.embed_tokens.weight.data[eos_token_id]
+                else:
+                    raise NotImplementedError(f"Make sure beacon_embed_init is either eos or bos, found {self.config.beacon_embed_init}")
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.embed_tokens = value
+
+    @add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING)
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        # BEACON: always use cache
+        use_cache = True
+
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # retrieve input_ids and inputs_embeds
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+        elif input_ids is not None:
+            batch_size, seq_length = input_ids.shape[:2]
+        elif inputs_embeds is not None:
+            batch_size, seq_length = inputs_embeds.shape[:2]
+        else:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+        past_key, past_value, beacon_size, beacon_indices = past_key_values[0]
+
+        # BEACON: separately embed ordinal tokens and beacon tokens because ordinal tokens do not receive gradients
+        if beacon_size > 0:
+            # NOTE: when beacon_pos == "interleave", the beacon_indices points to all beacon tokens in the current window (cached activations + input_ids), so we shall slice out the part corresponding to the input_ids
+            cur_beacon_indices = beacon_indices[-input_ids.shape[1]:]
+
+            ordinal_input_ids = input_ids[:, cur_beacon_indices == 0]
+            beacon_input_ids = input_ids[:, cur_beacon_indices > 0]
+            ordinal_inputs_embeds = self.embed_tokens(ordinal_input_ids)
+            beacon_input_embeds = self.beacon_embed_tokens(beacon_input_ids - self.config.vocab_size)
+            # create a new embedding tensor
+            inputs_embeds = beacon_input_embeds.new_zeros(*input_ids.shape, beacon_input_embeds.shape[-1])
+            inputs_embeds[:, cur_beacon_indices == 0] = ordinal_inputs_embeds
+            inputs_embeds[:, cur_beacon_indices > 0] = beacon_input_embeds
+
+        else:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        # embed positions
+        hidden_states = inputs_embeds
+
+        # print(f"input_ids:          {input_ids}")
+        # print(f"beacon_indices:     {beacon_indices}")
+        # print(f"position_ids:       {position_ids}")
+        # print(f"attention_mask:\n{attention_mask == 0}")
+        # x = input()
+        # if x == "s":
+        #     return
+
+        # decoder layers
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+        # BEACON: still use tuple to organize cache
+        next_decoder_cache = () if use_cache else None
+
+        for idx, decoder_layer in enumerate(self.layers):
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            # BEACON: slice out the past_key_value of the corresponding layer
+            past_key_value = past_key_values[idx] if past_key_values is not None else None
+
+            if self.gradient_checkpointing and self.training:
+                layer_outputs = self._gradient_checkpointing_func(
+                    decoder_layer.__call__,
+                    hidden_states,
+                    attention_mask,
+                    position_ids,
+                    past_key_value,
+                    output_attentions,
+                    use_cache,
+                )
+            else:
+                layer_outputs = decoder_layer(
+                    hidden_states,
+                    attention_mask=attention_mask,
+                    position_ids=position_ids,
+                    past_key_value=past_key_value,
+                    output_attentions=output_attentions,
+                    use_cache=use_cache,
+                )
+
+            hidden_states = layer_outputs[0]
+
+            if use_cache:
+                next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)
+
+            if output_attentions:
+                all_self_attns += (layer_outputs[1],)
+
+        hidden_states = self.norm(hidden_states)
+
+        # add hidden states from the last decoder layer
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        next_cache = next_decoder_cache if use_cache else None
+
+        if not return_dict:
+            return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=next_cache,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+        )
+
+
+class LlamaForCausalLM(LlamaPreTrainedModel):
+    _tied_weights_keys = ["lm_head.weight"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = LlamaModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    def get_decoder(self):
+        return self.model
+
+    @classmethod
+    def from_pretrained(cls, *args, **kwargs):
+        """Override the default from_pretrained to extend vocab size according to beacon_size."""
+        kwargs.update(output_loading_info=True)
+        model, loading_info = super().from_pretrained(*args, **kwargs)
+
+        # NOTE: set memory after from_pretrained because there may be another transformer model inside the Memory object, which may cause weird erros during loading
+        config = model.config
+        model.memory = Memory(
+            model_config=config,
+            k_seq_dim=2,
+            v_seq_dim=2,
+        )
+
+        missing_keys = loading_info["missing_keys"]
+        # NOTE: the beacon parameters may or may not be loaded from the checkpoint
+        # if it is loaded from the checkpoint, we should not re-initilize it
+        model.model._init_beacon_embed(missing_keys)
+        # initialize weights of possible q,k,v,o,mlp
+        for layer in model.model.layers:
+            layer.self_attn._init_beacon_proj(missing_keys)
+            layer.mlp._init_beacon_proj(missing_keys)
+
+        return model
+
+    def _native_forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        shift_labels: Optional[bool] = True,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BeaconModelOutput]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # when we directly call _native_forward, the past_key_values would be None
+        if past_key_values is None:
+            # NOTE: set beacon size to 0 to avoid using any beacon parameters, see LlamaAttention.forward
+            past_key_values = [(None, None, 0, None) for _ in range(self.config.num_hidden_layers)]
+
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        hidden_states = outputs[0]
+        logits = self.lm_head(hidden_states)
+        logits = logits.float()
+
+        loss = None
+        batch_loss = None
+        valid_token_num = None
+        
+        if labels is not None:
+            loss, batch_loss, valid_token_num = compute_loss(logits, labels, shift=shift_labels)
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+
+        return BeaconModelOutput(
+            loss=loss,
+            batch_loss=batch_loss,
+            valid_token_num=valid_token_num,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def _beacon_forward(self, 
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ):
+        # t1 = time.time()
+
+        # initialize cache
+        self.memory.prepare(
+            input_ids=input_ids, 
+            attention_mask=attention_mask, 
+            labels=labels
+        )
+
+        # t2 = time.time()
+
+        # after the first window, one token at a time
+        while not self.memory.finish:
+
+            # t3 = time.time()
+
+            input_ids, attention_mask, position_ids, past_key_values, labels = self.memory.step()
+
+            # t4 = time.time()
+
+            outputs = self._native_forward(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                inputs_embeds=inputs_embeds,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+                labels=labels,
+                # NOTE: the labels have been shifted so that all tokens in the window have the proper loss
+                shift_labels=False,
+            )
+
+            # t5 = time.time()
+
+            # update past_key_values
+            self.memory.update_memory(outputs.past_key_values)
+
+            # t6 = time.time()
+
+            if labels is not None:
+                # update loss
+                self.memory.update_loss(outputs.batch_loss, outputs.valid_token_num)
+
+            # t7 = time.time()
+
+            # print(f"step time: {t4-t3}, forward time: {t5-t4}, update time: {t6-t5}, loss time: {t7-t6}")
+            # input()
+
+        # t8 = time.time()
+
+        # output loss, past_key_values, and perplexity
+        outputs = self.memory.output(outputs)
+
+        # t9 = time.time()
+
+        # print(f"output time:            {t9-t8}")
+        # input()
+
+        return outputs
+    
+    def forward(self, **kwargs):
+        """Forward computation over a batch of sequences.
+        """
+        # only allow gradient when training
+        with optional_grad_ctx(with_grad=self.training):
+            # we can disable beacon to use the original llama
+            if hasattr(self, "_enable_beacon") and self._enable_beacon == False:
+                return self._native_forward(**kwargs)
+            else:
+                return self._beacon_forward(**kwargs)
+
+    def prepare_inputs_for_generation(
+        self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
+    ):
+        if past_key_values:
+            input_ids = input_ids[:, -1:]
+
+        position_ids = kwargs.get("position_ids", None)
+        if attention_mask is not None and position_ids is None:
+            # create position_ids on the fly for batch generation
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            if past_key_values:
+                position_ids = position_ids[:, -1].unsqueeze(-1)
+
+        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+        if inputs_embeds is not None and past_key_values is None:
+            model_inputs = {"inputs_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids}
+
+        model_inputs.update(
+            {
+                "position_ids": position_ids,
+                "past_key_values": past_key_values,
+                "use_cache": kwargs.get("use_cache"),
+                "attention_mask": attention_mask,
+            }
+        )
+        return model_inputs
+
+    @staticmethod
+    def _reorder_cache(past_key_values, beam_idx):
+        reordered_past = ()
+        for layer_past in past_key_values:
+            reordered_past += (
+                tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
+            )
+        return reordered_past

Long_LLM/activation_beacon/new/src/metrics.py

+import os
+import json
+import inspect
+import numpy as np
+from functools import partial
+from rouge import Rouge
+from tqdm import tqdm
+from transformers.utils import logging
+from .utils import makedirs, split_file_dir_name_ext, normalize_text
+
+logger = logging.get_logger(__name__)
+
+
+class Metric:
+    """Class for computing metrics and some post-processings."""
+    @classmethod
+    def get_metric_fn(cls, metrics, **kwds):
+        assert isinstance(metrics, list) or isinstance(metrics, tuple), "You must pass metric_names in a list or tuple!"
+        return_metrics = {}
+        # get all methods
+        metric_fns = []
+
+        all_metric_names = [x[0] for x in inspect.getmembers(cls, predicate=inspect.isfunction) if not x[0].startswith("get_")]
+        for metric_name in metrics:
+            if metric_name in all_metric_names:
+                metric_fns.append(partial(getattr(cls, metric_name), **kwds))
+            else:
+                raise NotImplementedError(f"Metric {metric_name} not implemented!")
+
+        def compute_metrics(*args, **kwargs):
+            for metric_fn in metric_fns:
+                # call corresponding method
+                metric = metric_fn(*args, **kwargs)
+                # NOTE: some metric_fn are only used for post-processing and saving results, which return None by default
+                if metric is not None:
+                    return_metrics.update(metric)
+            return return_metrics
+        return compute_metrics
+    
+    def get_save_path(eval_data, output_dir=None, field="result", save_name=None):
+        """
+        if output_dir is None:
+            -> {eval_data_dir}/{eval_data_name}.{field}.{save_name}.{eval_data_ext}
+        else:
+            -> {output_dir}/{eval_data_name}.{field}.{save_name}.{eval_data_ext}
+        """
+        eval_data_dir, eval_data_name, eval_data_ext = split_file_dir_name_ext(eval_data)
+        if output_dir is None:
+            output_dir = eval_data_dir
+        fields = [eval_data_name, field]
+        if save_name is not None:
+            fields.append(save_name)
+        save_path = os.path.join(output_dir, ".".join(fields) + eval_data_ext)
+        makedirs(save_path)
+        return save_path
+
+    def save_result(preds, labels, save_path, indices=None, **kwargs):
+        if len(preds) != len(labels):
+            logger.warning(f"There are {len(preds)} samples in predictions while {len(labels)} samples in labels!")
+            labels = labels[:min(len(preds), len(labels))]
+            preds = preds[:min(len(preds), len(labels))]
+        
+        with open(save_path, "w", encoding="utf-8") as f:
+            for i, (pred, label) in enumerate(zip(preds, labels)):
+                item = {
+                    "prediction": pred,
+                    "target": label,
+                }
+                if indices is not None:
+                    item["index"] = indices[i]
+                f.write(json.dumps(item, ensure_ascii=False) + "\n")
+
+    def rouge(preds, labels, **kwargs):
+        rouge = Rouge()
+
+        if len(preds) != len(labels):
+            logger.warning(f"There are {len(preds)} samples in predictions while {len(labels)} samples in labels!")
+            labels = labels[:min(len(preds), len(labels))]
+            preds = preds[:min(len(preds), len(labels))]
+
+        preds = normalize_text(preds)
+        labels = normalize_text(labels)
+
+        # filter empty preditions
+        preds = [":)" if len(pred) == 0 else pred for pred in preds]
+
+        score = rouge.get_scores(preds, labels, avg=True)
+
+        metric = {
+            "rouge-1": score["rouge-1"]["f"],
+            "rouge-2": score["rouge-2"]["f"],
+            "rouge-l": score["rouge-2"]["f"],
+        }
+        return metric
+    
+    # def acc(eval_data=None, **kwds):
+    #     if eval_data is not None:
+    #         data_labels = Metric._prepare_label(eval_data)
+
+    #     def compute_metric(indices, preds, labels=None, **kwargs):
+    #         if labels is None:
+    #             labels = data_labels
+
+    #         if len(preds) != len(labels):
+    #             logger.warning(f"There are {len(preds)} queries in predictions while {len(labels)} queries in labels!")
+
+    #         labels = [labels[query_id] for query_id in indices]
+
+    #         preds = normalize_text(preds)
+    #         labels = normalize_text(labels)
+
+    #         overlap = 0
+    #         for pred, label in zip(preds, labels):
+    #             if pred == label:
+    #                 overlap += 1
+
+    #         metric = {
+    #             "acc": overlap / len(preds),
+    #         }
+    #         return metric
+    #     return compute_metric

Long_LLM/activation_beacon/new/src/mistral/__init__.py

+from .modeling_mistral import MistralForCausalLM
+from .configuration_mistral import MistralConfig
\ No newline at end of file

Long_LLM/activation_beacon/new/src/mistral/configuration_mistral.py

+# coding=utf-8
+# Copyright 2023 Mistral AI and the HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+""" Mistral model configuration"""
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+MISTRAL_PRETRAINED_CONFIG_ARCHIVE_MAP = {
+    "mistralai/Mistral-7B-v0.1": "https://huggingface.co/mistralai/Mistral-7B-v0.1/resolve/main/config.json",
+    "mistralai/Mistral-7B-Instruct-v0.1": "https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.1/resolve/main/config.json",
+}
+
+
+class MistralConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`MistralModel`]. It is used to instantiate an
+    Mistral model according to the specified arguments, defining the model architecture. Instantiating a configuration
+    with the defaults will yield a similar configuration to that of the Mistral-7B-v0.1 or Mistral-7B-Instruct-v0.1.
+
+    [mistralai/Mistral-7B-v0.1](https://huggingface.co/mistralai/Mistral-7B-v0.1)
+    [mistralai/Mistral-7B-Instruct-v0.1](https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.1)
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+
+    Args:
+        vocab_size (`int`, *optional*, defaults to 32000):
+            Vocabulary size of the Mistral model. Defines the number of different tokens that can be represented by the
+            `inputs_ids` passed when calling [`MistralModel`]
+        hidden_size (`int`, *optional*, defaults to 4096):
+            Dimension of the hidden representations.
+        intermediate_size (`int`, *optional*, defaults to 14336):
+            Dimension of the MLP representations.
+        num_hidden_layers (`int`, *optional*, defaults to 32):
+            Number of hidden layers in the Transformer encoder.
+        num_attention_heads (`int`, *optional*, defaults to 32):
+            Number of attention heads for each attention layer in the Transformer encoder.
+        num_key_value_heads (`int`, *optional*, defaults to 8):
+            This is the number of key_value heads that should be used to implement Grouped Query Attention. If
+            `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
+            `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When
+            converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
+            by meanpooling all the original heads within that group. For more details checkout [this
+            paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to `8`.
+        hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
+            The non-linear activation function (function or string) in the decoder.
+        max_position_embeddings (`int`, *optional*, defaults to `4096*32`):
+            The maximum sequence length that this model might ever be used with. Mistral's sliding window attention
+            allows sequence of up to 4096*32 tokens.
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        rms_norm_eps (`float`, *optional*, defaults to 1e-06):
+            The epsilon used by the rms normalization layers.
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Whether or not the model should return the last key/values attentions (not used by all models). Only
+            relevant if `config.is_decoder=True`.
+        pad_token_id (`int`, *optional*):
+            The id of the padding token.
+        bos_token_id (`int`, *optional*, defaults to 1):
+            The id of the "beginning-of-sequence" token.
+        eos_token_id (`int`, *optional*, defaults to 2):
+            The id of the "end-of-sequence" token.
+        tie_word_embeddings (`bool`, *optional*, defaults to `False`):
+            Whether the model's input and output word embeddings should be tied.
+        rope_theta (`float`, *optional*, defaults to 10000.0):
+            The base period of the RoPE embeddings.
+        sliding_window (`int`, *optional*, defaults to 4096):
+            Sliding window attention window size. If not specified, will default to `4096`.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+
+    ```python
+    >>> from transformers import MistralModel, MistralConfig
+
+    >>> # Initializing a Mistral 7B style configuration
+    >>> configuration = MistralConfig()
+
+    >>> # Initializing a model from the Mistral 7B style configuration
+    >>> model = MistralModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "mistral"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        vocab_size=32000,
+        hidden_size=4096,
+        intermediate_size=14336,
+        num_hidden_layers=32,
+        num_attention_heads=32,
+        num_key_value_heads=8,
+        hidden_act="silu",
+        max_position_embeddings=4096 * 32,
+        initializer_range=0.02,
+        rms_norm_eps=1e-6,
+        use_cache=True,
+        pad_token_id=None,
+        bos_token_id=1,
+        eos_token_id=2,
+        tie_word_embeddings=False,
+        rope_theta=10000.0,
+        sliding_window=4096,
+        rope_scaling=None,
+        attention_dropout=0.0,
+        beacon_window=1024,
+        beacon_stride=1024,
+        beacon_attn="full-coverage",
+        beacon_ratio=[2,4,8,16,32],
+        beacon_ratio_mix="step-random",
+        beacon_param=[],
+        beacon_embed_init="eos",
+        beacon_sink_size=0,
+        beacon_attend_prev=True,
+        beacon_pos="interleave",
+        beacon_parallel_window=1,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.sliding_window = sliding_window
+
+        # for backward compatibility
+        if num_key_value_heads is None:
+            num_key_value_heads = num_attention_heads
+
+        self.num_key_value_heads = num_key_value_heads
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.rope_theta = rope_theta
+        self.attention_dropout = attention_dropout
+
+        self.rope_scaling = rope_scaling
+        self._rope_scaling_validation()
+
+        self.beacon_window = beacon_window
+        self.beacon_stride = beacon_stride
+        self.beacon_attn = beacon_attn
+        self.beacon_ratio = beacon_ratio
+        self.beacon_ratio_mix = beacon_ratio_mix
+        self.beacon_param = beacon_param
+        self.beacon_embed_init = beacon_embed_init
+        self.beacon_sink_size = beacon_sink_size
+        self.beacon_attend_prev = beacon_attend_prev
+        self.beacon_pos = beacon_pos
+        self.beacon_parallel_window = beacon_parallel_window
+
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+    def _rope_scaling_validation(self):
+        """
+        Validate the `rope_scaling` configuration.
+        """
+        if self.rope_scaling is None:
+            return
+
+        if not isinstance(self.rope_scaling, dict) or len(self.rope_scaling) != 2:
+            raise ValueError(
+                "`rope_scaling` must be a dictionary with with two fields, `type` and `factor`, "
+                f"got {self.rope_scaling}"
+            )
+        rope_scaling_type = self.rope_scaling.get("type", None)
+        rope_scaling_factor = self.rope_scaling.get("factor", None)
+        if rope_scaling_type is None or rope_scaling_type not in ["linear", "dynamic"]:
+            raise ValueError(
+                f"`rope_scaling`'s type field must be one of ['linear', 'dynamic'], got {rope_scaling_type}"
+            )
+        if rope_scaling_factor is None or not isinstance(rope_scaling_factor, float) or rope_scaling_factor <= 1.0:
+            raise ValueError(f"`rope_scaling`'s factor field must be a float > 1, got {rope_scaling_factor}")
+

Long_LLM/activation_beacon/new/src/mistral/modeling_mistral.py

+# coding=utf-8
+# Copyright 2024 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# 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.
+""" PyTorch Mistral model."""
+import inspect
+import math
+import warnings
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from torch import nn
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask, _prepare_4d_causal_attention_mask_for_sdpa
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, SequenceClassifierOutputWithPast
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import (
+    add_start_docstrings,
+    add_start_docstrings_to_model_forward,
+    is_flash_attn_2_available,
+    is_flash_attn_greater_or_equal_2_10,
+    logging,
+    replace_return_docstrings,
+)
+from transformers.integrations import is_deepspeed_zero3_enabled
+from .configuration_mistral import MistralConfig
+
+
+if is_flash_attn_2_available():
+    from flash_attn import flash_attn_func, flash_attn_varlen_func
+    from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input  # noqa
+
+    _flash_supports_window_size = "window_size" in list(inspect.signature(flash_attn_func).parameters)
+
+from ..modeling_beacon import Memory
+from ..modeling_utils import optional_grad_ctx, compute_loss, BeaconModelOutput
+
+
+logger = logging.get_logger(__name__)
+
+_CONFIG_FOR_DOC = "MistralConfig"
+
+# Copied from transformers.models.llama.modeling_llama._get_unpad_data
+def _get_unpad_data(attention_mask):
+    seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
+    indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
+    max_seqlen_in_batch = seqlens_in_batch.max().item()
+    cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
+    return (
+        indices,
+        cu_seqlens,
+        max_seqlen_in_batch,
+    )
+
+
+# Copied from transformers.models.llama.modeling_llama.MistralRMSNorm with Mistral->Mistral
+class MistralRMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        """
+        MistralRMSNorm is equivalent to T5LayerNorm
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+
+# Copied from transformers.models.llama.modeling_llama.rotate_half
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+class MistralRotaryEmbedding(nn.Module):
+    def __init__(self, dim, max_position_embeddings=32768, base=10000, device=None):
+        super().__init__()
+
+        self.dim = dim
+        self.max_position_embeddings = max_position_embeddings
+        self.base = base
+        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+        # Build here to make `torch.jit.trace` work.
+        self._set_cos_sin_cache(
+            seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
+        )
+
+    def _set_cos_sin_cache(self, seq_len, device, dtype):
+        self.max_seq_len_cached = seq_len
+        t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.int64).type_as(self.inv_freq)
+
+        freqs = torch.outer(t, self.inv_freq)
+        # Different from paper, but it uses a different permutation in order to obtain the same calculation
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
+        self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
+
+    def forward(self, q, k, position_ids):
+        seq_len = max(position_ids.max().item() + 1, k.shape[2])
+
+        # x: [bs, num_attention_heads, seq_len, head_size]
+        if seq_len > self.max_seq_len_cached:
+            self._set_cos_sin_cache(seq_len=seq_len, device=k.device, dtype=k.dtype)
+
+        # batch_size, 1, key_len, head_dim
+        k_cos = self.cos_cached[position_ids].to(dtype=k.dtype).unsqueeze(1)
+        k_sin = self.sin_cached[position_ids].to(dtype=k.dtype).unsqueeze(1)
+
+        q_cos = k_cos[..., -q.shape[2]:, :]
+        q_sin = k_sin[..., -q.shape[2]:, :]
+
+        q_embed = (q * q_cos) + (rotate_half(q) * q_sin)
+        k_embed = (k * k_cos) + (rotate_half(k) * k_sin)
+        return q_embed, k_embed
+
+
+class MistralLinearScalingRotaryEmbedding(MistralRotaryEmbedding):
+    """MistralRotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev"""
+
+    def __init__(self, dim, max_position_embeddings=32768, base=10000, device=None, scaling_factor=1.0):
+        self.scaling_factor = scaling_factor
+        super().__init__(dim, max_position_embeddings, base, device)
+
+    def _set_cos_sin_cache(self, seq_len, device, dtype):
+        self.max_seq_len_cached = seq_len
+        t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
+        t = t / self.scaling_factor
+
+        freqs = torch.outer(t, self.inv_freq)
+        # Different from paper, but it uses a different permutation in order to obtain the same calculation
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
+        self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
+
+
+class MistralDynamicNTKScalingRotaryEmbedding(MistralRotaryEmbedding):
+    """MistralRotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla"""
+
+    def __init__(self, dim, max_position_embeddings=32768, base=10000, device=None, scaling_factor=1.0):
+        self.scaling_factor = scaling_factor
+        super().__init__(dim, max_position_embeddings, base, device)
+
+    def _set_cos_sin_cache(self, seq_len, device, dtype):
+        self.max_seq_len_cached = seq_len
+
+        if seq_len > self.max_position_embeddings:
+            base = self.base * (
+                (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1)
+            ) ** (self.dim / (self.dim - 2))
+            inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
+            self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+        t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
+
+        freqs = torch.outer(t, self.inv_freq)
+        # Different from paper, but it uses a different permutation in order to obtain the same calculation
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
+        self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
+
+
+class MistralYarnRotaryEmbedding(nn.Module):
+    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0, beta_slow=2, beta_fast=128):
+        super().__init__()
+
+        self.base = base
+        self.dim = dim
+        self.scaling_factor = scaling_factor
+        self.beta_slow = beta_slow
+        self.beta_fast = beta_fast
+        self.max_position_embeddings = max_position_embeddings
+
+        self._set_cos_sin_cache(
+            seq_len=max_position_embeddings, device=device, dtype=torch.get_default_dtype()
+        )
+
+    def _get_factor(self, device, dtype):
+        # the dimension whose index is smaller than fast_dim rotates more than beta_fast
+        fast_dim = self.dim / 2 * (math.log(self.max_position_embeddings / (2 * math.pi * self.beta_fast)) / math.log(self.base))
+        fast_dim = max(math.floor(fast_dim), 0)
+        # the dimension whose index is bigger than slow_dim rotates less than beta_slow
+        slow_dim = self.dim / 2 * (math.log(self.max_position_embeddings / (2 * math.pi * self.beta_slow)) / math.log(self.base))
+        slow_dim = min(math.ceil(slow_dim), self.dim - 1)
+
+        if fast_dim == slow_dim:
+            slow_dim += 0.001
+
+        # NOTE: very important to use full precision here so that the factor is correct
+        dim_arange = torch.arange(0, self.dim // 2, device=device, dtype=torch.float32)
+        dim_factor = (dim_arange - fast_dim) / (slow_dim - fast_dim)
+        dim_factor = torch.clamp(dim_factor, 0, 1)
+
+        # align with the paper notation
+        return (1 - dim_factor)
+
+    def _get_temperature(self):
+        if self.scaling_factor <= 1:
+            return 1.0
+        return 0.07 * math.log(self.scaling_factor) + 1.0
+    
+    def _set_cos_sin_cache(self, seq_len, device, dtype):
+        dim_arange = torch.arange(0, self.dim, 2, device=device) / self.dim
+        # dim / 2
+        freq = self.base ** dim_arange
+        theta = 1 / freq
+        interleave_theta = theta / self.scaling_factor
+
+        factor = self._get_factor(device, dtype)
+        yarn_theta = factor * theta + (1 - factor) * interleave_theta
+        self.register_buffer("inv_freq", yarn_theta, persistent=False)
+
+        # print(factor, self.inv_freq)
+
+        t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype)
+        freqs = torch.outer(t, self.inv_freq)
+        emb = torch.cat((freqs, freqs), dim=-1)
+
+        # get attention temperature
+        temperature = self._get_temperature()
+
+        self.register_buffer("cos_cached", (emb.cos() * temperature).to(dtype), persistent=False)
+        self.register_buffer("sin_cached", (emb.sin() * temperature).to(dtype), persistent=False)
+        self.max_seq_len_cached = seq_len
+    
+    def forward(self, q, k, position_ids):
+        seq_len = max(position_ids.max().item() + 1, k.shape[2])
+
+        # x: [bs, num_attention_heads, seq_len, head_size]
+        if seq_len > self.max_seq_len_cached:
+            self.scaling_factor = seq_len / self.max_position_embeddings
+            self._set_cos_sin_cache(seq_len=seq_len, device=k.device, dtype=k.dtype)
+
+        k_cos = self.cos_cached[position_ids].to(dtype=k.dtype).unsqueeze(1)
+        k_sin = self.sin_cached[position_ids].to(dtype=k.dtype).unsqueeze(1)
+
+        q_cos = k_cos[..., -q.shape[2]:, :]
+        q_sin = k_sin[..., -q.shape[2]:, :]
+
+        q_embed = (q * q_cos) + (rotate_half(q) * q_sin)
+        k_embed = (k * k_cos) + (rotate_half(k) * k_sin)
+        return q_embed, k_embed
+
+
+# Copied from transformers.models.mistral.modeling_mistral.MistralMLP with Mistral->Mistral
+class MistralMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        self.act_fn = ACT2FN[config.hidden_act]
+
+        if "mlp" in config.beacon_param:            
+            self.beacon_up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+            self.beacon_up_proj.weight.data.zero_()
+            self.beacon_up_proj._is_hf_initialized = True
+
+            self.beacon_down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+            self.beacon_down_proj.weight.data.zero_()
+            self.beacon_down_proj._is_hf_initialized = True
+    
+    def _init_beacon_proj(self, missing_keys):
+        """Initialize the beacon projection weight with that of the ordinal projection."""
+        if "mlp" in self.config.beacon_param:
+            if is_deepspeed_zero3_enabled():
+                # FIXME: after deepspeed initialization, some weights becomes non-zero
+                # For Mistral, there are rows that are full of zeros
+                # For Mistral, there are values bigger than 1e29...
+
+                import deepspeed
+                params = [self.up_proj.weight, self.down_proj.weight, self.beacon_up_proj.weight, self.beacon_down_proj.weight]
+                with deepspeed.zero.GatheredParameters(params, modifier_rank=0):
+                    if (self.beacon_up_proj.weight.sum(-1) == 0).any() or (self.beacon_up_proj.weight > 1e29).any():
+                        self.beacon_up_proj.weight.data[:] = self.up_proj.weight.data
+                        self.beacon_down_proj.weight.data[:] = self.down_proj.weight.data
+            else:
+                if any("beacon_up_proj" in missing_key for missing_key in missing_keys):
+                    # only copy the value in-place, without tieing the weight
+                    self.beacon_up_proj.weight.data[:] = self.up_proj.weight.data
+                    self.beacon_down_proj.weight.data[:] = self.down_proj.weight.data
+
+    def forward(self, x, beacon_size, beacon_indices):
+        if "mlp" in self.config.beacon_param:
+            # NOTE: when beacon_pos == "interleave", the beacon_indices points to all beacon tokens in the current window (cached activations + input_ids), so we shall slice out the part corresponding to the input_ids
+            if beacon_size > 0:
+                cur_beacon_indices = beacon_indices[-x.shape[1]:]
+                ordinal_hidden_states = x[:, cur_beacon_indices == 0]
+                beacon_hidden_states = x[:, cur_beacon_indices == 1]
+
+                ordinal_down_proj = self.down_proj(self.act_fn(self.gate_proj(ordinal_hidden_states)) * self.up_proj(ordinal_hidden_states))
+                beacon_down_proj = self.beacon_down_proj(self.act_fn(self.gate_proj(beacon_hidden_states)) * self.beacon_up_proj(beacon_hidden_states))
+
+                down_proj = beacon_down_proj.new_ones(x.shape)
+                down_proj[:, beacon_indices == 0] = ordinal_down_proj
+                down_proj[:, beacon_indices == 1] = beacon_down_proj
+            else:
+                down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+        else:
+            down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+
+        return down_proj
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+class MistralAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config: MistralConfig, layer_idx: Optional[int] = None):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        if layer_idx is None:
+            logger.warning_once(
+                f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will "
+                "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "
+                "when creating this class."
+            )
+
+        self.attention_dropout = config.attention_dropout
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = self.hidden_size // self.num_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.max_position_embeddings = config.max_position_embeddings
+        self.rope_theta = config.rope_theta
+        self.is_causal = True
+
+        if (self.head_dim * self.num_heads) != self.hidden_size:
+            raise ValueError(
+                f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
+                f" and `num_heads`: {self.num_heads})."
+            )
+        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
+
+        self._init_rope()
+
+        # NOTE: add extra parameters for beacon tokens
+        # skip post initialization to speed up loading
+        if "q" in config.beacon_param:
+            self.beacon_q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=self.q_proj.bias is not None)
+            # NOTE: initialize the beacon parameters as zero
+            self.beacon_q_proj.weight.data.zero_()
+            self.beacon_q_proj._is_hf_initialized = True
+        if "k" in config.beacon_param:
+            self.beacon_k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=self.k_proj.bias is not None)
+            self.beacon_k_proj.weight.data.zero_()
+            self.beacon_k_proj._is_hf_initialized = True
+        if "v" in config.beacon_param:
+            self.beacon_v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=self.v_proj.bias is not None)
+            self.beacon_v_proj.weight.data.zero_()
+            self.beacon_v_proj._is_hf_initialized = True
+        if "o" in config.beacon_param:
+            self.beacon_o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=self.o_proj.bias is not None)
+            self.beacon_o_proj.weight.data.zero_()
+            self.beacon_o_proj._is_hf_initialized = True
+
+    def _init_rope(self):
+        if self.config.rope_scaling is None:
+            self.rotary_emb = MistralRotaryEmbedding(
+                self.head_dim,
+                max_position_embeddings=self.max_position_embeddings,
+                base=self.rope_theta,
+            )
+        else:
+            scaling_type = self.config.rope_scaling["type"]
+            scaling_factor = self.config.rope_scaling["factor"]
+            if scaling_type == "linear":
+                self.rotary_emb = MistralLinearScalingRotaryEmbedding(
+                    self.head_dim,
+                    max_position_embeddings=self.max_position_embeddings,
+                    scaling_factor=scaling_factor,
+                    base=self.rope_theta,
+                )
+            elif scaling_type == "dynamic":
+                self.rotary_emb = MistralDynamicNTKScalingRotaryEmbedding(
+                    self.head_dim,
+                    max_position_embeddings=self.max_position_embeddings,
+                    scaling_factor=scaling_factor,
+                    base=self.rope_theta,
+                )
+            elif scaling_type == "yarn":
+                self.rotary_emb = MistralYarnRotaryEmbedding(
+                    self.head_dim,
+                    max_position_embeddings=self.max_position_embeddings,
+                    scaling_factor=scaling_factor,
+                    base=self.rope_theta,
+                )
+            else:
+                raise ValueError(f"Unknown RoPE scaling type {scaling_type}")
+
+    def _init_beacon_proj(self, missing_keys):
+        """Initialize the beacon projection weight with that of the ordinal projection."""
+        beacon_param = self.config.beacon_param
+        
+        if is_deepspeed_zero3_enabled():
+            # FIXME: after deepspeed initialization, some weights becomes non-zero
+            # For Mistral, there are rows that are full of zeros
+            # For Mistral, there are values bigger than 1e29...
+
+            import deepspeed
+            if "q" in beacon_param:
+                params = [self.beacon_q_proj.weight, self.q_proj.weight]
+                if self.q_proj.bias is not None:
+                    params.extend([self.beacon_q_proj.bias, self.q_proj.bias])
+                with deepspeed.zero.GatheredParameters(params, modifier_rank=0):
+                    # FIXME: after deepspeed initialization, some weights becomes non-zero, but there are rows that are full of zeros
+                    if (self.beacon_q_proj.weight.sum(-1) == 0).any() or (self.beacon_q_proj.weight > 1e29).any():
+                        self.beacon_q_proj.weight.data[:] = self.q_proj.weight.data
+                        if self.q_proj.bias is not None:
+                            self.beacon_q_proj.bias.data[:] = self.q_proj.bias.data
+            if "k" in beacon_param:
+                params = [self.beacon_k_proj.weight, self.k_proj.weight]
+                if self.k_proj.bias is not None:
+                    params.extend([self.beacon_k_proj.bias, self.k_proj.bias])
+                with deepspeed.zero.GatheredParameters(params, modifier_rank=0):
+                    # FIXME: after deepspeed initialization, some weights becomes non-zero, but there are rows that are full of zeros
+                    if (self.beacon_k_proj.weight.sum(-1) == 0).any() or (self.beacon_k_proj.weight > 1e29).any():
+                        self.beacon_k_proj.weight.data[:] = self.k_proj.weight.data
+                        if self.k_proj.bias is not None:
+                            self.beacon_k_proj.bias.data[:] = self.k_proj.bias.data
+            if "v" in beacon_param:
+                params = [self.beacon_v_proj.weight, self.v_proj.weight]
+                if self.v_proj.bias is not None:
+                    params.extend([self.beacon_v_proj.bias, self.v_proj.bias])
+                with deepspeed.zero.GatheredParameters(params, modifier_rank=0):
+                    # FIXME: after deepspeed initialization, some weights becomes non-zero, but there are rows that are full of zeros
+                    if (self.beacon_v_proj.weight.sum(-1) == 0).any() or (self.beacon_v_proj.weight > 1e29).any():
+                        self.beacon_v_proj.weight.data[:] = self.v_proj.weight.data
+                        if self.v_proj.bias is not None:
+                            self.beacon_v_proj.bias.data[:] = self.v_proj.bias.data
+            if "o" in beacon_param:
+                params = [self.beacon_o_proj.weight, self.o_proj.weight]
+                if self.o_proj.bias is not None:
+                    params.extend([self.beacon_o_proj.bias, self.o_proj.bias])
+                with deepspeed.zero.GatheredParameters(params, modifier_rank=0):
+                    # FIXME: after deepspeed initialization, some weights becomes non-zero, but there are rows that are full of zeros
+                    if (self.beacon_o_proj.weight.sum(-1) == 0).any() or (self.beacon_o_proj.weight > 1e29).any():
+                        self.beacon_o_proj.weight.data[:] = self.o_proj.weight.data
+                        if self.o_proj.bias is not None:
+                            self.beacon_o_proj.bias.data[:] = self.o_proj.bias.data
+        else:
+            # only copy the value in-place, without tieing the weight
+            if "q" in beacon_param and any("beacon_q_proj" in missing_key for missing_key in missing_keys):
+                # FIXME: some beacon weights are not initialized as zero for mistral model, why? 
+                # if (self.beacon_q_proj.weight == 0).all():
+                    self.beacon_q_proj.weight.data[:] = self.q_proj.weight.data
+                    if self.q_proj.bias is not None:
+                        self.beacon_q_proj.bias.data[:] = self.q_proj.bias.data
+            if "k" in beacon_param and any("beacon_k_proj" in missing_key for missing_key in missing_keys):
+                # if (self.beacon_k_proj.weight == 0).all():
+                    self.beacon_k_proj.weight.data[:] = self.k_proj.weight.data
+                    if self.k_proj.bias is not None:
+                        self.beacon_k_proj.bias.data[:] = self.k_proj.bias.data
+            if "v" in beacon_param and any("beacon_v_proj" in missing_key for missing_key in missing_keys):
+                # if (self.beacon_v_proj.weight == 0).all():
+                    self.beacon_v_proj.weight.data[:] = self.v_proj.weight.data
+                    if self.v_proj.bias is not None:
+                        self.beacon_v_proj.bias.data[:] = self.v_proj.bias.data
+            if "o" in beacon_param and any("beacon_o_proj" in missing_key for missing_key in missing_keys):
+                # if (self.beacon_o_proj.weight == 0).all():
+                    self.beacon_o_proj.weight.data[:] = self.o_proj.weight.data
+                    if self.o_proj.bias is not None:
+                        self.beacon_o_proj.bias.data[:] = self.o_proj.bias.data
+
+    def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
+        return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
+    
+    def qkv_proj_with_beacon(self, hidden_states, beacon_size, beacon_indices):
+        if beacon_size > 0:
+            # NOTE: when beacon_pos == "interleave", the beacon_indices points to all beacon tokens in the current window (cached activations + input_ids), so we shall slice out the part corresponding to the input_ids
+            cur_beacon_indices = beacon_indices[-hidden_states.shape[1]:]
+
+            ordinal_hidden_states = hidden_states[:, cur_beacon_indices == 0]
+            beacon_hidden_states = hidden_states[:, cur_beacon_indices == 1]
+
+            if "q" in self.config.beacon_param:
+                ordinal_query_states = self.q_proj(ordinal_hidden_states)
+                beacon_query_states = self.beacon_q_proj(beacon_hidden_states)
+                query_states = beacon_query_states.new_zeros((ordinal_query_states.shape[0], cur_beacon_indices.shape[0], ordinal_query_states.shape[2]))
+                query_states[:, cur_beacon_indices == 0] = ordinal_query_states
+                query_states[:, cur_beacon_indices == 1] = beacon_query_states
+                # NOTE: replicate hidden states for beacon tokens in case of parallel windows
+                if (cur_beacon_indices == 2).any():
+                    query_states[:, cur_beacon_indices == 2] = beacon_query_states[:, :(cur_beacon_indices == 2).sum()]
+
+            else:
+                query_states = self.q_proj(hidden_states)
+
+            if "k" in self.config.beacon_param:
+                ordinal_key_states = self.k_proj(ordinal_hidden_states)
+                beacon_key_states = self.beacon_k_proj(beacon_hidden_states)
+                key_states = beacon_key_states.new_zeros((ordinal_key_states.shape[0], cur_beacon_indices.shape[0], ordinal_key_states.shape[2]))
+                key_states[:, cur_beacon_indices == 0] = ordinal_key_states
+                key_states[:, cur_beacon_indices == 1] = beacon_key_states
+                # NOTE: replicate hidden states for beacon tokens in case of parallel windows
+                if (cur_beacon_indices == 2).any():
+                    key_states[:, cur_beacon_indices == 2] = beacon_key_states[:, :(cur_beacon_indices == 2).sum()]
+
+            else:
+                key_states = self.k_proj(hidden_states)
+            
+            if "v" in self.config.beacon_param:
+                ordinal_value_states = self.v_proj(ordinal_hidden_states)
+                beacon_value_states = self.beacon_v_proj(beacon_hidden_states)
+                value_states = beacon_value_states.new_zeros((ordinal_value_states.shape[0], cur_beacon_indices.shape[0], ordinal_value_states.shape[2]))
+                value_states[:, cur_beacon_indices == 0] = ordinal_value_states
+                value_states[:, cur_beacon_indices == 1] = beacon_value_states
+                # NOTE: replicate hidden states for beacon tokens in case of parallel windows
+                if (cur_beacon_indices == 2).any():
+                    value_states[:, cur_beacon_indices == 2] = beacon_value_states[:, :(cur_beacon_indices == 2).sum()]
+            else:
+                value_states = self.v_proj(hidden_states)
+
+        else:
+            query_states = self.q_proj(hidden_states)
+            key_states = self.k_proj(hidden_states)
+            value_states = self.v_proj(hidden_states)
+
+        return query_states, key_states, value_states
+    
+    def o_proj_with_beacon(self, attn_output, beacon_size, beacon_indices):
+        if beacon_size > 0:
+            # NOTE: when beacon_pos == "interleave", the beacon_indices points to all beacon tokens in the current window (cached activations + input_ids), so we shall slice out the part corresponding to the input_ids
+            cur_beacon_indices = beacon_indices[-attn_output.shape[1]:]
+
+            if "o" in self.config.beacon_param:
+                ordinal_attn_output = self.o_proj(attn_output[:, cur_beacon_indices == 0])
+                beacon_attn_output = self.beacon_o_proj(attn_output[:, cur_beacon_indices == 1])
+                attn_output = beacon_attn_output.new_zeros(attn_output.shape)
+                attn_output[:, cur_beacon_indices == 0] = ordinal_attn_output
+                attn_output[:, cur_beacon_indices == 1] = beacon_attn_output
+                # NOTE: replicate hidden states for beacon tokens in case of parallel windows
+                # if (cur_beacon_indices == 2).any():
+                #     attn_output[:, cur_beacon_indices == 2] = beacon_attn_output[:, :(cur_beacon_indices == 2).sum()]
+            else:
+                attn_output = self.o_proj(attn_output)
+        else:
+            attn_output = self.o_proj(attn_output)
+        return attn_output
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Cache] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        if "padding_mask" in kwargs:
+            warnings.warn(
+                "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
+            )
+
+        bsz, q_len, _ = hidden_states.size()
+        kv_seq_len = hidden_states.shape[-2]
+        past_key, past_value, beacon_size, beacon_indices = past_key_value
+
+        if past_key is not None:
+            past_seq_len = past_key.shape[2]
+            kv_seq_len += past_seq_len
+        else:
+            past_seq_len = 0
+
+        query_states, key_states, value_states = self.qkv_proj_with_beacon(hidden_states, beacon_size, beacon_indices)
+
+        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+        # return keys and values before rope
+        # NOTE: incrementally return keys and values for efficiency 
+        past_key_value = (key_states, value_states, beacon_size, beacon_indices)
+
+        if past_key is not None:
+            # reuse k, v, self_attention
+            key_states = torch.cat([past_key, key_states], dim=2)
+            value_states = torch.cat([past_value, value_states], dim=2)
+
+        query_states, key_states = self.rotary_emb(query_states, key_states, position_ids)
+
+        key_states = repeat_kv(key_states, self.num_key_value_groups)
+        value_states = repeat_kv(value_states, self.num_key_value_groups)
+
+        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
+
+        if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
+            raise ValueError(
+                f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"
+                f" {attn_weights.size()}"
+            )
+
+        if attention_mask is not None:
+            if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
+                raise ValueError(
+                    f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
+                )
+            attn_weights = attn_weights + attention_mask
+
+        # upcast attention to fp32
+        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
+        attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)
+        attn_output = torch.matmul(attn_weights, value_states)
+
+        if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
+            raise ValueError(
+                f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
+                f" {attn_output.size()}"
+            )
+
+        attn_output = attn_output.transpose(1, 2).contiguous()
+
+        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
+
+        attn_output = self.o_proj_with_beacon(attn_output, beacon_size, beacon_indices)
+
+        if not output_attentions:
+            attn_weights = None
+
+        return attn_output, attn_weights, past_key_value
+
+
+class MistralSdpaAttention(MistralAttention):
+    """
+    Mistral attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from
+    `MistralAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to
+    SDPA API.
+    """
+
+    # Adapted from MistralAttention.forward
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Cache] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        if output_attentions:
+            # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented.
+            logger.warning_once(
+                "MistralModel is using MistralSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, "
+                'but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
+            )
+            return super().forward(
+                hidden_states=hidden_states,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_value=past_key_value,
+                output_attentions=output_attentions,
+                use_cache=use_cache,
+            )
+        bsz, q_len, _ = hidden_states.size()
+        kv_seq_len = hidden_states.shape[-2]
+        past_key, past_value, beacon_size, beacon_indices = past_key_value
+        if past_key is not None:
+            past_seq_len = past_key.shape[2]
+            kv_seq_len += past_seq_len
+        else:
+            past_seq_len = 0
+
+        query_states, key_states, value_states = self.qkv_proj_with_beacon(hidden_states, beacon_size, beacon_indices)
+
+        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+        # return keys and values before rope
+        # NOTE: incrementally return keys and values for efficiency 
+        past_key_value = (key_states, value_states, beacon_size, beacon_indices)
+
+        if past_key is not None:
+            # reuse k, v, self_attention
+            key_states = torch.cat([past_key, key_states], dim=2)
+            value_states = torch.cat([past_value, value_states], dim=2)
+
+        query_states, key_states = self.rotary_emb(query_states, key_states, position_ids)
+
+        key_states = repeat_kv(key_states, self.num_key_value_groups)
+        value_states = repeat_kv(value_states, self.num_key_value_groups)
+
+        if attention_mask is not None:
+            if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
+                raise ValueError(
+                    f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
+                )
+
+        # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,
+        # Reference: https://github.com/pytorch/pytorch/issues/112577.
+        if query_states.device.type == "cuda" and attention_mask is not None:
+            query_states = query_states.contiguous()
+            key_states = key_states.contiguous()
+            value_states = value_states.contiguous()
+
+        attn_output = torch.nn.functional.scaled_dot_product_attention(
+            query_states,
+            key_states,
+            value_states,
+            attn_mask=attention_mask,
+            dropout_p=self.attention_dropout if self.training else 0.0,
+            # The q_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case q_len == 1.
+            is_causal=self.is_causal and attention_mask is None and q_len > 1,
+        )
+
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
+        attn_output = self.o_proj_with_beacon(attn_output, beacon_size, beacon_indices)
+
+        return attn_output, None, past_key_value
+
+
+class MistralFlashAttention2(MistralAttention):
+    """
+    Mistral flash attention module. This module inherits from `MistralAttention` as the weights of the module stays
+    untouched. The only required change would be on the forward pass where it needs to correctly call the public API of
+    flash attention and deal with padding tokens in case the input contains any of them.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
+        # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.
+        # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left).
+        self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Cache] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        output_attentions = False
+
+        bsz, q_len, _ = hidden_states.size()
+        kv_seq_len = hidden_states.shape[-2]
+
+        past_key, past_value, beacon_size, beacon_indices = past_key_value
+        if past_key is not None:
+            past_seq_len = past_key.shape[2]
+            kv_seq_len += past_seq_len
+        else:
+            past_seq_len = 0
+
+        query_states, key_states, value_states = self.qkv_proj_with_beacon(hidden_states, beacon_size, beacon_indices)
+
+        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+        # return keys and values before rope
+        # NOTE: incrementally return keys and values for efficiency 
+        past_key_value = (key_states, value_states, beacon_size, beacon_indices)
+
+        if past_key is not None:
+            # reuse k, v, self_attention
+            key_states = torch.cat([past_key, key_states], dim=2)
+            value_states = torch.cat([past_value, value_states], dim=2)
+
+        query_states, key_states = self.rotary_emb(query_states, key_states, position_ids)
+
+        key_states = repeat_kv(key_states, self.num_key_value_groups)
+        value_states = repeat_kv(value_states, self.num_key_value_groups)
+        
+        # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache
+        # to be able to avoid many of these transpose/reshape/view.
+        query_states = query_states.transpose(1, 2)
+        key_states = key_states.transpose(1, 2)
+        value_states = value_states.transpose(1, 2)
+
+        dropout_rate = self.attention_dropout if self.training else 0.0
+
+        # In PEFT, usually we cast the layer norms in float32 for training stability reasons
+        # therefore the input hidden states gets silently casted in float32. Hence, we need
+        # cast them back in the correct dtype just to be sure everything works as expected.
+        # This might slowdown training & inference so it is recommended to not cast the LayerNorms
+        # in fp32. (MistralRMSNorm handles it correctly)
+
+        input_dtype = query_states.dtype
+        if input_dtype == torch.float32:
+            if torch.is_autocast_enabled():
+                target_dtype = torch.get_autocast_gpu_dtype()
+            # Handle the case where the model is quantized
+            elif hasattr(self.config, "_pre_quantization_dtype"):
+                target_dtype = self.config._pre_quantization_dtype
+            else:
+                target_dtype = self.q_proj.weight.dtype
+
+            logger.warning_once(
+                f"The input hidden states seems to be silently casted in float32, this might be related to"
+                f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
+                f" {target_dtype}."
+            )
+
+            query_states = query_states.to(target_dtype)
+            key_states = key_states.to(target_dtype)
+            value_states = value_states.to(target_dtype)
+
+        attn_output = self._flash_attention_forward(
+            query_states, 
+            key_states, 
+            value_states, 
+            attention_mask, 
+            q_len, 
+            dropout=dropout_rate
+        )
+
+        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()
+        attn_output = self.o_proj_with_beacon(attn_output, beacon_size, beacon_indices)
+
+        if not output_attentions:
+            attn_weights = None
+
+        return attn_output, attn_weights, past_key_value
+
+    def _flash_attention_forward(
+        self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None
+    ):
+        """
+        Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token
+        first unpad the input, then computes the attention scores and pad the final attention scores.
+
+        Args:
+            query_states (`torch.Tensor`):
+                Input query states to be passed to Flash Attention API
+            key_states (`torch.Tensor`):
+                Input key states to be passed to Flash Attention API
+            value_states (`torch.Tensor`):
+                Input value states to be passed to Flash Attention API
+            attention_mask (`torch.Tensor`):
+                The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
+                position of padding tokens and 1 for the position of non-padding tokens.
+            dropout (`float`):
+                Attention dropout
+            softmax_scale (`float`, *optional*):
+                The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
+        """
+        if not self._flash_attn_uses_top_left_mask:
+            causal = self.is_causal
+        else:
+            # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in MistralFlashAttention2 __init__.
+            causal = self.is_causal and query_length != 1
+
+        # Contains at least one padding token in the sequence
+        if attention_mask is not None:
+            batch_size = query_states.shape[0]
+            query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
+                query_states, key_states, value_states, attention_mask, query_length
+            )
+
+            cu_seqlens_q, cu_seqlens_k = cu_seq_lens
+            max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
+
+            attn_output_unpad = flash_attn_varlen_func(
+                query_states,
+                key_states,
+                value_states,
+                cu_seqlens_q=cu_seqlens_q,
+                cu_seqlens_k=cu_seqlens_k,
+                max_seqlen_q=max_seqlen_in_batch_q,
+                max_seqlen_k=max_seqlen_in_batch_k,
+                dropout_p=dropout,
+                softmax_scale=softmax_scale,
+                causal=causal,
+            )
+
+            attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)
+        else:
+            attn_output = flash_attn_func(
+                query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=causal
+            )
+
+        return attn_output
+
+    def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
+        indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
+        batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
+
+        key_layer = index_first_axis(
+            key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
+        )
+        value_layer = index_first_axis(
+            value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
+        )
+        if query_length == kv_seq_len:
+            query_layer = index_first_axis(
+                query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim), indices_k
+            )
+            cu_seqlens_q = cu_seqlens_k
+            max_seqlen_in_batch_q = max_seqlen_in_batch_k
+            indices_q = indices_k
+        elif query_length == 1:
+            max_seqlen_in_batch_q = 1
+            cu_seqlens_q = torch.arange(
+                batch_size + 1, dtype=torch.int32, device=query_layer.device
+            )  # There is a memcpy here, that is very bad.
+            indices_q = cu_seqlens_q[:-1]
+            query_layer = query_layer.squeeze(1)
+        else:
+            # The -q_len: slice assumes left padding.
+            attention_mask = attention_mask[:, -query_length:]
+            query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)
+
+        return (
+            query_layer,
+            key_layer,
+            value_layer,
+            indices_q,
+            (cu_seqlens_q, cu_seqlens_k),
+            (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
+        )
+
+
+MISTRAL_ATTENTION_CLASSES = {
+    "eager": MistralAttention,
+    "sdpa": MistralSdpaAttention,
+    "flash_attention_2": MistralFlashAttention2,
+}
+
+
+class MistralDecoderLayer(nn.Module):
+    def __init__(self, config: MistralConfig, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+
+        if config.sliding_window is not None and config._attn_implementation != "flash_attention_2":
+            logger.warning_once(
+                f"Sliding Window Attention is enabled but not implemented for `{config._attn_implementation}`; "
+                "unexpected results may be encountered."
+            )
+        self.self_attn = MISTRAL_ATTENTION_CLASSES[config._attn_implementation](config, layer_idx)
+
+        self.mlp = MistralMLP(config)
+        self.input_layernorm = MistralRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = MistralRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        output_attentions: Optional[bool] = False,
+        use_cache: Optional[bool] = False,
+        **kwargs,
+    ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
+        """
+        Args:
+            hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
+            attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
+                `(batch, sequence_length)` where padding elements are indicated by 0.
+            output_attentions (`bool`, *optional*):
+                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+                returned tensors for more detail.
+            use_cache (`bool`, *optional*):
+                If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
+                (see `past_key_values`).
+            past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
+        """
+        if "padding_mask" in kwargs:
+            warnings.warn(
+                "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
+            )
+
+        # NOTE: get beacon_size in case the mlp is included in beacon_param
+        past_key, past_value, beacon_size, beacon_indices = past_key_value
+
+        residual = hidden_states
+
+        hidden_states = self.input_layernorm(hidden_states)
+
+        # Self Attention
+        hidden_states, self_attn_weights, present_key_value = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_value=past_key_value,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+        )
+        hidden_states = residual + hidden_states
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states, beacon_size, beacon_indices)
+        hidden_states = residual + hidden_states
+
+        outputs = (hidden_states,)
+
+        if output_attentions:
+            outputs += (self_attn_weights,)
+
+        if use_cache:
+            outputs += (present_key_value,)
+
+        return outputs
+
+
+MISTRAL_START_DOCSTRING = r"""
+    This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
+    library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
+    etc.)
+
+    This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
+    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
+    and behavior.
+
+    Parameters:
+        config ([`MistralConfig`]):
+            Model configuration class with all the parameters of the model. Initializing with a config file does not
+            load the weights associated with the model, only the configuration. Check out the
+            [`~PreTrainedModel.from_pretrained`] method to load the model weights.
+"""
+
+
+@add_start_docstrings(
+    "The bare Mistral Model outputting raw hidden-states without any specific head on top.",
+    MISTRAL_START_DOCSTRING,
+)
+class MistralPreTrainedModel(PreTrainedModel):
+    config_class = MistralConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["MistralDecoderLayer"]
+    _skip_keys_device_placement = "past_key_values"
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+    _supports_cache_class = True
+
+    def _init_weights(self, module):
+        std = self.config.initializer_range
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+
+MISTRAL_INPUTS_DOCSTRING = r"""
+    Args:
+        input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+            Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
+            it.
+
+            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+            [`PreTrainedTokenizer.__call__`] for details.
+
+            [What are input IDs?](../glossary#input-ids)
+        attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
+
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **masked**.
+
+            [What are attention masks?](../glossary#attention-mask)
+
+            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+            [`PreTrainedTokenizer.__call__`] for details.
+
+            If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see
+            `past_key_values`).
+
+            If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]
+            and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more
+            information on the default strategy.
+
+            - 1 indicates the head is **not masked**,
+            - 0 indicates the head is **masked**.
+        position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
+            config.n_positions - 1]`.
+
+            [What are position IDs?](../glossary#position-ids)
+        past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*):
+            Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
+            blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values`
+            returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`.
+
+            Two formats are allowed:
+            - a [`~cache_utils.Cache`] instance;
+            - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of
+            shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy
+            cache format.
+
+            The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the
+            legacy cache format will be returned.
+
+            If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't
+            have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids`
+            of shape `(batch_size, sequence_length)`.
+        inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
+            Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
+            is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
+            model's internal embedding lookup matrix.
+        use_cache (`bool`, *optional*):
+            If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
+            `past_key_values`).
+        output_attentions (`bool`, *optional*):
+            Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
+            tensors for more detail.
+        output_hidden_states (`bool`, *optional*):
+            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
+            more detail.
+        return_dict (`bool`, *optional*):
+            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+"""
+
+
+@add_start_docstrings(
+    "The bare Mistral Model outputting raw hidden-states without any specific head on top.",
+    MISTRAL_START_DOCSTRING,
+)
+class MistralModel(MistralPreTrainedModel):
+    """
+    Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`MistralDecoderLayer`]
+
+    Args:
+        config: MistralConfig
+    """
+
+    def __init__(self, config: MistralConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+
+        # BEACON: add beacon embedding
+        self.beacon_embed_tokens = nn.Embedding(1, config.hidden_size, self.padding_idx)
+        self.beacon_embed_tokens._is_hf_initialized = True
+
+        self.layers = nn.ModuleList(
+            [MistralDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self._attn_implementation = config._attn_implementation
+        self.norm = MistralRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        self.gradient_checkpointing = False
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def _init_beacon_embed(self, missing_keys):
+        """Initialize the beacon token embedding with that of the eos token."""
+        if is_deepspeed_zero3_enabled():
+            import deepspeed
+            params = [self.beacon_embed_tokens.weight, self.embed_tokens.weight]
+            with deepspeed.zero.GatheredParameters(params, modifier_rank=0):
+                # deepspeed will initialize the parameters to zero
+                if (self.beacon_embed_tokens.weight == 0).all():
+                    if self.config.beacon_embed_init == "bos":
+                        self.beacon_embed_tokens.weight.data[:] = self.embed_tokens.weight.data[self.config.bos_token_id]
+                    elif self.config.beacon_embed_init == "eos":
+                        if isinstance(self.config.eos_token_id, list):
+                            eos_token_id = self.config.eos_token_id[0]
+                        else:
+                            eos_token_id = self.config.eos_token_id
+                        self.beacon_embed_tokens.weight.data[:] = self.embed_tokens.weight.data[eos_token_id]
+                    else:
+                        raise NotImplementedError(f"Make sure beacon_embed_init is either eos or bos, found {self.config.beacon_embed_init}")
+        else:
+            if any("beacon_embed_tokens" in missing_key for missing_key in missing_keys):
+                if self.config.beacon_embed_init == "bos":
+                    self.beacon_embed_tokens.weight.data[:] = self.embed_tokens.weight.data[self.config.bos_token_id]
+                elif self.config.beacon_embed_init == "eos":
+                    if isinstance(self.config.eos_token_id, list):
+                        eos_token_id = self.config.eos_token_id[0]
+                    else:
+                        eos_token_id = self.config.eos_token_id
+                    self.beacon_embed_tokens.weight.data[:] = self.embed_tokens.weight.data[eos_token_id]
+                else:
+                    raise NotImplementedError(f"Make sure beacon_embed_init is either eos or bos, found {self.config.beacon_embed_init}")
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.embed_tokens = value
+
+    @add_start_docstrings_to_model_forward(MISTRAL_INPUTS_DOCSTRING)
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        # BEACON: always use cache
+        use_cache = True
+
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # retrieve input_ids and inputs_embeds
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+        elif input_ids is not None:
+            batch_size, seq_length = input_ids.shape[:2]
+        elif inputs_embeds is not None:
+            batch_size, seq_length = inputs_embeds.shape[:2]
+        else:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+        past_key, past_value, beacon_size, beacon_indices = past_key_values[0]
+
+        # BEACON: separately embed ordinal tokens and beacon tokens because ordinal tokens do not receive gradients
+        if beacon_size > 0:
+            # NOTE: when beacon_pos == "interleave", the beacon_indices points to all beacon tokens in the current window (cached activations + input_ids), so we shall slice out the part corresponding to the input_ids
+            cur_beacon_indices = beacon_indices[-input_ids.shape[1]:]
+
+            ordinal_input_ids = input_ids[:, cur_beacon_indices == 0]
+            beacon_input_ids = input_ids[:, cur_beacon_indices > 0]
+            ordinal_inputs_embeds = self.embed_tokens(ordinal_input_ids)
+            beacon_input_embeds = self.beacon_embed_tokens(beacon_input_ids - self.config.vocab_size)
+            # create a new embedding tensor
+            inputs_embeds = beacon_input_embeds.new_zeros(*input_ids.shape, beacon_input_embeds.shape[-1])
+            inputs_embeds[:, cur_beacon_indices == 0] = ordinal_inputs_embeds
+            inputs_embeds[:, cur_beacon_indices > 0] = beacon_input_embeds
+
+        else:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        # embed positions
+        hidden_states = inputs_embeds
+
+        # print(f"input_ids:          {input_ids}")
+        # print(f"beacon_indices:     {beacon_indices}")
+        # print(f"position_ids:       {position_ids}")
+        # print(f"attention_mask:\n{attention_mask == 0}")
+        # x = input()
+        # if x == "s":
+        #     return
+
+        # decoder layers
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+        # BEACON: still use tuple to organize cache
+        next_decoder_cache = () if use_cache else None
+
+        for idx, decoder_layer in enumerate(self.layers):
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            # BEACON: slice out the past_key_value of the corresponding layer
+            past_key_value = past_key_values[idx] if past_key_values is not None else None
+
+            if self.gradient_checkpointing and self.training:
+                layer_outputs = self._gradient_checkpointing_func(
+                    decoder_layer.__call__,
+                    hidden_states,
+                    attention_mask,
+                    position_ids,
+                    past_key_value,
+                    output_attentions,
+                    use_cache,
+                )
+            else:
+                layer_outputs = decoder_layer(
+                    hidden_states,
+                    attention_mask=attention_mask,
+                    position_ids=position_ids,
+                    past_key_value=past_key_value,
+                    output_attentions=output_attentions,
+                    use_cache=use_cache,
+                )
+
+            hidden_states = layer_outputs[0]
+
+            if use_cache:
+                next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)
+
+            if output_attentions:
+                all_self_attns += (layer_outputs[1],)
+
+        hidden_states = self.norm(hidden_states)
+
+        # add hidden states from the last decoder layer
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        next_cache = next_decoder_cache if use_cache else None
+
+        if not return_dict:
+            return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=next_cache,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+        )
+
+
+class MistralForCausalLM(MistralPreTrainedModel):
+    _tied_weights_keys = ["lm_head.weight"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = MistralModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    def get_decoder(self):
+        return self.model
+
+    @classmethod
+    def from_pretrained(cls, *args, **kwargs):
+        """Override the default from_pretrained to extend vocab size according to beacon_size."""
+        kwargs.update(output_loading_info=True)
+        model, loading_info = super().from_pretrained(*args, **kwargs)
+
+        # NOTE: set memory after from_pretrained because there may be another transformer model inside the Memory object, which may cause weird erros during loading
+        config = model.config
+        model.memory = Memory(
+            model_config=config,
+            k_seq_dim=2,
+            v_seq_dim=2,
+        )
+
+        missing_keys = loading_info["missing_keys"]
+        # NOTE: the beacon parameters may or may not be loaded from the checkpoint
+        # if it is loaded from the checkpoint, we should not re-initilize it
+        model.model._init_beacon_embed(missing_keys)
+        # initialize weights of possible q,k,v,o,mlp
+        for layer in model.model.layers:
+            layer.self_attn._init_beacon_proj(missing_keys)
+            layer.mlp._init_beacon_proj(missing_keys)
+
+        return model
+
+    def _native_forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        shift_labels: Optional[bool] = True,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BeaconModelOutput]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # when we directly call _native_forward, the past_key_values would be None
+        if past_key_values is None:
+            # NOTE: set beacon size to 0 to avoid using any beacon parameters, see MistralAttention.forward
+            past_key_values = [(None, None, 0, None) for _ in range(self.config.num_hidden_layers)]
+
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        hidden_states = outputs[0]
+        logits = self.lm_head(hidden_states)
+        logits = logits.float()
+
+        loss = None
+        batch_loss = None
+        valid_token_num = None
+        
+        if labels is not None:
+            loss, batch_loss, valid_token_num = compute_loss(logits, labels, shift=shift_labels)
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+
+        return BeaconModelOutput(
+            loss=loss,
+            batch_loss=batch_loss,
+            valid_token_num=valid_token_num,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def _beacon_forward(self, 
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ):
+        # t1 = time.time()
+
+        # initialize cache
+        self.memory.prepare(
+            input_ids=input_ids, 
+            attention_mask=attention_mask, 
+            labels=labels
+        )
+
+        # t2 = time.time()
+
+        # after the first window, one token at a time
+        while not self.memory.finish:
+
+            # t3 = time.time()
+
+            input_ids, attention_mask, position_ids, past_key_values, labels = self.memory.step()
+
+            # t4 = time.time()
+
+            outputs = self._native_forward(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                inputs_embeds=inputs_embeds,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+                labels=labels,
+                # NOTE: the labels have been shifted so that all tokens in the window have the proper loss
+                shift_labels=False,
+            )
+
+            # t5 = time.time()
+
+            # update past_key_values
+            self.memory.update_memory(outputs.past_key_values)
+
+            # t6 = time.time()
+
+            if labels is not None:
+                # update loss
+                self.memory.update_loss(outputs.batch_loss, outputs.valid_token_num)
+
+            # t7 = time.time()
+
+            # print(f"step time: {t4-t3}, forward time: {t5-t4}, update time: {t6-t5}, loss time: {t7-t6}")
+            # input()
+
+        # t8 = time.time()
+
+        # output loss, past_key_values, and perplexity
+        outputs = self.memory.output(outputs)
+
+        # t9 = time.time()
+
+        # print(f"output time:            {t9-t8}")
+        # input()
+
+        return outputs
+    
+    def forward(self, **kwargs):
+        """Forward computation over a batch of sequences.
+        """
+        # only allow gradient when training
+        with optional_grad_ctx(with_grad=self.training):
+            # we can disable beacon to use the original mistral
+            if hasattr(self, "_enable_beacon") and self._enable_beacon == False:
+                return self._native_forward(**kwargs)
+            else:
+                return self._beacon_forward(**kwargs)
+
+    def prepare_inputs_for_generation(
+        self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
+    ):
+        if past_key_values:
+            input_ids = input_ids[:, -1:]
+
+        position_ids = kwargs.get("position_ids", None)
+        if attention_mask is not None and position_ids is None:
+            # create position_ids on the fly for batch generation
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            if past_key_values:
+                position_ids = position_ids[:, -1].unsqueeze(-1)
+
+        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+        if inputs_embeds is not None and past_key_values is None:
+            model_inputs = {"inputs_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids}
+
+        model_inputs.update(
+            {
+                "position_ids": position_ids,
+                "past_key_values": past_key_values,
+                "use_cache": kwargs.get("use_cache"),
+                "attention_mask": attention_mask,
+            }
+        )
+        return model_inputs
+
+    @staticmethod
+    def _reorder_cache(past_key_values, beam_idx):
+        reordered_past = ()
+        for layer_past in past_key_values:
+            reordered_past += (
+                tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
+            )
+        return reordered_past

Long_LLM/activation_beacon/new/src/modeling_beacon.py

+import os
+import torch
+import time
+import numpy as np
+import torch.distributed as dist
+from transformers.utils import logging
+from transformers import AutoTokenizer
+from itertools import cycle
+from typing import List
+
+logger = logging.get_logger(__name__)
+
+
+class Memory(torch.nn.Module):
+    def __init__(
+        self, 
+        model_config, 
+        k_seq_dim:int=2, 
+        v_seq_dim:int=2, 
+    ):
+        """Setup necessary attributes."""
+        super().__init__()
+
+        self.config = model_config
+
+        # initialize necessary parameters
+        self.k_seq_dim = k_seq_dim
+        self.v_seq_dim = v_seq_dim
+        self.rng = np.random.default_rng(42)
+
+        self.beacon_token = torch.tensor([self.config.vocab_size])
+
+        self._post_validation()
+        self.reset()
+
+    def _post_validation(self, verbose=True):
+        assert self.config.beacon_window >= self.config.beacon_stride, f"Make sure the beacon_window {self.config.beacon_window} >= beacon_stride {self.config.beacon_stride}!"
+        for ratio in self.config.beacon_ratio:
+            assert ratio >= 0, f"Make sure all beacon ratios are greater than or equal to 0, found {self.config.beacon_ratio}!"
+        assert self.config.beacon_attn in ["segmentation", "step-expansion", "full-coverage"], f"beacon_attn {self.config.beacon_attn} not implemented!"
+        assert self.config.beacon_ratio_mix in ["instance-random", "step-random", "sequence"] or "adapt-" in self.config.beacon_ratio_mix, f"beacon_ratio_mix {self.config.beacon_ratio_mix} not implemented!"
+        # assert self.config.beacon_pos in ["append", "interleave"], f"beacon_pos {self.config.beacon_pos} not implemented!"
+        if self.config.beacon_pos == "interleave":
+            assert self.config.beacon_window == self.config.beacon_stride, f"Make sure the beacon_window equals to beacon_stride when using interleaving mode."
+        if self.config.beacon_parallel_window > 1:
+            assert self.config._attn_implementation != "flash_attention_2", f"Currently parallel window does not support flash_attention_2!"
+
+        self._cpu = torch.device("cpu")
+
+        if verbose:
+            info = f"applying activation beacon on {self.config.beacon_param} (the beacon embedding is initialized from {'bos' if self.config.beacon_embed_init == 'bos' else 'eos'} embedding, the beacon tokens are positioned with '{self.config.beacon_pos}' method), with window size {self.config.beacon_window}, stride {self.config.beacon_stride}, {self.config.beacon_attn} attention{' (attending to previous beacons)' if self.config.beacon_attend_prev else ' (no attending to previous beacons)'}, sink size {self.config.beacon_sink_size}, compression ratio {self.config.beacon_ratio} (mixed by {self.config.beacon_ratio_mix})..."
+            logger.info(info)
+
+    def set(self, verbose=True, **kwargs):
+        """
+        Set attributes out of the constructor.
+        """
+        for k, v in kwargs.items():
+            setattr(self.config, k, v)
+        self._post_validation(verbose=verbose)
+
+    def reset(self):
+        """Initialize attributes for a new sequence."""
+        # the cursor pointing to the start of the current window
+        self._start_idx = 0
+        # the cursor pointing to the end of the current window
+        self._end_idx = 0
+        # the beacon sizes of all strides
+        self._all_beacon_sizes = []
+        # the loss per batch
+        self._batch_loss = None
+        # the valid token number per batch
+        self._valid_token_num = None
+        # the step index for processing the input_ids
+        self._step_idx = 0
+        # used in set_compression_ratio
+        self._compression_ratio = None
+        # the previous inputs is a full window or not, defaults to True
+        self._is_full_window = True
+        # the number of raw activations to preserve in update_memory (only useful when beacon_stride < beacon_window)
+        self._raw_size_to_cache = 0
+
+        # the number of tokens in previous stride that should be compressed by the upcoming beacon
+        self._interleave_remainder = 0
+        # compression ratio for the unfinished window
+        self._interleave_compression_ratio = None
+        self._beacon_indices = None
+
+        self.all_input_ids = None
+        self.all_attention_mask = None
+        self.all_labels = None
+
+        # the raw activations of recent tokens
+        self.raw_activations = [(None, None) for _ in range(self.config.num_hidden_layers)]
+        # the attention sink activations
+        self.sink_activations = [(None, None) for _ in range(self.config.num_hidden_layers)]
+        # the beacon activations
+        self.beacon_activations = [(None, None) for _ in range(self.config.num_hidden_layers)]
+
+    @property
+    def all_sequence_length(self):
+        if self.all_input_ids is None:
+            return 0
+        else:
+            return self.all_input_ids.shape[1]
+
+    @property
+    def batch_size(self):
+        if self.all_input_ids is None:
+            return 0
+        else:
+            return self.all_input_ids.shape[0]
+
+    @property
+    def finish(self):
+        is_finish = self._end_idx == self.all_sequence_length
+        return is_finish
+
+    @property
+    def dtype(self):
+        return self.config.torch_dtype
+
+    @property
+    def min_value(self):
+        return torch.finfo(self.dtype).min
+
+    @property
+    def max_position_embeddings(self):
+        max_position_embeddings = self.config.max_position_embeddings
+        if getattr(self.config, "rope_scaling", None) is not None:
+            scaling_factor = self.config.rope_scaling["factor"]
+            max_position_embeddings = max_position_embeddings * scaling_factor
+        return max_position_embeddings
+ 
+    def get_memory_size(self):
+        """
+        Sink memory size, beacon memory size and raw memory size.
+        """
+        sink_memory_size = 0
+        beacon_memory_size = 0
+        raw_memory_size = 0
+        if self.sink_activations[0][0] is not None:
+            sink_memory_size += self.sink_activations[0][0].shape[self.k_seq_dim]
+        if self.beacon_activations[0][0] is not None:
+            beacon_memory_size += self.beacon_activations[0][0].shape[self.k_seq_dim]
+        if self.raw_activations[0][0] is not None:
+            raw_memory_size += self.raw_activations[0][0].shape[self.k_seq_dim]
+        return sink_memory_size, beacon_memory_size, raw_memory_size
+
+    def prepare(self, input_ids, attention_mask, labels):
+        """
+        Prepare inputs for the model. These inputs belong to the same sequence.
+        """
+        # assert input_ids.shape[0] == 1, "Make sure the batch size is 1!"
+        # assert attention_mask is None or (attention_mask == 1).all(), "Make sure there is no padding!"
+
+        self._device = input_ids.device
+
+        # accumulate input_ids
+        if self.all_input_ids is None:
+            self.all_input_ids = input_ids.cpu()
+        else:
+            self.all_input_ids = torch.cat([self.all_input_ids, input_ids.cpu()], dim=1)
+
+        # accumulate attention_mask
+        if attention_mask is None:
+            attention_mask = torch.ones_like(input_ids, device=torch.device("cpu"))
+        if self.all_attention_mask is None:
+            self.all_attention_mask = attention_mask.cpu()
+        else:
+            self.all_attention_mask = torch.cat([self.all_attention_mask, attention_mask.cpu()], dim=1)
+
+        # accumulate labels if exisits
+        if labels is not None:
+            # rotate labels in advance so that the loss of the last token is not ignored in every window
+            labels = torch.cat([labels[:, 1:].cpu(), torch.tensor([-100]).expand(labels.shape[0], 1)], dim=1)
+            if self.all_labels is None:
+                self.all_labels = labels.cpu()
+            else:
+                self.all_labels = torch.cat([self.all_labels, labels], dim=1)
+            assert self.all_input_ids.shape[1] == self.all_labels.shape[1], f"Found inconsistent all_input_ids {self.all_input_ids.shape} and all_labels {self.all_labels.shape}!"
+
+    def set_compression_ratio(self, start_idx, end_idx):
+        """Choose a condensing ratio from self.config.beacon_ratio"""
+        def filter_ratio(ratios, stride):
+            valid_ratios = []
+            for ratio in ratios:
+                # stride must be bigger than condensing ratio because we there must be at least one beacon
+                if stride < ratio:
+                    continue
+                # the stride must be evenly divisible by condensing ratio
+                if ratio > 0 and (stride % ratio) != 0:
+                    continue
+                # when training, ratio=0 is valid if previous windows contain beacon or later windows contain beacon
+                if ratio == 0 and self.training:
+                    previous_has_zero = -1 in self._all_beacon_sizes
+                    following_has_nonzero = (start_idx + stride + self.config.beacon_window) <= self.all_sequence_length
+                    if previous_has_zero or (not following_has_nonzero):
+                        continue
+                valid_ratios.append(ratio)
+            assert len(valid_ratios), f"Cannot find valid condensing ratio (among {ratios}) for stride {stride}!"
+            return valid_ratios
+
+        def get_max_length(ratios):
+            max_lengths = []
+            for compression_ratio in ratios:
+                if compression_ratio > 0:
+                    # NOTE: here we must use the scaled position embeddings
+                    max_lengths.append((self.max_position_embeddings - self.config.beacon_window) * compression_ratio + self.config.beacon_window)
+                else:
+                    max_lengths.append(self.max_position_embeddings)
+            return max_lengths
+
+        if len(self.config.beacon_ratio) == 1:
+            return self.config.beacon_ratio[0]
+
+        ratio_mix = self.config.beacon_ratio_mix
+
+        beacon_ratio = filter_ratio(self.config.beacon_ratio, self.config.beacon_stride)
+
+        if ratio_mix == "instance-random":
+            if self._compression_ratio is None:
+                beacon_ratio = self.rng.choice(beacon_ratio).tolist()
+                self._compression_ratio = beacon_ratio
+            else:
+                beacon_ratio = self._compression_ratio
+
+        elif ratio_mix == "step-random":
+            beacon_ratio = self.rng.choice(beacon_ratio).tolist()
+        
+        elif ratio_mix == "sequence":
+            if self._compression_ratio is None:
+                self._compression_ratio = cycle(beacon_ratio)
+            beacon_ratio = next(self._compression_ratio)
+
+        elif "adapt" in ratio_mix:
+            if self._compression_ratio is None:
+                future_length = int(ratio_mix.split("-")[1])
+                sequence_length = self.all_input_ids.shape[1] + future_length
+                max_lengths = get_max_length(beacon_ratio)
+                # ascendingly sort the max lengths
+                valid_max_lengths_and_indices = [x for x in enumerate(max_lengths) if x[1] >= sequence_length]
+                if len(valid_max_lengths_and_indices):
+                    minimum_length_index = min(valid_max_lengths_and_indices, key=lambda x: x[1])[0]
+                    # use the minimal possible length for this sequence (the smallest fold ratio)
+                    beacon_ratio = beacon_ratio[minimum_length_index]
+                else:
+                    beacon_ratio = max(beacon_ratio)
+                    # logger.warning(f"Failed to find valid fold window and size for sequence length {sequence_length}, as the maximum theoretical length is {max(max_lengths)}. Fall back to use the maximum one: {beacon_ratio}.")
+                self._compression_ratio = beacon_ratio
+            else:
+                beacon_ratio = self._compression_ratio
+
+        return beacon_ratio
+
+    def step(self):
+        # parallel does not support stride < window
+        # parallel does not support non-compression
+        # the input_ids is not long enough for parallel
+        if \
+        (self.config.beacon_parallel_window > 1) and \
+        (self.config.beacon_stride == self.config.beacon_window) and \
+        (0 not in self.config.beacon_ratio) and \
+        (self.all_input_ids[:, self._end_idx:].shape[1] >= self.config.beacon_parallel_window * self.config.beacon_window):
+            input_ids_list = []
+            attention_mask_list = []
+            position_ids_list = []
+            labels_list = []
+
+            beacon_size_list = []
+            beacon_indices_list = []
+
+            for i in range(self.config.beacon_parallel_window):
+                if i == 0:
+                    _input_ids, _attention_mask, _position_ids, _past_key_values, _labels = self._step()
+                else:
+                    _input_ids, _attention_mask, _position_ids, _past_key_values, _labels = self._step(ignore_memory=True)
+
+                input_ids_list.append(_input_ids)
+                attention_mask_list.append(_attention_mask)
+                position_ids_list.append(_position_ids)
+                labels_list.append(_labels)
+                beacon_size_list.append(_past_key_values[0][2])
+                beacon_indices_list.append(_past_key_values[0][3])
+
+                if i == 0:
+                    past_key_values = _past_key_values
+                    if past_key_values[0][0] is None:
+                        mem_size = 0
+                    else:
+                        mem_size = past_key_values[0][0].shape[self.k_seq_dim]
+
+                else:
+                    # no memory
+                    assert _past_key_values[0][0] is None
+            
+            batch_size = self.all_input_ids.shape[0]
+            # NOTE: we do not need to repliace beacon tokens for the last window
+            seq_len = sum(x.shape[1] for x in input_ids_list) + sum(beacon_size_list) - beacon_size_list[-1]
+
+            input_ids = _input_ids.new_zeros((batch_size, seq_len)) + self.beacon_token.to(_input_ids.device)
+            # all 0
+            attention_mask = _attention_mask.new_zeros((batch_size, 1, seq_len, mem_size + seq_len)) + self.min_value
+            position_ids = torch.arange(mem_size + seq_len, device=self._device).expand(batch_size, mem_size + seq_len)
+            # 2 indicates the beacon token is used for replication
+            beacon_indices = beacon_indices_list[0].new_zeros(seq_len) + 2
+            if _labels is not None:
+                # -100 because no loss on beacon tokens
+                labels = _labels.new_zeros((batch_size, seq_len)) - 100
+            else:
+                labels = None
+
+            start_idx = 0
+            position_offset = mem_size
+            for i in range(self.config.beacon_parallel_window):
+                beacon_size = beacon_size_list[i]
+
+                # populate input_ids
+                _input_ids = input_ids_list[i]
+                cur_seq_len = _input_ids.shape[1]
+                input_ids[:, start_idx: start_idx + cur_seq_len] = _input_ids
+                
+                # populate attention_mask and position_ids
+                _attention_mask = attention_mask_list[i]
+                _position_ids = position_ids_list[i]
+                # the attention mask in the first window contains the mask for memory, which is redundant here
+                if i == 0:
+                    _attention_mask = _attention_mask[:, :, :, mem_size:]
+                    _position_ids = _position_ids[:, mem_size:] - mem_size
+
+                attention_mask[:, :, start_idx: start_idx + cur_seq_len, mem_size + start_idx: mem_size + start_idx + cur_seq_len] = _attention_mask
+                position_ids[:, mem_size + start_idx: mem_size + start_idx + cur_seq_len] = _position_ids + position_offset
+
+                # populate beacon_indices
+                _beacon_indices = beacon_indices_list[i]
+                beacon_indices[start_idx: start_idx + cur_seq_len] = _beacon_indices
+
+                # populate labels
+                if labels is not None:
+                    # populate labels
+                    _labels = labels_list[i]
+                    labels[:, start_idx: start_idx + cur_seq_len] = _labels
+
+                # NOTE: when there is sink activations, we need to bias the position_ids for the first window
+                if i == 0 and self.config.beacon_sink_size > 0 and self.sink_activations[0][0] is None:
+                    position_offset += 1
+
+                # modify the attention and position for replicated beacon tokens
+                if i != self.config.beacon_parallel_window - 1:
+                    replicate_beacon_row_start = start_idx + cur_seq_len
+                    replicate_beacon_col_start = mem_size + start_idx + cur_seq_len
+                    # NOTE: any attention mask is okay for replicated beacon tokens, but for convenience we use the causal mask
+                    attention_mask[:, :, replicate_beacon_row_start: replicate_beacon_row_start + beacon_size, replicate_beacon_col_start: replicate_beacon_col_start + beacon_size] = _attention_mask.new_full((beacon_size, beacon_size), self.min_value).triu(1)
+                    # NOTE: all future tokens can attend to the replicated beacon tokens
+                    attention_mask[:, :, replicate_beacon_row_start + beacon_size:, replicate_beacon_col_start: replicate_beacon_col_start + beacon_size] = 0
+                    # NOTE: the position of replicated beacon tokens start from 0
+                    position_ids[:, mem_size + start_idx + cur_seq_len: mem_size + start_idx + cur_seq_len + beacon_size] = torch.arange(position_offset, position_offset + beacon_size, device=_input_ids.device)[None:]
+
+                start_idx += cur_seq_len + beacon_size
+                position_offset += beacon_size
+
+            # the memory is visible to all subsequent tokens
+            attention_mask[:, :, :, :max(mem_size, self.config.beacon_sink_size)] = 0
+
+            # NOTE: modify beacon_indices
+            for i, (key, value, _, _) in enumerate(past_key_values):
+                past_key_values[i] = (key, value, sum(beacon_size_list), beacon_indices)
+
+            # NOTE: update _beacon_indices so that the next-token logits can be properly sliced out in self.output()
+            self._beacon_indices = beacon_indices
+            
+            return input_ids, attention_mask, position_ids, past_key_values, labels
+
+        else:
+            return self._step()
+
+    def _step(self, ignore_memory=False):
+        """
+        Yield inputs for the current sliding window, including the input_ids, attention_mask, position_ids, and past_key_values.
+        """
+        #============================================#
+        # Check whether the inputs fulfills a window.
+        #============================================#
+
+        # the starting position of the current window w.r.t. the start of the current input sequence
+        start_idx = self._start_idx
+        # the end position of the current window w.r.t. the start of the current input sequence
+        end_idx = start_idx + self.config.beacon_window
+        # indicates if the current window is completely filled by raw activations and new tokens
+        # we only append beacon tokens for full windows
+        if end_idx > self.all_sequence_length:
+            # the input is shorter than the initial window size
+            end_idx = self.all_sequence_length
+            is_full_window = False
+        else:
+            is_full_window = True
+
+        # NOTE: in training, the entire sequence is input to the model at once
+        # In the last window, we do not need to append beacons because they will not be used at all
+        if self.training and end_idx == self.all_sequence_length:
+            next_start_idx = start_idx
+            raw_size_to_cache = -1
+            beacon_size = 0
+            compression_ratio = 1
+            is_full_window = False
+
+        else:
+            #============================================#
+            # Set compression ratio
+            #============================================#
+            if self.config.beacon_pos == "append":
+                if is_full_window:
+                    # determine compression ratio for the current window
+                    beacon_stride = self.config.beacon_stride
+                    compression_ratio = self.set_compression_ratio(start_idx=start_idx, end_idx=end_idx)
+
+                    if compression_ratio > 0:
+                        # the stride must be evenly divisible by compression_ratio
+                        beacon_size = beacon_stride // compression_ratio
+                    else:
+                        # the raw activations are used as beacon activations
+                        beacon_size = -1
+
+                    # forward start_idx and end_idx
+                    next_start_idx = start_idx + beacon_stride
+                    # how many raw activations to save
+                    raw_size_to_cache = end_idx - next_start_idx
+                else:
+                    # no stride because the sequence has finished
+                    next_start_idx = start_idx
+                    # cache all raw activations
+                    raw_size_to_cache = -1
+                    beacon_size = 0
+                    compression_ratio = 0
+
+            elif self.config.beacon_pos == "interleave":
+                # the number of raw tokens in the input_ids
+                input_size = end_idx - self._end_idx
+                # set compression ratio once the previous window has finished, otherwise, reuse the interleave_compression_ratio if the input belongs to an unfinished window
+                if self._is_full_window:
+                    compression_ratio = self.set_compression_ratio(start_idx=start_idx, end_idx=end_idx)
+                    self._interleave_compression_ratio = compression_ratio
+                else:
+                    compression_ratio = self._interleave_compression_ratio
+
+                # the beacon size is non-zero even if the window is not full
+                if compression_ratio > 0:
+                    # this number of beacon tokens will be inserted among the raw tokens
+                    beacon_size = (input_size + self._interleave_remainder) // compression_ratio
+                else:
+                    # the raw activations are used as beacon activations
+                    beacon_size = -1
+
+                if is_full_window:
+                    # move forward one window
+                    next_start_idx = start_idx + self.config.beacon_stride
+                    # no save raw activations
+                    raw_size_to_cache = 0
+                else:
+                    # no stride because the sequence has not finished
+                    next_start_idx = start_idx
+                    # cache all recent raw activations to be used in the next window
+                    raw_size_to_cache = -1
+
+        #============================================#
+        # Slice out input_ids (raw tokens in the current window)
+        #============================================#
+        input_ids = self.all_input_ids[:, self._end_idx: end_idx].to(self._device)
+        attention_mask = self.all_attention_mask[:, self._end_idx: end_idx].to(self._device)
+        if self.all_labels is not None:
+            labels = self.all_labels[:, self._end_idx: end_idx].to(self._device)
+        else:
+            labels = None
+        batch_size = input_ids.shape[0]
+
+        #============================================#
+        # Insert beacon tokens if necessary.
+        #============================================#
+        # t1 = time.time()
+
+        if self.config.beacon_pos == "append":
+            # append beacons if necessary
+            if is_full_window and beacon_size > 0:
+                input_ids = torch.cat([input_ids, self.beacon_token.expand(batch_size, beacon_size).to(input_ids.device, dtype=input_ids.dtype)], dim=1)
+                # NOTE: prepend 1 to attention_mask because we have past_key_values
+                attention_mask = torch.cat([attention_mask, attention_mask.new_ones(batch_size, beacon_size)], dim=1)
+                if labels is not None:
+                    labels = torch.cat([labels, labels.new_zeros(batch_size, beacon_size) - 100], dim=1)
+
+        elif self.config.beacon_pos == "interleave":
+            input_len = input_ids.shape[1]
+            if beacon_size > 0:
+                # insert beacon tokens in between raw tokens
+                input_ids_with_beacons = input_ids.new_full((input_ids.shape[0], input_len + beacon_size), self.beacon_token.item())
+                raw_token_indices = torch.arange(input_ids_with_beacons.shape[1], device=input_ids.device)
+                interleave_start_idx = compression_ratio - self._interleave_remainder
+                raw_token_indices = raw_token_indices[raw_token_indices % (compression_ratio + 1) != interleave_start_idx].unsqueeze(0).expand_as(input_ids)
+                input_ids_with_beacons = input_ids_with_beacons.scatter(dim=1, index=raw_token_indices, src=input_ids)
+                input_ids = input_ids_with_beacons
+                # attention mask
+                attention_mask_with_beacons = attention_mask.new_full((attention_mask.shape[0], attention_mask.shape[1] + beacon_size), 1)
+                attention_mask_with_beacons = attention_mask_with_beacons.scatter(dim=1, index=raw_token_indices, src=attention_mask)
+                attention_mask = attention_mask_with_beacons
+                # labels
+                if labels is not None:
+                    labels_with_beacons = labels.new_full((labels.shape[0], labels.shape[1] + beacon_size), -100)
+                    labels_with_beacons = labels_with_beacons.scatter(dim=1, index=raw_token_indices, src=labels)
+                    labels = labels_with_beacons
+
+            if compression_ratio > 0:
+                # update the reminder
+                self._interleave_remainder = (input_len + self._interleave_remainder) % compression_ratio
+
+            # NOTE: skip computing loss in the very first window because the beacon tokens will be used in the next window
+            if self.training and self._step_idx == 0 and not (self.config.beacon_pos == 'interleave' and self.config.beacon_attn == 'full-coverage'):
+                labels[:] = -100
+
+        # t2 = time.time()
+
+        #============================================#
+        # Prepare beacon_indices for interleave beacon_pos, a boolean mask where True indicates the beacon tokens.
+        # The mask is applied on the inputs of the entire window, including the cached activations and the input_ids.
+        #============================================#
+        beacon_indices = (input_ids[0] == self.beacon_token.item()).long()
+        if self._is_full_window:
+            self._beacon_indices = torch.tensor([], dtype=torch.long, device=input_ids.device)
+        # the beacon_indices always tracks the beacon tokens in both the cached activations and the input_ids
+        beacon_indices = torch.cat([self._beacon_indices, beacon_indices])
+        # record the beacon_indices for the next window
+        self._beacon_indices = beacon_indices
+        if is_full_window and beacon_size == -1:
+            # NOTE: the first beacon_stride raw tokens serve as beacon tokens
+            # we use -1 to indicate these raw tokens, so that the attention mask and position ids will not be modified
+            beacon_indices[:self.config.beacon_stride] = -1
+
+        # t3 = time.time()
+
+        #============================================#
+        # Prepare past_key_values.
+            # beacon_size: how many beacon tokens are there in the input_ids
+            # beacon_indices: the boolean mask for the entire window where True indicates the beacon tokens (for append, the beacon_indices corresponds to input_ids, while for 'interleave', the beacon_indices corresponds to the entire window including both the input_ids and the cached activations)
+        #============================================#
+        past_key_values = []
+        for layer_idx in range(self.config.num_hidden_layers):
+            if ignore_memory:
+                key, value = None, None
+            else:
+                sink_key, sink_value = self.sink_activations[layer_idx]
+                beacon_key, beacon_value = self.beacon_activations[layer_idx]
+                raw_key, raw_value = self.raw_activations[layer_idx]
+
+                key = cat_tensor([
+                    sink_key, beacon_key, raw_key,
+                ], dim=self.k_seq_dim)
+                value = cat_tensor([
+                    sink_value, beacon_value, raw_value,
+                ], dim=self.v_seq_dim)
+
+            layer_past_key_values = (key, value, beacon_size, beacon_indices)
+            past_key_values.append(layer_past_key_values)
+
+        # t4 = time.time()
+        
+        #============================================#
+        # Prepare attention_mask and position_ids.
+        #============================================#
+        first_key = past_key_values[0][0]
+        mem_size = first_key.shape[self.k_seq_dim] if first_key is not None else 0
+        if mem_size > 0:
+            attention_mask = torch.cat([attention_mask.new_ones(batch_size, mem_size), attention_mask], dim=1)
+
+        input_length = input_ids.shape[1]
+        position_ids = torch.arange(attention_mask.shape[-1], dtype=torch.long, device=self._device).repeat(batch_size, 1)
+
+        if self.config._attn_implementation == "flash_attention_2":
+            assert self.config.beacon_attn == "full-coverage", f"Make sure to set beacon_attn='full-coverage' when using flash attention! Found {self.config.beacon_attn}."
+            if 0 in attention_mask:
+                pass
+            else:
+                attention_mask = None
+        elif self.config._attn_implementation == "sdpa" and self.config.beacon_pos == "append" and beacon_size <= 0 and (input_length == 1 or mem_size == 0):
+            attention_mask = None
+        else:
+            attention_mask, position_ids = self._make_4d_attention_mask_and_position_ids(
+                attention_mask, 
+                position_ids, 
+                mem_size, 
+                beacon_size, 
+                compression_ratio, 
+            )
+
+        # t5 = time.time()
+
+        # print(f"prepare inputs {t2-t1}, prepare indices {t3-t2}, prepare memory {t4-t3}, prepare attention mask {t5-t4}")
+
+        #============================================#
+        # Update necessary attributes.
+        #============================================#
+        # keep track of whether the current inputs is a full_window
+        self._is_full_window = is_full_window
+        # keep track of the raw_size_to_cache
+        self._raw_size_to_cache = raw_size_to_cache
+        # involked in self.output()
+        self._all_beacon_sizes.append(beacon_size)
+        # update end_idx
+        self._start_idx = next_start_idx
+        self._end_idx = end_idx
+        self._step_idx += 1
+
+        # print(f"beacon_size:        {beacon_size}")
+        # print(f"raw_size_to_cache:  {raw_size_to_cache}")
+        # print(f"input_ids:          {input_ids}")
+        # print(f"beacon_indices:     {beacon_indices}")
+        # print(f"position_ids:       {position_ids}")
+        # print(f"attention_mask:\n{attention_mask}")
+        # x = input()
+        # if x == "s":
+        #     return
+
+        return input_ids, attention_mask, position_ids, past_key_values, labels
+
+    def update_memory(self, past_key_values):
+        """
+        Accumulate beacon activations and raw activations.
+        """
+        for layer_idx, (key, value, beacon_size, beacon_indices) in enumerate(past_key_values):
+            # NOTE: the past_key_values are incrementally returned (only the new keys and values are returned)
+            previous_raw_key, previous_raw_value = self.raw_activations[layer_idx]
+
+            if self.beacon_activations[layer_idx][0] is None and self.config.beacon_sink_size > 0:
+                # save the sink activations
+                # NOTE: we do not slice the key/value activations, which may cause duplication when beacon_ratio=-1 for the first window, but it's okay
+                self.sink_activations[layer_idx] = [
+                    slice_tensor(key, end=self.config.beacon_sink_size, dim=self.k_seq_dim),
+                    slice_tensor(value, end=self.config.beacon_sink_size, dim=self.v_seq_dim),
+                ]
+
+            if not self._is_full_window:
+                # this means the current input does not fulfill a window
+                # thus, the key and value are all raw activations, and we accumulate them until the window is fulfilled
+                assert self._raw_size_to_cache == -1
+                raw_key = cat_tensor([
+                    previous_raw_key,
+                    key
+                ], dim=self.k_seq_dim)
+                raw_value = cat_tensor([
+                    previous_raw_value, 
+                    value
+                ], dim=self.v_seq_dim)
+                self.raw_activations[layer_idx] = (raw_key, raw_value)
+
+            else:
+                # NOTE: use the correct previous_beacon_key and value!
+                previous_beacon_key, previous_beacon_value = self.beacon_activations[layer_idx]
+                
+                beacon_key, beacon_value, raw_key, raw_value = self._extract_beacon_and_raw_memory(
+                    key, 
+                    value, 
+                    previous_beacon_key, 
+                    previous_beacon_value, 
+                    previous_raw_key, 
+                    previous_raw_value, 
+                    beacon_indices,
+                )
+
+                self.beacon_activations[layer_idx] = (beacon_key, beacon_value)
+                self.raw_activations[layer_idx] = (raw_key, raw_value)
+
+    def update_loss(self, batch_loss, valid_token_num):
+        """
+        Accumulate loss for later perplexity computation and backward pass.
+        """
+        if self._batch_loss is None:
+            # NOTE: multiply valid_token_num because batch_loss is divided by it in advance
+            self._batch_loss = batch_loss * valid_token_num
+            self._valid_token_num = valid_token_num
+        else:
+            # NOTE: avoid in-place operations, otherwise there will be gradient errors in training
+            self._batch_loss = self._batch_loss + batch_loss * valid_token_num
+            self._valid_token_num = self._valid_token_num + valid_token_num
+
+    def output(self, model_outputs):
+        """
+        Override loss with accumulated loss. Update the next-token logits.
+        """
+        # override loss
+        if self._batch_loss is not None:
+            # here the batch_loss is the summation of all token losses in each element
+            loss = self._batch_loss.sum() / self._valid_token_num.sum()
+
+            # NOTE: prevent nan
+            batch_loss = self._batch_loss / self._valid_token_num
+            if (self._valid_token_num == 0).any():
+                batch_loss = batch_loss.masked_fill(self._valid_token_num == 0, 0.)
+
+            # NOTE: we must use dict to override values, otherwise trainer cannot find loss
+            model_outputs["loss"] = loss
+            model_outputs["batch_loss"] = batch_loss
+            model_outputs["valid_token_num"] = self._valid_token_num
+
+        # override last_hidden_states (used in generation)
+        beacon_size = self._all_beacon_sizes[-1]
+        # remove logits corresponding to beacon tokens
+        if beacon_size > 0:
+            logits = model_outputs["logits"]
+            beacon_indices = self._beacon_indices[-logits.shape[1]:]
+            model_outputs["logits"] = logits[:, beacon_indices == 0]
+
+        return model_outputs
+
+    def _make_4d_attention_mask_and_position_ids(
+        self, 
+        attention_mask, 
+        position_ids,
+        mem_size, 
+        beacon_size, 
+        compression_ratio, 
+    ):
+        """
+        Convert attention_mask into causal 4D attention_mask (batch_size, head_num, query_len, key_len).
+        """
+        tgt_size = attention_mask.size(-1) - mem_size
+        dtype = self.dtype
+        min_value = self.min_value
+        device = self._device
+        batch_size, src_size = attention_mask.size()
+
+        # square for memory, and lower triangular for input_ids
+        causal_mask = torch.full((tgt_size, tgt_size), min_value, device=device, dtype=dtype)
+        mask_cond = torch.arange(causal_mask.size(-1), device=device)
+        causal_mask.masked_fill_(mask_cond < (mask_cond + 1).view(causal_mask.size(-1), -1), 0)
+        causal_mask = torch.cat([torch.zeros(tgt_size, mem_size, dtype=dtype, device=device), causal_mask], dim=-1)
+        causal_mask = causal_mask[None, None, ...].expand(batch_size, 1, tgt_size, src_size)
+        # 1 for non-padding tokens
+        expand_mask = attention_mask[:, None, None, :].expand(batch_size, 1, tgt_size, src_size)
+        invert_mask = 1.0 - expand_mask
+        invert_mask.masked_fill_(invert_mask.bool(), min_value)
+
+        attention_mask = causal_mask.masked_fill(invert_mask.bool(), min_value)
+
+        if self.config.beacon_attn == "step-expansion":
+            # each beacon can attend to one more sub-interval than its predecessor
+
+            if self.config.beacon_pos == "append" and beacon_size > 0:
+                window_size = self.config.beacon_window
+                window_size_with_beacon = window_size + beacon_size
+                beacon_start_idx = -beacon_size
+                # batch_size, head_num, window_size
+                reference_attention_mask = attention_mask[..., -beacon_size - 1, -window_size_with_beacon: -beacon_size]
+
+                # compression_ratio, 2 * compression_ratio, ..., beacon_size * compression_ratio
+                beacon_arange = torch.arange(1, beacon_size + 1, device=device) * compression_ratio
+                # 0, 1, 2, ..., window_size - 1
+                ordinal_arange = torch.arange(window_size, device=device)
+                # beacon_size, window_size
+                valid_pos = ordinal_arange.expand(beacon_size, window_size) < beacon_arange.unsqueeze(-1)
+                # beacon_size, window_size
+                ordinal_attention_mask = torch.where(valid_pos, 0, min_value)
+                # NOTE: add reference attention_mask so that padding tokens are considered
+                ordinal_attention_mask = ordinal_attention_mask[None, None, ...] + reference_attention_mask.unsqueeze(-2)
+
+                if self.config.beacon_attend_prev:
+                    beacon_attention_mask = attention_mask.new_full((beacon_size, beacon_size), min_value).triu(1)
+                    # the beacon token is next to the last ordinal token it attends to
+                    ordinal_position_ids = position_ids[:, -window_size_with_beacon: -beacon_size]
+                    beacon_position_ids = ordinal_position_ids[:, compression_ratio - 1::compression_ratio] + torch.arange(1, beacon_size + 1, device=device)[None]
+                    position_ids[:, beacon_start_idx:] = beacon_position_ids
+                else:
+                    beacon_attention_mask = attention_mask.new_full((beacon_size, beacon_size), min_value).fill_diagonal_(0)
+                    # the beacon token is next to the last ordinal token it attends to
+                    ordinal_position_ids = position_ids[:, -window_size_with_beacon: -beacon_size]
+                    beacon_position_ids = ordinal_position_ids[:, compression_ratio - 1::compression_ratio] + 1
+                    position_ids[:, beacon_start_idx:] = beacon_position_ids
+
+                attention_mask[..., beacon_start_idx:, -window_size_with_beacon: -beacon_size] = ordinal_attention_mask
+                attention_mask[..., beacon_start_idx:, beacon_start_idx:] = beacon_attention_mask
+
+            # NOTE: the attention mask should be modified when there is beacon token within the window, not in the input_ids
+            elif self.config.beacon_pos == "interleave" and (self._beacon_indices == 1).any():
+                assert self.config.beacon_attend_prev == False, f"Make sure beacon_attend_prev is False if using 'interleave' beacon pos!"
+
+                beacon_indices = self._beacon_indices
+
+                cur_position_ids = position_ids[:, -len(beacon_indices):]
+                base_position = cur_position_ids[:, 0] - 1
+                # NOTE: alternate position so that the position of raw tokens are consistent
+                position_template = cur_position_ids.new_ones(cur_position_ids.shape)
+                position_template[:, compression_ratio + 1::compression_ratio + 1] = 0
+                cur_position_ids = base_position + position_template.cumsum(-1)
+                position_ids[:, -len(beacon_indices):] = cur_position_ids
+
+                cur_input_length = len(beacon_indices)
+                cur_attention_mask = attention_mask[..., -cur_input_length:, -cur_input_length:]
+                # mask all beacon columns
+                cur_attention_mask[..., beacon_indices] = min_value
+                # beacon tokens can attend to themselves
+                input_ids_attention_mask = cur_attention_mask[..., -tgt_size:, -tgt_size:]
+                input_ids_attention_mask[..., range(tgt_size), range(tgt_size)] = 0
+
+        elif self.config.beacon_attn == "segmentation":
+            # each beacon can attend to its corresponding sub-interval
+
+            if self.config.beacon_pos == "append" and beacon_size > 0:
+                window_size = self.config.beacon_window
+                window_size_with_beacon = window_size + beacon_size
+                beacon_start_idx = -beacon_size
+                # batch_size, head_num, window_size
+                reference_attention_mask = attention_mask[..., -beacon_size - 1, -window_size_with_beacon: -beacon_size]
+
+                # beacon_size, compression_ratio
+                indices = torch.arange(compression_ratio * beacon_size, device=device).view(beacon_size, -1)
+                # beacon_size, window_size
+                ordinal_attention_mask = attention_mask.new_full((beacon_size, window_size), min_value)
+                ordinal_attention_mask.scatter_(dim=-1, index=indices, value=0)
+
+                # NOTE: add reference attention_mask so that padding tokens are considered
+                ordinal_attention_mask = ordinal_attention_mask[None, None, ...] + reference_attention_mask.unsqueeze(-2)
+
+                if self.config.beacon_attend_prev:
+                    beacon_attention_mask = attention_mask.new_full((beacon_size, beacon_size), min_value).triu(1)
+                    # the beacon token is next to the last ordinal token it attends to
+                    beacon_position_ids = position_ids.new_full(beacon_size, fill_value=compression_ratio + mem_size)
+                    beacon_position_ids = beacon_position_ids + torch.arange(beacon_size)
+                    position_ids[:, beacon_start_idx:] = beacon_position_ids
+                else:
+                    beacon_attention_mask = attention_mask.new_full((beacon_size, beacon_size), min_value).fill_diagonal_(0)
+                    # the beacon token is next to the last ordinal token it attends to
+                    beacon_position_ids = position_ids.new_full(beacon_size, fill_value=compression_ratio + mem_size)
+                    position_ids[:, beacon_start_idx:] = beacon_position_ids
+
+                attention_mask[..., beacon_start_idx:, -window_size_with_beacon: -beacon_size] = ordinal_attention_mask
+                attention_mask[..., beacon_start_idx:, beacon_start_idx:] = beacon_attention_mask
+                # beacons of different ratios are blind to others
+                attention_mask[..., beacon_start_idx:, -beacon_size: beacon_start_idx] = min_value
+
+            elif self.config.beacon_pos == "interleave":
+                raise NotImplementedError
+
+        elif self.config.beacon_attn == "full-coverage":
+            pass
+
+        return attention_mask, position_ids
+
+    def _extract_beacon_and_raw_memory(
+        self, 
+        key, 
+        value, 
+        previous_beacon_key, 
+        previous_beacon_value, 
+        previous_raw_key, 
+        previous_raw_value, 
+        beacon_indices,
+    ):
+        """Extract beacon and raw memory from the returned key and value when the window is full."""
+        key = cat_tensor([
+            previous_raw_key, 
+            key
+        ], dim=self.k_seq_dim)
+        value = cat_tensor([
+            previous_raw_value, 
+            value
+        ], dim=self.v_seq_dim)
+
+        # NOTE: we use magic slice instead of boolean index here for efficiency
+        beacon_key = slice_tensor(key, index=torch.logical_or(beacon_indices == 1, beacon_indices == -1), dim=self.k_seq_dim)
+        beacon_key = cat_tensor([previous_beacon_key, beacon_key], dim=self.k_seq_dim)
+        beacon_value = slice_tensor(value, index=torch.logical_or(beacon_indices == 1, beacon_indices == -1), dim=self.v_seq_dim)
+        beacon_value = cat_tensor([previous_beacon_value, beacon_value], dim=self.v_seq_dim)
+
+        if self._raw_size_to_cache > 0:
+            raw_key = slice_tensor(key, index=beacon_indices == 0, dim=self.k_seq_dim)
+            raw_key = slice_tensor(raw_key, start=-raw_size_to_cache, dim=self.k_seq_dim)
+
+            raw_value = slice_tensor(value, index=beacon_indices == 0, dim=self.v_seq_dim)
+            raw_value = slice_tensor(raw_value, start=-raw_size_to_cache, dim=self.v_seq_dim)        
+
+        else:
+            raw_key = None
+            raw_value = None
+
+        return beacon_key, beacon_value, raw_key, raw_value
+
+
+def slice_tensor(x, start=None, end=None, step=None, index=None, dim=2):
+    if x is None:
+        return None
+    if end == 0:
+        return None
+    if start == x.shape[dim]:
+        return None
+    if start is not None and start == end:
+        return None
+    if dim == 2:
+        if index is not None:
+            return x[:, :, index]
+        elif start is None and end is not None:
+            if step is None:
+                return x[:, :, :end, ...]
+            else:
+                return x[:, :, :end:step, ...]
+        elif start is not None and end is None:
+            if step is None:
+                return x[:, :, start:, ...]
+            else:
+                return x[:, :, start::step, ...]
+        elif start is not None and end is not None:
+            if step is None:
+                return x[:, :, start:end, ...]
+            else:
+                return x[:, :, start:end:step, ...]
+    elif dim == 1:
+        if index is not None:
+            return x[:, :, index]
+        elif start is None and end is not None:
+            if step is None:
+                return x[:, :end, ...]
+            else:
+                return x[:, :end:step, ...]
+        elif start is not None and end is None:
+            if step is None:
+                return x[:, start:, ...]
+            else:
+                return x[:, start::step, ...]
+        elif start is not None and end is not None:
+            if step is None:
+                return x[:, start:end, ...]
+            else:
+                return x[:, start:end:step, ...]
+    else:
+        raise NotImplementedError
+
+def cat_tensor(list_of_tensors, dim=-1):
+    list_of_tensors = [t for t in list_of_tensors if t is not None]
+    if len(list_of_tensors) > 1:
+        result = torch.cat(list_of_tensors, dim=dim)
+    elif len(list_of_tensors) == 1:
+        result = list_of_tensors[0]
+    else:
+        result = None
+    return result
+
+def slice_activations(activations, start=None, end=None, k_seq_dim=2, v_seq_dim=2):
+    new_activations = []
+    for key, value in activations:
+        new_key = slice_tensor(key, start=start, end=end, dim=k_seq_dim)
+        new_value = slice_tensor(value, start=start, end=end, dim=v_seq_dim)
+        new_activations.append([new_key, new_value])
+    return new_activations
+
+def cat_activations(list_of_activations, k_seq_dim=2, v_seq_dim=2):
+    assert all(len(x) == len(list_of_activations[0]) for x in list_of_activations), f"Make sure all activations have the same number of layers! Found {[len(x) for x in list_of_activations]}."
+
+    new_activations = []
+    for layer_idx in range(len(list_of_activations[0])):
+        keys = [x[layer_idx][0] for x in list_of_activations]
+        values = [x[layer_idx][1] for x in list_of_activations]
+
+        new_key = cat_tensor(keys, dim=k_seq_dim)
+        new_value = cat_tensor(values, dim=v_seq_dim)
+        new_activations.append([new_key, new_value])
+    return new_activations
+
+def interleave_activations(main_activations, augment_activations, main_spans, augment_spans, k_seq_dim=2, v_seq_dim=2, device=torch.device("cuda")):
+    """ Interleave main_activations and augment_activations according to main_span and augment_span.
+
+    Args:
+        main_span: a list of tuples (start_idx, end_idx). when start_idx and end_idx is None, the augment_activations will be plugged in.
+        augment_span: a list of tuples (start_idx, end_idx)
+    """
+    assert len(main_activations) == len(augment_activations) , f"Make sure main and augment activations have the same number of layers! Found {len(main_activations)} and {len(augment_activations)}!"
+    assert sum(x[0] is None and x[1] is None for x in main_spans) == len(augment_spans), f"Make sure the number of slots for augmentation (start_idx=None and end_idx=None in main_spans) matches the number of augmentations. Found {sum(x for x in main_spans if x[0] is None and x[1] is None)} slots but {len(augment_spans)} augmentations!"
+
+    new_activations = []
+    for layer_idx in range(len(main_activations)):
+        main_key, main_value = main_activations[layer_idx]
+        augment_key, augment_value = augment_activations[layer_idx]
+
+        sliced_keys = []
+        sliced_values = []
+
+        augment_idx = 0
+        for start, end in main_spans:
+            if start is None and end is None:
+                # this means the augment key/value should be plugged in
+                augment_start, augment_end = augment_spans[augment_idx]
+                sliced_key = slice_tensor(
+                    augment_key, 
+                    start=augment_start, 
+                    end=augment_end,
+                    dim=k_seq_dim
+                ).to(device)
+                sliced_value = slice_tensor(
+                    augment_value, 
+                    start=augment_start, 
+                    end=augment_end,
+                    dim=v_seq_dim
+                ).to(device)
+
+            else:
+                sliced_key = slice_tensor(
+                    main_key, 
+                    start=start, 
+                    end=end,
+                    dim=k_seq_dim
+                )
+                sliced_value = slice_tensor(
+                    main_value, 
+                    start=start, 
+                    end=end,
+                    dim=v_seq_dim
+                )
+
+            sliced_keys.append(sliced_key)
+            sliced_values.append(sliced_value)
+
+        new_key = cat_tensor(sliced_keys, dim=k_seq_dim)
+        new_value = cat_tensor(sliced_values, dim=v_seq_dim)
+        new_activations.append([new_key, new_value])
+
+    return new_activations
+
+def softmax(x:np.ndarray, axis=-1, temperature=1):
+    if isinstance(x, list):
+        x = np.array(x)
+    x = x / temperature
+    x = x - x.max(axis=axis, keepdims=True)
+    y = np.exp(x)
+    return y / y.sum(axis=axis, keepdims=True)
+
+def l1_norm(x):
+    sum_x = sum(x)
+    x = [y/sum_x for y in x]
+    return x
\ No newline at end of file

Long_LLM/activation_beacon/new/src/modeling_retrieval.py

+import torch
+import faiss
+import numpy as np
+from typing import List, Mapping, Optional, Union
+from accelerate import Accelerator
+from dataclasses import dataclass
+from collections import defaultdict
+from transformers import AutoTokenizer, AutoModel
+from transformers.utils import logging
+from transformers.modeling_utils import PreTrainedModel
+from transformers.tokenization_utils import PreTrainedTokenizer
+from torch.utils.data import DataLoader
+from semantic_text_splitter import TextSplitter
+
+from src import apply_chat_template
+
+logger = logging.get_logger(__name__)
+
+
+
+class BM25Retriever:
+    def __init__(self, k1:float=0.9, b:float=0.4) -> None:
+        self.name = "bm25"
+        self.k1 = k1
+        self.b = b
+
+        self.remove_all()
+    
+    @property
+    def num_keys(self):
+        return self.N
+
+    def add(self, docs: List[Union[str, List[int]]], stop_tokens: set={}):
+        """Build in-memory BM25 index."""
+        for doc in docs:
+            if isinstance(doc, str):
+                doc = doc.split()
+            df = {}
+            tf = defaultdict(int)
+            for token in doc:
+                if token not in stop_tokens:
+                    tf[token] += 1
+                    df[token] = 1
+            self.tfs.append(dict(tf))
+            for token in df:
+                self.dfs[token] += 1
+                # store the doc offset in the inverted lists of the corresponding token
+                self.inverted_lists[token].append(self.N)
+
+            self.N += 1
+            self.doc_lengths.append(len(doc))
+    
+    def remove_all(self):
+        """Remove all keys from the index."""
+        self.dfs = defaultdict(float)
+        self.tfs = []
+        self.inverted_lists = defaultdict(list)
+        self.doc_lengths = []
+        self.N = 0
+
+    def search(self, queries: Union[str, List[int], List[str], List[List[int]]], hits: int=100, k1: Optional[float]=None, b: Optional[float]=None):
+        """Search over the BM25 index."""
+        if k1 is None:
+            k1 = self.k1
+        if b is None:
+            b = self.b
+        
+        hits = min(self.N, hits)
+        
+        global_scores = np.zeros(self.N, dtype=np.float32)
+        
+        if isinstance(queries, str):
+            queries = [queries]
+        elif isinstance(queries, list) and isinstance(queries[0], int):
+            queries = [queries]
+        
+        all_scores = np.zeros((len(queries), hits), dtype=np.float32)
+        all_indices = np.zeros((len(queries), hits), dtype=np.int64)
+
+        doc_lengths = np.array(self.doc_lengths)
+        
+        for i, query in enumerate(queries):
+            if isinstance(query, str):
+                query = query.split(" ")
+                # TODO: stem
+
+            for token in query:
+                if token in self.inverted_lists:
+                    candidates = self.inverted_lists[token]
+                else:
+                    continue
+
+                tfs = np.array([self.tfs[candidate][token] for candidate in candidates], dtype=np.float32)
+                df = self.dfs[token]
+                idf = np.log((self.N - df + 0.5) / (df + 0.5) + 1)
+
+                candidate_scores = idf * (k1 + 1) * tfs / (tfs + k1 * (1 - b + b * doc_lengths[candidates]))
+                global_scores[candidates] += candidate_scores
+
+            indice = np.argpartition(-global_scores, hits - 1)[:hits]
+            score = global_scores[indice]
+            
+            sorted_idx = np.argsort(score)[::-1]
+            indice = indice[sorted_idx]
+            score = score[sorted_idx]
+
+            invalid_pos = score == 0
+            indice[invalid_pos] = -1
+            score[invalid_pos] = -float('inf')
+
+            all_scores[i] = score
+            all_indices[i] = indice
+        return all_scores, all_indices
+
+
+
+class DenseRetriever:
+    def __init__(self, encoder:str='BAAI/bge-large-en', pooling_method:List[str]=["cls"], dense_metric:str="cos", query_max_length:int=1024, key_max_length:int=1024, dtype:str="fp16", cache_dir:Optional[str]=None) -> None:
+        self.name = encoder
+
+        self.pooling_method = pooling_method
+        self.dense_metric = dense_metric
+        self.query_max_length = query_max_length
+        self.key_max_length = key_max_length
+        logger.info(f"Loading tokenizer and model from {encoder}...")
+
+        if dtype == "bf16":
+            dtype = torch.bfloat16
+        elif dtype == "fp16":
+            dtype = torch.float16
+        else:
+            dtype = torch.float32
+
+        self.tokenizer = AutoTokenizer.from_pretrained(encoder, cache_dir=cache_dir)
+        self.encoder = AutoModel.from_pretrained(encoder, cache_dir=cache_dir, torch_dtype=dtype, device_map={'': "cuda"}).eval()
+
+        self.ndim = self.encoder.config.hidden_size
+        self._index = None
+
+    @property
+    def device(self):
+        return self.encoder.device
+
+    @property
+    def num_keys(self):
+        if self._index is not None:
+            return self._index.index.ntotal
+        else:
+            return 0
+
+    def _prepare(self, inputs: Union[str, List[str], Mapping], field="key"):
+        """Convert inputs into tokenized input_ids"""
+        if isinstance(inputs, str) or (isinstance(inputs, list) and isinstance(inputs[0], str)):
+            if field == "key":
+                inputs = self.tokenizer(
+                    inputs, return_tensors="pt", padding=True, truncation=True, max_length=self.key_max_length)
+                inputs = inputs.to(self.device)
+            elif field == "query":
+                inputs = self.tokenizer(
+                    inputs, return_tensors="pt", padding=True, truncation=True, max_length=self.query_max_length)
+                inputs = inputs.to(self.device)
+            else:
+                raise NotImplementedError
+        elif isinstance(inputs, Mapping) and "input_ids" in inputs:
+            if field == "key":
+                for k, v in inputs.items():
+                    inputs[k] = v[:, :self.key_max_length].to(self.device)
+            elif field == "query":
+                for k, v in inputs.items():
+                    inputs[k] = v[:, :self.query_max_length].to(self.device)
+            else:
+                raise NotImplementedError
+        else:
+            raise ValueError(f"Expected inputs of type str, list[str], or dict, got {type(inputs)}!")
+        return inputs
+
+    def _pool(self, embeddings, attention_mask):
+        if "mean" in self.pooling_method:
+            embeddings = embeddings.masked_fill(
+                ~attention_mask[..., None].bool(), 0.0)
+            embedding = embeddings.sum(
+                dim=1) / attention_mask.sum(dim=1, keepdim=True)
+        elif "cls" in self.pooling_method:
+            embedding = embeddings[:, 0]
+        else:
+            raise NotImplementedError(
+                f"Pooling_method {self.pooling_method} not implemented!")
+        return embedding
+
+    @torch.no_grad()
+    def encode(self, inputs: Union[str, List[str], Mapping], field:str="key"):
+        """Encode inputs into embeddings
+
+        Args:
+            inputs: can be string, list of strings, or BatchEncoding results from tokenizer
+
+        Returns:
+            Tensor: [batch_size, d_embed]
+        """
+        inputs = self._prepare(inputs, field=field)
+        encoder = self.encoder
+
+        embeddings = encoder(**inputs).last_hidden_state    # B, L, D
+        embedding = self._pool(embeddings, inputs["attention_mask"])
+        if self.dense_metric == "cos":
+            embedding = torch.nn.functional.normalize(embedding, p=2, dim=1)
+        return embedding
+
+    def remove_all(self):
+        """Remove all keys from the index."""
+        if self._index is not None:
+            self._index.index.reset()
+
+    @torch.no_grad()
+    def add(self, docs: List[str], index_factory:str="Flat", batch_size=500):
+        """Build faiss index.
+        
+        Args:
+            shard_across_devices: split the corpus onto all devices and encode them
+        """
+        if len(docs) == 0:
+            return
+
+        metric = self.dense_metric
+        doc_embeddings = np.zeros((len(docs), self.ndim), dtype=np.float32)
+
+        for i in range(0, len(docs), batch_size):
+            batch_docs = docs[i: i + batch_size]
+            embeddings = self.encode(batch_docs)    # batch_size, ndim
+            doc_embeddings[i: i + batch_size] = embeddings.cpu().numpy()
+
+        if self._index is None:
+            index = FaissIndex(self.device)
+            index.build(doc_embeddings, index_factory, metric)
+            self._index = index
+        else:
+            self._index.add(doc_embeddings)
+
+    @torch.no_grad()
+    def search(self, queries: Union[str, List[str]], hits:int=10):
+        assert self._index is not None, "Make sure there is an indexed corpus!"
+
+        embeddings = self.encode(queries, field="query").cpu().numpy().astype(np.float32, order="C")
+        scores, indices = self._index.search(embeddings, hits)
+        return scores, indices
+    
+
+
+class FaissIndex:
+    def __init__(self, device) -> None:
+        if isinstance(device, torch.device):
+            if device.index is None:
+                device = "cpu"
+            else:
+                device = device.index
+        self.device = device
+
+    def build(self, doc_embeddings, index_factory, metric):
+        if metric == "l2":
+            metric = faiss.METRIC_L2
+        elif metric in ["ip", "cos"]:
+            metric = faiss.METRIC_INNER_PRODUCT
+        else:
+            raise NotImplementedError(f"Metric {metric} not implemented!")
+        
+        index = faiss.index_factory(doc_embeddings.shape[1], index_factory, metric)
+        
+        if self.device != "cpu":
+            co = faiss.GpuClonerOptions()
+            co.useFloat16 = True
+            # logger.info("using fp16 on GPU...")
+            index = faiss.index_cpu_to_gpu(faiss.StandardGpuResources(), self.device, index, co)
+
+        index.train(doc_embeddings)
+        index.add(doc_embeddings)
+        self.index = index
+        return index
+    
+    def add(self, doc_embeddings):
+        self.index.add(doc_embeddings)
+
+    def load(self, index_path):
+        logger.info(f"loading index from {index_path}...")
+        index = faiss.read_index(index_path)
+        if self.device != "cpu":
+            co = faiss.GpuClonerOptions()
+            co.useFloat16 = True
+            index = faiss.index_cpu_to_gpu(faiss.StandardGpuResources(), self.device, index, co)
+        self.index = index
+        return index
+
+    def save(self, index_path):
+        logger.info(f"saving index at {index_path}...")
+        if isinstance(self.index, faiss.GpuIndex):
+            index = faiss.index_gpu_to_cpu(self.index)
+        else:
+            index = self.index
+        faiss.write_index(index, index_path)
+
+    def search(self, query, hits):
+        return self.index.search(query, k=hits)