new commits

README.md

@@ -11,3 +11,36 @@ license: apache-2.0
 ---
 
 # CelebChat
+CelebChat is a Hugging Face Space where the user can talk with nearly 50 virtual celebrities.
+
+## System details
+![alt text](https://github.com/liuhaozhe6788/CelebChat/blob/main/img/flow_chart.jpg?raw=true)
+
+## Citation
+```
+@article{bertsch2023unlimiformer,
+  title={Unlimiformer: Long-Range Transformers with Unlimited Length Input},
+  author={Bertsch, Amanda and Alon, Uri and Neubig, Graham and Gormley, Matthew R},
+  journal={arXiv preprint arXiv:2305.01625},
+  year={2023}
+}
+
+@misc{https://doi.org/10.48550/arxiv.2210.11416,
+  doi = {10.48550/ARXIV.2210.11416},
+  
+  url = {https://arxiv.org/abs/2210.11416},
+  
+  author = {Chung, Hyung Won and Hou, Le and Longpre, Shayne and Zoph, Barret and Tay, Yi and Fedus, William and Li, Eric and Wang, Xuezhi and Dehghani, Mostafa and Brahma, Siddhartha and Webson, Albert and Gu, Shixiang Shane and Dai, Zhuyun and Suzgun, Mirac and Chen, Xinyun and Chowdhery, Aakanksha and Narang, Sharan and Mishra, Gaurav and Yu, Adams and Zhao, Vincent and Huang, Yanping and Dai, Andrew and Yu, Hongkun and Petrov, Slav and Chi, Ed H. and Dean, Jeff and Devlin, Jacob and Roberts, Adam and Zhou, Denny and Le, Quoc V. and Wei, Jason},
+  
+  keywords = {Machine Learning (cs.LG), Computation and Language (cs.CL), FOS: Computer and information sciences, FOS: Computer and information sciences},
+  
+  title = {Scaling Instruction-Finetuned Language Models},
+  
+  publisher = {arXiv},
+  
+  year = {2022},
+  
+  copyright = {Creative Commons Attribution 4.0 International}
+}
+
+```

app.py

@@ -64,7 +64,7 @@ def main():
     st.session_state["celeb_bot"] = CelebBot(st.session_state["celeb_name"], 
                     celeb_gender,
                     get_tokenizer(st.session_state["QA_model_path"]), 
-                    get_seq2seq_model(st.session_state["QA_model_path"]) if "flan-t5" in st.session_state["QA_model_path"] else get_causal_model(st.session_state["QA_model_path"]), 
+                    get_seq2seq_model(st.session_state["QA_model_path"], _tokenizer=get_tokenizer(st.session_state["QA_model_path"])) if "flan-t5" in st.session_state["QA_model_path"] else get_causal_model(st.session_state["QA_model_path"]), 
                     get_tokenizer(st.session_state["sentTr_model_path"]), 
                     get_auto_model(st.session_state["sentTr_model_path"]), 
                     *preprocess_text(st.session_state["celeb_name"], knowledge, "en_core_web_lg")

celebbot.py

@@ -2,21 +2,17 @@ import datetime
 import numpy as np
 import torch
 import torch.nn.functional as F
-import os 
-import json
 import speech_recognition as sr
 import re
 import time
-import spacy
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModel
 import pickle
-import streamlit as st
 from sklearn.metrics.pairwise import cosine_similarity
 
 
 # Build the AI
 class CelebBot():
-    def __init__(self, name, gender, QA_tokenizer, QA_model, sentTr_tokenizer, sentTr_model, spacy_model, knowledge_sents):
+
+    def __init__(self, name, gender, QA_tokenizer, QA_model, sentTr_tokenizer, sentTr_model, spacy_model, knowledge_sents, top_k = 8):
         self.name = name
         self.gender = gender
         print("--- starting up", self.name, self.gender, "---")
@@ -29,6 +25,7 @@ class CelebBot():
         self.spacy_model = spacy_model
 
         self.all_knowledge = knowledge_sents
+        self.top_k = top_k
 
     def speech_to_text(self):
         recognizer = sr.Recognizer()
@@ -83,7 +80,7 @@ class CelebBot():
         all_knowledge_embeddings = self.sentence_embeds_inference(self.all_knowledge)
         similarity = cosine_similarity(all_knowledge_embeddings.cpu(), question_embeddings.cpu())
         similarity = np.reshape(similarity, (1, -1))[0]
-        K = min(8, len(self.all_knowledge))
+        K = min(self.top_k, len(self.all_knowledge))
         top_K = np.sort(np.argpartition(similarity, -K)[-K: ])
         all_knowledge_assertions = np.array(self.all_knowledge)[top_K]
 
@@ -175,7 +172,7 @@ class CelebBot():
               knowledge = self.retrieve_knowledge_assertions()
 
             query = f"Context: {instruction} {knowledge}\n\nChat History: {chat_his}Question: {self.text}\n\nAnswer:"
-            input_ids = self.QA_tokenizer(f"{query}", return_tensors="pt").input_ids.to(self.QA_model.device)
+            input_ids = self.QA_tokenizer(f"{query}", truncation=False, return_tensors="pt").input_ids.to(self.QA_model.device)
             outputs = self.QA_model.generate(input_ids, max_length=1024, min_length=8, repetition_penalty=2.5)
             self.text = self.QA_tokenizer.decode(outputs[0], skip_special_tokens=True)
         return self.text

requirements.txt

@@ -30,4 +30,5 @@ sentence-transformers==2.2.2
 evaluate==0.4.1
 https://huggingface.co/spacy/en_core_web_lg/resolve/main/en_core_web_lg-any-py3-none-any.whl
 protobuf==3.20
-streamlit_mic_recorder==0.0.2
+streamlit_mic_recorder==0.0.2
+faiss-cpu==1.7.4

run_cli.py

@@ -5,7 +5,7 @@ import spacy
 from celebbot import CelebBot
 from utils import *
 
-DEBUG = False
+DEBUG = True
 QA_MODEL_ID = "google/flan-t5-large"
 SENTTR_MODEL_ID = "sentence-transformers/all-mpnet-base-v2"
 
@@ -52,6 +52,7 @@ def main():
             print("me --> ", ai.text)
 
             answers.append(ai.question_answer())
+            print("bot --> ", ai.text)
 
             if not DEBUG:
                 ai.text_to_speech()

run_eval.py

@@ -4,15 +4,19 @@ import spacy
 import json
 import evaluate
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModel
+from unlimiformer import Unlimiformer, UnlimiformerArguments
 import torch 
 
 from utils import *
 from celebbot import CelebBot
 
-QA_MODEL_ID = "google/flan-t5-large"
+QA_MODEL_ID = "google/flan-t5-xl"
 SENTTR_MODEL_ID = "sentence-transformers/all-mpnet-base-v2"
 celeb_names = ["Cate Blanchett", "David Beckham", "Emma Watson", "Lady Gaga", "Madonna", "Mark Zuckerberg"]
 
+USE_UNLIMIFORMER = True
+TOP_K = 8
+
 celeb_data = get_celeb_data("data.json")
 references = [val['answers'] for key, val in list(celeb_data.items()) if key in celeb_names]
 references = list(itertools.chain.from_iterable(references))
@@ -20,7 +24,29 @@ predictions = []
 
 device = 'cpu'
 QA_tokenizer = AutoTokenizer.from_pretrained(QA_MODEL_ID)
-QA_model = AutoModelForSeq2SeqLM.from_pretrained(QA_MODEL_ID).to(device)
+QA_model = AutoModelForSeq2SeqLM.from_pretrained(QA_MODEL_ID)
+if USE_UNLIMIFORMER:
+    defaults = UnlimiformerArguments()
+    unlimiformer_kwargs = {
+        'layer_begin': defaults.layer_begin, 
+        'layer_end': defaults.layer_end,
+        'unlimiformer_head_num': defaults.unlimiformer_head_num, 
+        'exclude_attention': defaults.unlimiformer_exclude, 
+        'chunk_overlap': defaults.unlimiformer_chunk_overlap,
+        'model_encoder_max_len': defaults.unlimiformer_chunk_size,
+        'verbose': defaults.unlimiformer_verbose, 'tokenizer': QA_tokenizer,
+        'unlimiformer_training': defaults.unlimiformer_training,
+        'use_datastore': defaults.use_datastore,
+        'flat_index': defaults.flat_index,
+        'test_datastore': defaults.test_datastore,
+        'reconstruct_embeddings': defaults.reconstruct_embeddings,
+        'gpu_datastore': defaults.gpu_datastore,
+        'gpu_index': defaults.gpu_index
+    }
+    QA_model =Unlimiformer.convert_model(QA_model, **unlimiformer_kwargs).to(device)
+else:
+    QA_model = QA_model.to(device)
+
 sentTr_tokenizer = AutoTokenizer.from_pretrained(SENTTR_MODEL_ID)
 sentTr_model = AutoModel.from_pretrained(SENTTR_MODEL_ID).to(device)
 
@@ -37,7 +63,7 @@ for celeb_name in celeb_names:
     spacy_model = spacy.load("en_core_web_lg")
     knowledge_sents = [i.text.strip() for i in spacy_model(knowledge).sents]
 
-    ai = CelebBot(celeb_name, gender, QA_tokenizer, QA_model, sentTr_tokenizer, sentTr_model, spacy_model, knowledge_sents)
+    ai = CelebBot(celeb_name, gender, QA_tokenizer, QA_model, sentTr_tokenizer, sentTr_model, spacy_model, knowledge_sents, top_k=TOP_K)
     for q in celeb_data[celeb_name]["questions"]:
         ai.text = q
         response = ai.question_answer()

unlimiformer/__init__.py

@@ -0,0 +1,2 @@
+from .model import Unlimiformer
+from .usage import UnlimiformerArguments

unlimiformer/configs/data/contract_nli.json

@@ -0,0 +1,12 @@
+{
+  "dataset_name": "tau/sled",
+  "dataset_config_name": "contract_nli",
+  "max_source_length": 16384,
+  "max_prefix_length": 64,
+  "pad_prefix": true,
+  "generation_max_length": 8,
+  "num_train_epochs": 20,
+  "metric_names": ["exact_match"],
+  "metric_for_best_model": "exact_match",
+  "greater_is_better": true
+}

unlimiformer/configs/data/gov_report.json

@@ -0,0 +1,12 @@
+{
+  "dataset_name": "tau/sled",
+  "dataset_config_name": "gov_report",
+  "max_source_length": 16384,
+  "generation_max_length": 1024,
+  "max_prefix_length": 0,
+  "pad_prefix": false,
+  "num_train_epochs": 10,
+  "metric_names": ["rouge"],
+  "metric_for_best_model": "rouge/geometric_mean",
+  "greater_is_better": true
+}

unlimiformer/configs/data/hotpotqa.json

@@ -0,0 +1,12 @@
+{
+  "dataset_name": "tau/sled",
+  "dataset_config_name": "hotpotqa",
+  "max_source_length": 16384,
+  "max_prefix_length": 64,
+  "pad_prefix": true,
+  "generation_max_length": 128,
+  "num_train_epochs": 9,
+  "metric_names": ["f1", "exact_match"],
+  "metric_for_best_model": "f1",
+  "greater_is_better": true
+}

unlimiformer/configs/data/hotpotqa_second_only.json

@@ -0,0 +1,12 @@
+{
+  "dataset_name": "tau/sled",
+  "dataset_config_name": "hotpotqa_second_only",
+  "max_source_length": 16384,
+  "max_prefix_length": 64,
+  "pad_prefix": true,
+  "generation_max_length": 128,
+  "num_train_epochs": 9,
+  "metric_names": ["f1", "exact_match"],
+  "metric_for_best_model": "f1",
+  "greater_is_better": true
+}

unlimiformer/configs/data/narative_qa.json

@@ -0,0 +1,12 @@
+{
+  "dataset_name": "tau/sled",
+  "dataset_config_name": "narrative_qa",
+  "max_source_length": 16384,
+  "max_prefix_length": 64,
+  "pad_prefix": true,
+  "num_train_epochs": 2,
+  "generation_max_length": 128,
+  "metric_names": ["f1"],
+  "metric_for_best_model": "f1",
+  "greater_is_better": true
+}

unlimiformer/configs/data/qasper.json

@@ -0,0 +1,12 @@
+{
+  "dataset_name": "tau/sled",
+  "dataset_config_name": "qasper",
+  "max_source_length": 16384,
+  "max_prefix_length": 64,
+  "pad_prefix": true,
+  "generation_max_length": 128,
+  "num_train_epochs": 20,
+  "metric_names": ["f1"],
+  "metric_for_best_model": "f1",
+  "greater_is_better": true
+}

unlimiformer/configs/data/qmsum.json

@@ -0,0 +1,12 @@
+{
+  "dataset_name": "tau/sled",
+  "dataset_config_name": "qmsum",
+  "max_source_length": 16384,
+  "max_prefix_length": 64,
+  "pad_prefix": true,
+  "num_train_epochs": 20,
+  "generation_max_length": 1024,
+  "metric_names": ["rouge"],
+  "metric_for_best_model": "rouge/geometric_mean",
+  "greater_is_better": true
+}

unlimiformer/configs/data/quality.json

@@ -0,0 +1,12 @@
+{
+  "dataset_name": "tau/sled",
+  "dataset_config_name": "quality",
+  "max_source_length": 16384,
+  "max_prefix_length": 160,
+  "pad_prefix": true,
+  "num_train_epochs": 20,
+  "generation_max_length": 128,
+  "metric_names": ["exact_match"],
+  "metric_for_best_model":  "exact_match",
+  "greater_is_better": true
+}

unlimiformer/configs/data/squad.json

@@ -0,0 +1,12 @@
+{
+  "dataset_name": "tau/sled",
+  "dataset_config_name": "squad",
+  "max_source_length": 16384,
+  "max_prefix_length": 64,
+  "pad_prefix": true,
+  "num_train_epochs": 3,
+  "generation_max_length": 128,
+  "metric_names": ["f1", "exact_match"],
+  "metric_for_best_model": "f1",
+  "greater_is_better": true
+}

unlimiformer/configs/data/squad_ordered_distractors.json

@@ -0,0 +1,12 @@
+{
+  "dataset_name": "tau/sled",
+  "dataset_config_name": "squad_ordered_distractors",
+  "max_source_length": 16384,
+  "max_prefix_length": 64,
+  "pad_prefix": true,
+  "num_train_epochs": 3,
+  "generation_max_length": 128,
+  "metric_names": ["f1", "exact_match"],
+  "metric_for_best_model": "f1",
+  "greater_is_better": true
+}

unlimiformer/configs/data/squad_shuffled_distractors.json

@@ -0,0 +1,12 @@
+{
+  "dataset_name": "tau/sled",
+  "dataset_config_name": "squad_shuffled_distractors",
+  "max_source_length": 16384,
+  "max_prefix_length": 64,
+  "pad_prefix": true,
+  "num_train_epochs": 3,
+  "generation_max_length": 128,
+  "metric_names": ["f1", "exact_match"],
+  "metric_for_best_model": "f1",
+  "greater_is_better": true
+}

unlimiformer/configs/data/summ_screen_fd.json

@@ -0,0 +1,12 @@
+{
+  "dataset_name": "tau/sled",
+  "dataset_config_name": "summ_screen_fd",
+  "max_source_length": 16384,
+  "max_prefix_length": 0,
+  "pad_prefix": false,
+  "num_train_epochs": 10,
+  "generation_max_length": 1024,
+  "metric_names": ["rouge"],
+  "metric_for_best_model": "rouge/geometric_mean",
+  "greater_is_better": true
+}

unlimiformer/configs/model/bart_base_sled.json

@@ -0,0 +1,6 @@
+{
+  "model_name_or_path": "tau/bart-base-sled",
+  "use_auth_token": false,
+  "max_target_length": 1024,
+  "fp16": true
+}

unlimiformer/configs/model/bart_large_sled.json

@@ -0,0 +1,6 @@
+{
+  "model_name_or_path": "tau/bart-large-sled",
+  "use_auth_token": false,
+  "max_target_length": 1024,
+  "fp16": true
+}

unlimiformer/configs/model/primera_govreport_sled.json

@@ -0,0 +1,9 @@
+{
+    "model_type": "tau/sled",
+    "underlying_config": "allenai/PRIMERA",
+    "context_size": 4096,
+    "window_fraction": 0.5,
+    "prepend_prefix": true,
+    "encode_prefix": true,
+    "sliding_method": "dynamic"
+  }

unlimiformer/configs/training/base_training_args.json

@@ -0,0 +1,22 @@
+{
+    "eval_steps_override": 0.5,
+    "save_steps_override": 0.5,
+    "evaluation_strategy": "steps",
+    "eval_fraction": 1000,
+    "predict_with_generate": true,
+    "gradient_checkpointing": true,
+    "do_train": true,
+    "do_eval": true,
+    "seed": 42,
+    "warmup_ratio": 0.1,
+    "save_total_limit": 2,
+    "preprocessing_num_workers": 1,
+    "load_best_model_at_end": true,
+    "lr_scheduler": "linear",   
+    "adam_epsilon": 1e-6,
+    "adam_beta1": 0.9,
+    "adam_beta2": 0.98,
+    "weight_decay": 0.001,
+    "patience": 10,
+    "extra_metrics": "bertscore"
+}

unlimiformer/index_building.py

@@ -0,0 +1,161 @@
+import faiss
+import faiss.contrib.torch_utils
+import time
+import logging
+
+import torch
+import numpy as np
+
+code_size = 64
+
+class DatastoreBatch():
+    def __init__(self, dim, batch_size, flat_index=False, gpu_index=False, verbose=False, index_device=None) -> None:
+        self.indices = []
+        self.batch_size = batch_size
+        self.device = index_device if index_device is not None else torch.device('cuda' if gpu_index else 'cpu')
+        for i in range(batch_size):
+            self.indices.append(Datastore(dim, use_flat_index=flat_index, gpu_index=gpu_index, verbose=verbose, device=self.device))
+    
+    def move_to_gpu(self):
+        for i in range(self.batch_size):
+            self.indices[i].move_to_gpu()
+
+    def add_keys(self, keys, num_keys_to_add_at_a_time=100000):
+        for i in range(self.batch_size):
+            self.indices[i].add_keys(keys[i], num_keys_to_add_at_a_time)
+        
+    def train_index(self, keys):
+        for index, example_keys in zip(self.indices, keys):
+            index.train_index(example_keys)
+    
+    def search(self, queries, k):
+        found_scores, found_values = [], []
+        for i in range(self.batch_size):
+            scores, values = self.indices[i].search(queries[i], k)
+            found_scores.append(scores)
+            found_values.append(values)
+        return torch.stack(found_scores, dim=0), torch.stack(found_values, dim=0)
+
+    def search_and_reconstruct(self, queries, k):
+        found_scores, found_values = [], []
+        found_vectors = []
+        for i in range(self.batch_size):
+            scores, values, vectors = self.indices[i].search_and_reconstruct(queries[i], k)
+            found_scores.append(scores)
+            found_values.append(values)
+            found_vectors.append(vectors)     
+        return torch.stack(found_scores, dim=0), torch.stack(found_values, dim=0), torch.stack(found_vectors, dim=0)
+
+class Datastore():
+    def __init__(self, dim, use_flat_index=False, gpu_index=False, verbose=False, device=None) -> None:
+        self.dimension = dim
+        self.device = device if device is not None else torch.device('cuda' if gpu_index else 'cpu')
+        self.logger = logging.getLogger('index_building')
+        self.logger.setLevel(20)
+        self.use_flat_index = use_flat_index
+        self.gpu_index = gpu_index
+
+        # Initialize faiss index
+        # TODO: is preprocessing efficient enough to spend time on?
+        if not use_flat_index:
+            self.index = faiss.IndexFlatIP(self.dimension) # inner product index because we use IP attention
+        
+        # need to wrap in index ID map to enable add_with_ids 
+        # self.index = faiss.IndexIDMap(self.index) 
+
+        self.index_size = 0
+        # if self.gpu_index:
+        #     self.move_to_gpu()
+        
+    def move_to_gpu(self):
+        if self.use_flat_index:
+            # self.keys = self.keys.to(self.device)
+            return
+        else:
+            co = faiss.GpuClonerOptions()
+            co.useFloat16 = True
+            self.index = faiss.index_cpu_to_gpu(faiss.StandardGpuResources(), self.device.index, self.index, co)
+    
+    def train_index(self, keys):
+        if self.use_flat_index:
+            self.add_keys(keys=keys, index_is_trained=True)
+        else:
+            keys = keys.cpu().float()
+            ncentroids = int(keys.shape[0] / 128)
+            self.index = faiss.IndexIVFPQ(self.index, self.dimension,
+                ncentroids, code_size, 8)
+            self.index.nprobe = min(32, ncentroids)
+            # if not self.gpu_index:
+            #     keys = keys.cpu()
+
+            self.logger.info('Training index')
+            start_time = time.time()
+            self.index.train(keys)
+            self.logger.info(f'Training took {time.time() - start_time} s')
+            self.add_keys(keys=keys, index_is_trained=True)
+            # self.keys = None
+            if self.gpu_index:
+                self.move_to_gpu()
+
+    def add_keys(self, keys, num_keys_to_add_at_a_time=1000000, index_is_trained=False):
+        self.keys = keys
+        if not self.use_flat_index and index_is_trained:
+            start = 0
+            while start < keys.shape[0]:
+                end = min(len(keys), start + num_keys_to_add_at_a_time)
+                to_add = keys[start:end]
+                # if not self.gpu_index:
+                #     to_add = to_add.cpu()
+                # self.index.add_with_ids(to_add, torch.arange(start+self.index_size, end+self.index_size))
+                self.index.add(to_add)
+                self.index_size += end - start
+                start += end
+                if (start % 1000000) == 0:
+                    self.logger.info(f'Added {start} tokens so far')
+        # else:
+        #     self.keys.append(keys)
+
+        # self.logger.info(f'Adding total {start} keys')
+        # self.logger.info(f'Adding took {time.time() - start_time} s')
+
+    def search_and_reconstruct(self, queries, k):
+        if len(queries.shape) == 1: # searching for only 1 vector, add one extra dim
+            self.logger.info("Searching for a single vector; unsqueezing")
+            queries = queries.unsqueeze(0)
+        # self.logger.info("Searching with reconstruct")
+        assert queries.shape[-1] == self.dimension # query vectors are same shape as "key" vectors
+        scores, values, vectors = self.index.index.search_and_reconstruct(queries.cpu().detach(), k)
+        # self.logger.info("Searching done")
+        return scores, values, vectors
+    
+    def search(self, queries, k):
+        # model_device = queries.device
+        # model_dtype = queries.dtype
+        if len(queries.shape) == 1: # searching for only 1 vector, add one extra dim
+            self.logger.info("Searching for a single vector; unsqueezing")
+            queries = queries.unsqueeze(0)
+        assert queries.shape[-1] == self.dimension # query vectors are same shape as "key" vectors
+        # if not self.gpu_index:
+        #     queries = queries.cpu()
+        # else:
+        #     queries = queries.to(self.device)
+        if self.use_flat_index:
+            if self.gpu_index:
+                scores, values = faiss.knn_gpu(faiss.StandardGpuResources(), queries, self.keys, k, 
+                    metric=faiss.METRIC_INNER_PRODUCT, device=self.device.index)
+            else:
+                scores, values = faiss.knn(queries, self.keys, k, metric=faiss.METRIC_INNER_PRODUCT)
+                scores = torch.from_numpy(scores).to(queries.dtype)
+                values = torch.from_numpy(values) #.to(model_dtype)
+        else:
+            scores, values = self.index.search(queries.float(), k)
+        
+        # avoid returning -1 as a value
+        # TODO: get a handle on the attention mask and mask the values that were -1
+        values = torch.where(torch.logical_or(values < 0, values >= self.keys.shape[0]), torch.zeros_like(values), values)
+        # self.logger.info("Searching done")
+        # return scores.to(model_dtype).to(model_device), values.to(model_device)
+        return scores, values
+
+    
+    

unlimiformer/metrics/__init__.py

@@ -0,0 +1 @@
+from .metrics import load_metric, download_metric

unlimiformer/metrics/metrics.py

@@ -0,0 +1,182 @@
+from typing import List, Dict
+import os
+import importlib
+from abc import ABC, abstractmethod
+import inspect
+import shutil
+
+import numpy as np
+
+from utils.decoding import decode
+from datasets import load_metric as hf_load_metric
+from huggingface_hub import hf_hub_download
+
+
+class Metric(ABC):
+    def __init__(self, **kwargs) -> None:
+        super().__init__()
+        self._kwargs = kwargs
+
+        self.prefix = os.path.splitext(os.path.basename(inspect.getfile(self.__class__)))[0]
+        self.requires_decoded = False
+
+    def __call__(self, id_to_pred, id_to_labels, is_decoded=False):
+        if self.requires_decoded and is_decoded is False:
+            id_to_pred = self._decode(id_to_pred)
+            id_to_labels = self._decode(id_to_labels)
+        return self._compute_metrics(id_to_pred, id_to_labels)
+
+    @abstractmethod
+    def _compute_metrics(self, id_to_pred, id_to_labels) -> Dict[str, float]:
+        return
+
+    def _decode(self, id_to_something):
+        tokenizer = self._kwargs.get("tokenizer")
+        data_args = self._kwargs.get("data_args")
+        return decode(id_to_something, tokenizer, data_args)
+
+
+class MetricCollection(Metric):
+    def __init__(self, metrics: List[Metric], **kwargs):
+        super().__init__(**kwargs)
+        self._metrics = metrics
+
+    def __call__(self, id_to_pred, id_to_labels):
+        return self._compute_metrics(id_to_pred, id_to_labels)
+
+    def _compute_metrics(self, id_to_pred, id_to_labels):
+        results = {}
+
+        id_to_pred_decoded = None
+        id_to_labels_decoded = None
+        for metric in self._metrics:
+            metric_prefix = f"{metric.prefix}/" if metric.prefix else ""
+            if metric.requires_decoded:
+                if id_to_pred_decoded is None:
+                    id_to_pred_decoded = self._decode(id_to_pred)
+                if id_to_labels_decoded is None:
+                    id_to_labels_decoded = self._decode(id_to_labels)
+
+                result = metric(id_to_pred_decoded, id_to_labels_decoded, is_decoded=True)
+            else:
+                result = metric(id_to_pred, id_to_labels)
+
+            results.update({f"{metric_prefix}{k}": np.mean(v) if type(v) is list else v for k, v in result.items() if type(v) is not str})
+
+        results["num_predicted"] = len(id_to_pred)
+        results["mean_prediction_length_characters"] = np.mean([len(pred) for pred in id_to_pred_decoded.values()])
+
+        elem = next(iter(id_to_pred.values()))
+        if not ((isinstance(elem, list) and isinstance(elem[0], str)) or isinstance(elem, str)):
+            tokenizer = self._kwargs["tokenizer"]
+            results["mean_prediction_length_tokens"] = np.mean(
+                [np.count_nonzero(np.array(pred) != tokenizer.pad_token_id) for pred in id_to_pred.values()]
+            )  # includes BOS/EOS tokens
+
+        results = {key: round(value, 4) for key, value in results.items()}
+        return results
+
+
+def load_metric(paths: List[str], **kwargs):
+    if paths is None or len(paths) == 0:
+        return None
+    if isinstance(paths, str):
+        paths = [paths]
+    else:
+        paths = [path for path in paths]
+
+    metric_cls_list = []
+
+    scrolls_custom_metrics = []
+    to_remove = []
+    for i, path in enumerate(paths):
+        if not os.path.isfile(path):
+            scrolls_custom_metrics.append(path)
+            to_remove.append(i)
+    for i in sorted(to_remove, reverse=True):
+        del paths[i]
+    if len(scrolls_custom_metrics) > 0:
+        scrolls_custom_metrics.insert(0, "")  # In order to have an identifying comma in the beginning
+        metric_cls_list.append(ScrollsWrapper(",".join(scrolls_custom_metrics), **kwargs))
+
+    for path in paths:
+        path = path.strip()
+        if len(path) == 0:
+            continue
+        if os.path.isfile(path) is False:
+            path = os.path.join("src", "metrics", f"{path}.py")
+
+        module = path[:-3].replace(os.sep, ".")
+
+        metric_cls = import_main_class(module)
+        metric_cls_list.append(metric_cls(**kwargs))
+
+    return MetricCollection(metric_cls_list, **kwargs)
+
+
+# Modified from datasets.load
+def import_main_class(module_path):
+    """Import a module at module_path and return its main class"""
+    module = importlib.import_module(module_path)
+
+    main_cls_type = Metric
+
+    # Find the main class in our imported module
+    module_main_cls = None
+    for name, obj in module.__dict__.items():
+        if isinstance(obj, type) and issubclass(obj, main_cls_type):
+            if inspect.isabstract(obj):
+                continue
+            module_main_cls = obj
+            break
+
+    return module_main_cls
+
+
+class ScrollsWrapper(Metric):
+    def __init__(self, comma_separated_metric_names, **kwargs) -> None:
+        super().__init__(**kwargs)
+        self.prefix = None
+
+        self._metric = hf_load_metric(download_metric(), comma_separated_metric_names, keep_in_memory=True)
+
+        self.requires_decoded = True
+
+    def _compute_metrics(self, id_to_pred, id_to_labels) -> Dict[str, float]:
+        return self._metric.compute(**self._metric.convert_from_map_format(id_to_pred, id_to_labels))
+
+class HFMetricWrapper(Metric):
+    def __init__(self, metric_name, **kwargs) -> None:
+        super().__init__(**kwargs)
+        self._metric = hf_load_metric(metric_name)
+        self.kwargs = HFMetricWrapper.metric_specific_kwargs.get(metric_name, {})
+        self.requires_decoded = True
+        self.prefix = metric_name
+        self.requires_decoded = True
+
+    def _compute_metrics(self, id_to_pred, id_to_labels) -> Dict[str, float]:
+        return self._metric.compute(**self.convert_from_map_format(id_to_pred, id_to_labels), **self.kwargs)
+
+    def convert_from_map_format(self, id_to_pred, id_to_labels):
+        index_to_id = list(id_to_pred.keys())
+        predictions = [id_to_pred[id_] for id_ in index_to_id]
+        references = [id_to_labels[id_] for id_ in index_to_id]
+        return {"predictions": predictions, "references": references}
+
+    metric_specific_kwargs = {
+        'bertscore': {
+            # 'model_type': 'microsoft/deberta-large-mnli' or the larger 'microsoft/deberta-xlarge-mnli'
+            'model_type': 'facebook/bart-large-mnli', # has context window of 1024,
+            'num_layers': 11 # according to: https://docs.google.com/spreadsheets/d/1RKOVpselB98Nnh_EOC4A2BYn8_201tmPODpNWu4w7xI/edit#gid=0
+        }
+    }
+
+
+def download_metric():
+    # here we load the custom metrics
+    scrolls_metric_path = hf_hub_download(repo_id="tau/scrolls", filename="metrics/scrolls.py", repo_type='dataset')
+    updated_scrolls_metric_path = (
+        os.path.dirname(scrolls_metric_path) + os.path.basename(scrolls_metric_path).replace(".", "_") + ".py"
+    )
+    shutil.copy(scrolls_metric_path, updated_scrolls_metric_path)
+    return updated_scrolls_metric_path

unlimiformer/model.py

@@ -0,0 +1,1157 @@
+import logging
+import numpy as np
+import torch
+from torch import nn
+from enum import Enum, auto
+from transformers import BartModel, BartForConditionalGeneration, \
+    T5Model, T5ForConditionalGeneration, \
+    LEDModel, LEDForConditionalGeneration, \
+    AutoModelForCausalLM, AutoModelForSeq2SeqLM, \
+    MODEL_WITH_LM_HEAD_MAPPING, MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING
+
+from typing import TypeVar, Generic
+
+from .index_building import Datastore, DatastoreBatch
+
+logger = logging.getLogger('Unlimiformer')
+logger.setLevel(20)
+
+ModelType = TypeVar('ModelType')
+class Unlimiformer(Generic[ModelType]):
+    def __init__(self, model: ModelType, 
+            layer_begin=-1, layer_end=None,
+            unlimiformer_head_num=None, 
+            exclude_attention=False, 
+            model_encoder_max_len=None,
+            chunk_overlap=0,
+            verbose=False, save_heatmap=False, 
+            tokenizer=None, unlimiformer_training=False,
+            use_datastore=False, 
+            flat_index=False,
+            test_datastore=False, reconstruct_embeddings=False, 
+            gpu_datastore=False, gpu_index=False,
+            index_devices=(0,), datastore_device=0,
+            ):
+        super().__init__()
+        self.model = model
+        model.unlimiformer = self
+        self.layer_begin = layer_begin
+        self.layer_end = layer_end
+        self.specific_head = unlimiformer_head_num
+        self.exclude_attention = exclude_attention
+        self.actual_model_window_size = None
+        self.model_encoder_max_len = model_encoder_max_len
+        self.chunk_overlap = chunk_overlap
+        self.verbose = verbose
+        self.save_heatmap = save_heatmap
+        self.tokenizer = tokenizer
+        self.unlimiformer_training = unlimiformer_training
+
+        self.use_datastore = use_datastore
+        self.flat_index = flat_index
+        self.reconstruct_embeddings = reconstruct_embeddings
+        self.gpu_datastore = gpu_datastore
+        self.gpu_index = gpu_index
+        # if torch.cuda.is_available() and gpu_index:
+        #     self.index_devices = [torch.device(f'cuda:{i}') for i in index_devices]
+        # else:
+        self.index_devices = [torch.device('cpu')]
+        self.datastore_device = torch.device('cpu')
+        self.test_datastore = test_datastore # flag for debugging
+
+        self.device = torch.device('cpu')
+        self.activation_capturer = None
+        self.is_encoder_decoder = model.config.is_encoder_decoder
+        self.hook_handles = []
+        self.is_input_encoding_pass = False
+        self.is_first_test_decoding_step = False
+        self.prev_tokens = None
+        self.last_beam_idx = None
+        self.heatmap = None
+        self.cur_decoder_layer_index = None
+        self.datastore = None
+
+        self.break_into(model)
+
+    def break_into(self, model):
+        self.actual_model_window_size = self.window_size()
+        if self.model_encoder_max_len is None:
+            self.model_encoder_max_len = self.actual_model_window_size
+        self.window_margin = int(self.model_encoder_max_len * self.chunk_overlap / 2)
+        self.num_heads = model.config.num_attention_heads
+        if self.specific_head is None:
+            self.head_nums = Ellipsis # torch.arange(0, self.num_heads, device=self.device)
+        else:
+            self.head_nums = self.specific_head
+        self.hooks_injected = False
+        self.training_hooks_injected = False
+        self.original_forward_func = model.forward
+
+        # Activate Unlimiformer when calling model.eval(), deactivate for model.train()
+        self.original_model_eval_func = model.eval
+        model.eval = self.pre_eval_hook
+        self.original_model_train_func = model.train
+        model.train = self.pre_train_hook
+
+    def pre_eval_hook(self):
+        self.remove_training_hooks(self.model)
+        self.inject_hooks(self.model)
+        self.original_model_eval_func()
+
+    def pre_train_hook(self, mode=True):
+        # mode=True means model.train() is called
+        # mode=False means model.eval() is called
+        torch.cuda.empty_cache()
+        if mode is True:
+            self.break_out(self.model)
+            if self.unlimiformer_training:
+                self.inject_training_hooks(self.model)
+        self.original_model_train_func(mode)
+        
+    def inject_hooks(self, model):
+        if self.hooks_injected:
+            return
+        # Inject our activation_capturer to capture the activations at every forward pass
+        attention_layers_to_capture = self.activation_to_capture(self.layer_begin, self.layer_end)
+        self.activation_capturer = []
+        for layer in attention_layers_to_capture:
+            if type(layer) is list:
+                layer_capturers = []
+                for k_or_v in layer:
+                    capturer = ActivationCapturer(k_or_v, capture_input=False)
+                    layer_capturers.append(capturer)
+                    self.register_hook(k_or_v, capturer)
+                self.activation_capturer.append(layer_capturers)
+            else:
+                capturer = ActivationCapturer(layer, capture_input=False)
+                self.register_hook(layer, capturer)
+                self.activation_capturer.append(capturer)
+
+        # Inject our main function after the main attention function
+        attention_layers_to_run = self.attention_op_to_run(self.layer_begin, self.layer_end)
+        for layer in attention_layers_to_run:
+            self.register_hook(layer, self.attention_forward_hook)
+
+        decoder_layers_to_run = self.attention_layer_to_run(self.layer_begin, self.layer_end)
+        self.original_decoder_layer_cross_attn_forward_funcs = []
+        for i, decoder_layer in enumerate(decoder_layers_to_run):
+            decoder_layer_cross_attention = self.cross_attention(decoder_layer)
+            self.original_decoder_layer_cross_attn_forward_funcs.append(decoder_layer_cross_attention.forward)
+            decoder_layer_cross_attention.forward = self.create_cross_attn_pre_forward_hook(decoder_layer_cross_attention.forward, i)
+
+        # Inject our hook function in the beginning of generation.
+        # When the "model.generate()" will be called, it will first call our "reset_generation()" function, 
+        # and only then call "model.generate()"
+        self.original_generate_func = model.generate
+        model.generate = self.pre_generate_hook
+
+        model.forward = self.pre_forward_hook
+        
+        self.original_reorder_cache_func = model._reorder_cache
+        model._reorder_cache = self.reorder_cache_hook
+        self.hooks_injected = True
+
+    def inject_training_hooks(self, model):
+        if self.training_hooks_injected:
+            return
+        # self.original_forward_func = model.forward
+        model.forward = self.pre_forward_hook
+
+        decoder_layers_to_run = self.attention_layer_to_run(self.layer_begin, self.layer_end)
+        
+        self.original_decoder_layer_self_attn_forward_funcs = []
+        for decoder_layer in decoder_layers_to_run:
+            attention = self.self_attention(decoder_layer)
+            self.original_decoder_layer_self_attn_forward_funcs.append(attention.forward)
+            attention.forward = self.create_self_attn_pre_forward_hook(attention.forward)
+
+        self.original_decoder_layer_cross_attn_forward_funcs = []
+        for i, decoder_layer in enumerate(decoder_layers_to_run):
+            decoder_layer_cross_attention = self.cross_attention(decoder_layer)
+            self.original_decoder_layer_cross_attn_forward_funcs.append(decoder_layer_cross_attention.forward)
+            decoder_layer_cross_attention.forward = self.create_cross_attn_pre_forward_hook(decoder_layer_cross_attention.forward, i)
+
+        self.original_decoder_layer_forward_funcs = []
+        for decoder_layer in decoder_layers_to_run:
+            self.original_decoder_layer_forward_funcs.append(decoder_layer.forward)
+            decoder_layer.forward = self.create_decoder_layer_func(decoder_layer.forward, decoder_layer)
+
+        self.inject_hooks_for_unaffected_layers(model, decoder_layers_to_run)
+
+        attention_layers_to_run = self.attention_op_to_run(self.layer_begin, self.layer_end)
+        for layer in attention_layers_to_run:
+            self.register_hook(layer, self.train_attention_forward_hook)
+
+        self.training_hooks_injected = True
+
+    def inject_hooks_for_unaffected_layers(self, model, decoder_layers_to_run):
+        self.original_non_injected_decoder_layer_forward_funcs = []
+        non_injected_decoder_layers = [l for l in self.attention_layer_to_run(0, None) 
+            if l not in decoder_layers_to_run]
+        for decoder_layer in non_injected_decoder_layers:
+            self.original_non_injected_decoder_layer_forward_funcs.append(decoder_layer.forward)
+            decoder_layer.forward = self.create_noninjected_decoder_layer_func(decoder_layer.forward, decoder_layer)
+
+    def create_self_attn_pre_forward_hook(self, original_self_attn_forward_func):
+        def self_attention_pre_forward_hook(*args, **kwargs):
+            kwargs['past_key_value'] = None
+            return original_self_attn_forward_func(*args, **kwargs)
+        
+        return self_attention_pre_forward_hook
+
+    def create_decoder_layer_func(self, decoder_layer_original_forward_func, decoder_layer):
+        def checkpointed_decoder_layer(
+                hidden_states: torch.Tensor,
+                attention_mask=None,
+                encoder_hidden_states=None,
+                encoder_attention_mask=None,
+                layer_head_mask=None,
+                cross_attn_layer_head_mask=None,
+                past_key_value=None,
+                output_attentions=False,
+                position_bias=None,
+                encoder_decoder_position_bias=None,
+                use_cache=True):
+
+           
+            def forward_with_all_keys(hidden_states, attention_mask, 
+                    encoder_hidden_states, encoder_attention_mask, layer_head_mask, 
+                    cross_attn_layer_head_mask, past_key_value, 
+                    output_attentions, use_cache, long_inputs, long_inputs_mask,
+                    position_bias, encoder_decoder_position_bias):
+                
+                key, value = self.create_key_value(long_inputs, decoder_layer)
+                decoder_layer_args = self.create_decoder_layer_args(
+                    hidden_states=hidden_states,
+                    attention_mask=attention_mask,
+                    encoder_hidden_states=encoder_hidden_states,
+                    encoder_attention_mask=encoder_attention_mask,
+                    layer_head_mask=layer_head_mask,
+                    cross_attn_layer_head_mask=cross_attn_layer_head_mask,
+                    past_key_value=past_key_value,
+                    output_attentions=output_attentions,
+                    position_bias=position_bias,
+                    encoder_decoder_position_bias=encoder_decoder_position_bias,
+                    use_cache=use_cache,
+                    key=key,value=value)
+                return decoder_layer_original_forward_func(**decoder_layer_args)
+
+            return torch.utils.checkpoint.checkpoint(
+                forward_with_all_keys, hidden_states, attention_mask, 
+                encoder_hidden_states, encoder_attention_mask, layer_head_mask, 
+                cross_attn_layer_head_mask, None, 
+                output_attentions, use_cache, self.long_inputs_encoded, self.long_inputs_mask,
+                position_bias, encoder_decoder_position_bias)
+
+        return checkpointed_decoder_layer
+
+    def create_noninjected_decoder_layer_func(self, decoder_layer_original_forward_func, decoder_layer):
+        def checkpointed_decoder_layer(
+                hidden_states: torch.Tensor,
+                attention_mask=None,
+                encoder_hidden_states=None,
+                encoder_attention_mask=None,
+                layer_head_mask=None,
+                cross_attn_layer_head_mask=None,
+                past_key_value=None,
+                output_attentions=False,
+                position_bias=None,
+                encoder_decoder_position_bias=None,
+                use_cache=True):
+
+           
+            def forward_with_all_keys(hidden_states, attention_mask, 
+                    encoder_hidden_states, encoder_attention_mask, layer_head_mask, 
+                    cross_attn_layer_head_mask, past_key_value, 
+                    output_attentions, use_cache, long_inputs, long_inputs_mask,
+                    position_bias, encoder_decoder_position_bias):
+                
+                decoder_layer_args = self.create_decoder_layer_args(
+                    hidden_states=hidden_states,
+                    attention_mask=attention_mask,
+                    encoder_hidden_states=encoder_hidden_states,
+                    encoder_attention_mask=encoder_attention_mask,
+                    layer_head_mask=layer_head_mask,
+                    cross_attn_layer_head_mask=cross_attn_layer_head_mask,
+                    past_key_value=past_key_value,
+                    output_attentions=output_attentions,
+                    position_bias=position_bias,
+                    encoder_decoder_position_bias=encoder_decoder_position_bias,
+                    use_cache=use_cache, key=None, value=None)
+                return decoder_layer_original_forward_func(**decoder_layer_args)
+
+            return torch.utils.checkpoint.checkpoint(
+                forward_with_all_keys, hidden_states, attention_mask, 
+                encoder_hidden_states, encoder_attention_mask, layer_head_mask, 
+                cross_attn_layer_head_mask, None, 
+                output_attentions, use_cache, self.long_inputs_encoded, self.long_inputs_mask,
+                position_bias, encoder_decoder_position_bias)
+
+        return checkpointed_decoder_layer
+
+    def register_hook(self, layer, func, pre=False):
+        handle = layer.register_forward_pre_hook(func) if pre else layer.register_forward_hook(func)
+        self.hook_handles.append(handle)
+
+    def break_out(self, model):
+        self.prompt_keys = []
+        self.prompt_values = []
+        self.prompt_attention_mask = []
+        self.generated_input_ids = []
+        torch.cuda.empty_cache()
+        if not self.hooks_injected:
+            return
+
+        for h in self.hook_handles:
+            h.remove()
+        model.generate = self.original_generate_func
+        model.forward = self.original_forward_func
+        model._reorder_cache = self.original_reorder_cache_func
+
+        decoder_layers_to_run = self.attention_layer_to_run(self.layer_begin, self.layer_end)
+        for decoder_layer, original_func in zip(decoder_layers_to_run, self.original_decoder_layer_cross_attn_forward_funcs):
+            self.cross_attention(decoder_layer).forward = original_func
+        self.hooks_injected = False
+
+    def remove_training_hooks(self, model):
+        self.long_inputs_encoded, self.long_inputs_mask = None, None
+        if not self.training_hooks_injected:
+            return
+        for h in self.hook_handles:
+            h.remove()
+        model.forward = self.original_forward_func
+
+        decoder_layers_to_run = self.attention_layer_to_run(self.layer_begin, self.layer_end)
+        for decoder_layer, original_func in zip(decoder_layers_to_run, self.original_decoder_layer_self_attn_forward_funcs):
+            self.self_attention(decoder_layer).forward = original_func
+        for decoder_layer, original_func in zip(decoder_layers_to_run, self.original_decoder_layer_cross_attn_forward_funcs):
+            self.cross_attention(decoder_layer).forward = original_func
+        for decoder_layer, original_func in zip(decoder_layers_to_run, self.original_decoder_layer_forward_funcs):
+            decoder_layer.forward = original_func
+
+        non_injected_decoder_layers = [l for l in self.attention_layer_to_run(0, None) 
+            if l not in decoder_layers_to_run]
+        for decoder_layer, original_func in zip(non_injected_decoder_layers, self.original_non_injected_decoder_layer_forward_funcs):
+            decoder_layer.forward = original_func
+
+        self.training_hooks_injected = False
+
+    def reset_memory(self, input_ids, attention_mask):
+        if self.use_datastore:
+            if self.is_encoder_decoder:
+                self.datastore = [DatastoreBatch(dim=self.model.config.hidden_size, batch_size=input_ids.shape[0], flat_index=self.flat_index, 
+                    gpu_index=self.gpu_index, index_device=self.index_devices[0])]
+                self.hidden_states = [[]]
+            else:
+                self.datastore = [DatastoreBatch(dim=self.model.config.hidden_size, batch_size=input_ids.shape[0], flat_index=self.flat_index, 
+                    gpu_index=self.gpu_index, index_device=self.index_devices[i % len(self.index_devices)]) 
+                    for i in range(self.model.config.num_hidden_layers)[self.layer_begin:self.layer_end]]
+                self.hidden_states = [[] for _ in range(self.model.config.num_hidden_layers)[self.layer_begin:self.layer_end]]
+            torch.cuda.empty_cache()
+        self.prompt_input_ids = input_ids
+        self.input_ids_size = input_ids.shape[-1]
+        self.prompt_keys, self.prompt_values = None, None
+        self.prev_tokens = [None for _ in range(len(self.original_decoder_layer_cross_attn_forward_funcs))]
+        self.last_beam_idx = None
+        self.cur_layer_key_value_placeholder = None
+        self.is_input_encoding_pass = True
+        if self.is_encoder_decoder:
+            dummy_labels = torch.zeros((input_ids.shape[0], 1), dtype=torch.long, device=input_ids.device)
+        else:
+            dummy_labels = None
+        if self.save_heatmap:
+            if self.heatmap is not None:
+                print(f'Generated: {self.tokenizer.decode(self.generated_input_ids[0])}')
+                self.plot_heatmap(self.heatmap[0].detach().cpu().numpy())
+            self.heatmap = torch.tensor([], dtype=torch.float, device=input_ids.device)
+        self.generated_input_ids = torch.tensor([], dtype=torch.long, device=input_ids.device)
+
+        self.prompt_keys = [[] for _ in range(self.model.config.num_hidden_layers)[self.layer_begin:self.layer_end]]
+        self.prompt_values = [[] for _ in range(self.model.config.num_hidden_layers)[self.layer_begin:self.layer_end]]
+        self.prompt_attention_mask = []
+        window_indices = self.window_indices(input_ids.shape[-1])
+
+        for context_start_ind, context_end_ind, update_start_ind, update_end_ind in window_indices:
+            logger.info(f'Encoding {context_start_ind} to {context_end_ind} out of {input_ids.shape[-1]}')
+            chunk = input_ids[:, context_start_ind:context_end_ind].to(self.device)
+            chunk_attention_mask = attention_mask[:, context_start_ind:context_end_ind].to(self.device)
+            with torch.inference_mode():
+                _ = self.model(chunk, attention_mask=chunk_attention_mask, labels=dummy_labels) # , return_dict=True, output_hidden_states=True)
+            if self.use_datastore:
+                # TODO: verify with BART as well
+                # hidden_states_to_index = [hidden_states.encoder_last_hidden_state] # list of length 1 of (batch, chunked_source_len, dim)
+                hidden_states_to_index = [
+                    layer_capturer.captured for layer_capturer in self.activation_capturer
+                ] 
+                # hidden_states_to_index = list(hidden_states.hidden_states)[:-1][self.layer_begin:self.layer_end]
+                to_add = [state[:, update_start_ind:update_end_ind].detach() for state in hidden_states_to_index]
+                to_apply_mask = chunk_attention_mask[:, update_start_ind:update_end_ind]
+                # to_apply_mask = to_apply_mask.log().to(to_add[0].dtype)
+                to_apply_mask = to_apply_mask.to(to_add[0].dtype)
+                if not self.reconstruct_embeddings:
+                    to_add_embeddings = to_add
+                    if not self.gpu_datastore:
+                        to_add_embeddings = [states.cpu() for states in to_add_embeddings]
+                        to_apply_mask = to_apply_mask.cpu()
+                    for i, layer_states in enumerate(to_add_embeddings):
+                        layer_states = layer_states * to_apply_mask.unsqueeze(-1)
+                        self.hidden_states[i].append(layer_states.to(self.datastore_device))
+                # list of len layers, inside it there is a list of len batch, each item is (masked_time, dim)
+                # for i, to_add_layer in enumerate(to_add):
+                #     keys = [key[mask.bool()] for key, mask in zip(to_add_layer, to_apply_mask)]
+                #     self.datastore[i].add_keys(keys)
+            if (not self.use_datastore) or self.test_datastore:
+                layers_kv = [
+                    self.process_key_value(layer_capturer) # (batch, head, time, dim)
+                    for layer_capturer in self.activation_capturer
+                ] # list of pairs of (batch, head, time, dim)
+
+                # list of (batch, head, chunked_time, dim)
+                key = [layer[0][:, :, update_start_ind:update_end_ind] for layer in layers_kv]
+                value = [layer[1][:, :, update_start_ind:update_end_ind] for layer in layers_kv]
+                chunk_attention_mask = chunk_attention_mask[:, update_start_ind:update_end_ind] # (batch, chunked_time)
+
+                # key = torch.stack(key, dim=0) # (num_layers, batch, head, time, dim)
+                # value = torch.stack(value, dim=0) # (num_layers, batch, head, time, dim)
+
+                for i, (layer_key, layer_value) in enumerate(zip(key, value)):
+                    self.prompt_keys[i].append(layer_key) # (num_layers, batch, head, chunked_source_len, dim)
+                    self.prompt_values[i].append(layer_value) # (num_layers, batch, head, chunked_source_len, dim)
+                self.prompt_attention_mask.append(chunk_attention_mask) # (batch, chunked_source_len)
+        
+        if self.use_datastore:
+            # keys are all in datastore already!
+            if not self.reconstruct_embeddings:
+                # self.hidden_states = [torch.cat(layer_hidden_states, axis=1) for layer_hidden_states in self.hidden_states]
+                concat_hidden_states = []
+                for i in range(len(self.hidden_states)):
+                    concat_hidden_states.append(torch.cat(self.hidden_states[i], axis=1))
+                    self.hidden_states[i] = None
+                self.hidden_states = concat_hidden_states
+            for datastore, layer_hidden_states in zip(self.datastore, self.hidden_states):
+                datastore.train_index(layer_hidden_states)
+        if (not self.use_datastore) or self.test_datastore:
+            for i, (layer_keys, layer_values) in enumerate(zip(self.prompt_keys, self.prompt_values)):
+                self.prompt_keys[i] = torch.cat(layer_keys, dim=-2)
+                self.prompt_values[i] = torch.cat(layer_values, dim=-2)
+            # self.prompt_keys = torch.cat(self.prompt_keys, dim=-2) # (num_layers, batch, head, source_len, dim)
+            # self.prompt_values = torch.cat(self.prompt_values, dim=-2) # (num_layers, batch, head, source_len, dim)
+            self.prompt_attention_mask = torch.cat(self.prompt_attention_mask, dim=-1) # (batch, source_len)
+
+        self.is_input_encoding_pass = False
+        if self.verbose:
+            print(f'Input: '
+                f'{self.tokenizer.decode(input_ids[0][:self.actual_model_window_size], skip_special_tokens=True)} ||| '
+                f'{self.tokenizer.decode(input_ids[0][self.actual_model_window_size:], skip_special_tokens=True)}')
+            print()
+
+    def chunked_encode_input(self, input_ids, attention_mask):
+        long_inputs_encoded = []
+        long_inputs_mask = []
+        window_indices = self.window_indices(input_ids.shape[-1])
+
+        self.is_input_encoding_pass = True
+        for context_start_ind, context_end_ind, update_start_ind, update_end_ind in window_indices:
+            chunk = input_ids[:, context_start_ind:context_end_ind]
+            chunk_attention_mask = attention_mask[:, context_start_ind:context_end_ind]
+            output = self.model.base_model.encoder(chunk, attention_mask=chunk_attention_mask, return_dict=True, output_hidden_states=True)
+            encoder_last_hidden_state = output.last_hidden_state # (batch, time, dim)
+            
+            # list of (batch, head, chunked_time, dim)
+            encoder_last_hidden_state = encoder_last_hidden_state[:, update_start_ind:update_end_ind] # (batch, chunked_time, dim)
+            chunk_attention_mask = chunk_attention_mask[:, update_start_ind:update_end_ind] # (batch, chunked_time)
+
+            long_inputs_encoded.append(encoder_last_hidden_state) # (batch, chunked_source_len, dim)
+            long_inputs_mask.append(chunk_attention_mask) # (batch, chunked_source_len)
+        
+        long_inputs_encoded = torch.cat(long_inputs_encoded, dim=1) # (batch, source_len, dim)
+        long_inputs_mask = torch.cat(long_inputs_mask, dim=1) # (batch, source_len)
+
+        self.is_input_encoding_pass = False
+        if self.verbose:
+            print(f'Input: '
+                f'{self.tokenizer.decode(input_ids[0][:self.actual_model_window_size], skip_special_tokens=True)} ||| '
+                f'{self.tokenizer.decode(input_ids[0][self.actual_model_window_size:], skip_special_tokens=True)}')
+            print()
+        return long_inputs_encoded, long_inputs_mask
+
+    def window_indices(self, total_seq_len):
+        # Copied from SLED (Ivgy et al., 2022)
+        # https://github.com/Mivg/SLED/blob/main/sled/modeling_sled.py#L467
+        if total_seq_len <= self.model_encoder_max_len:
+            return [(0, total_seq_len, 0, total_seq_len)]
+        else:
+            results = []
+            # if self.chunk_overlap == 0:
+            #     stride = self.model_encoder_max_len
+            stride = self.model_encoder_max_len - 2 * self.window_margin
+            context_start = update_start_ind = 0
+            context_end = self.model_encoder_max_len
+            if self.is_encoder_decoder:
+                update_end_ind = context_end - self.window_margin
+            else:
+                update_end_ind = context_end
+            # first window always should update from the beginning
+            results.append((context_start, context_end, update_start_ind, update_end_ind))  
+
+            while context_end < total_seq_len:
+                context_end = min(total_seq_len, context_end + stride)
+                context_start = (
+                    context_start + stride if context_end < total_seq_len else total_seq_len - self.model_encoder_max_len
+                )
+                update_start_ind = max(update_start_ind + stride, update_end_ind)
+                # last window always should update until the end
+                update_end_ind = (
+                    min(total_seq_len, update_end_ind + stride) if context_end < total_seq_len else total_seq_len
+                )
+
+                cs, ce, us, ue = context_start, context_end, update_start_ind - context_start, \
+                                 update_end_ind - context_start
+
+                results.append((cs, ce, us, ue))
+            return results
+
+    def pre_generate_hook(self, input_ids, **kwargs):
+        if 'attention_mask' not in kwargs:
+            kwargs['attention_mask'] = torch.ones_like(input_ids)
+        self.reset_memory(input_ids, kwargs['attention_mask'])
+        new_kwargs = kwargs
+        if 'attention_mask' in kwargs:
+            new_kwargs = {k: v for k, v in kwargs.items() if k != 'attention_mask'}
+            new_kwargs['attention_mask'] = kwargs['attention_mask'][:, :self.actual_model_window_size].to(self.device)
+        new_kwargs['use_cache'] = True
+        if self.is_encoder_decoder:
+            input_ids_prefix = input_ids[:, :self.actual_model_window_size]
+        else:
+            input_ids_prefix = input_ids[:, -self.actual_model_window_size:]
+        input_ids_prefix = input_ids_prefix.to(self.device)
+        return self.original_generate_func(input_ids_prefix, **new_kwargs)
+
+    def pre_forward_hook(self, input_ids=None, attention_mask=None, labels=None, **kwargs):
+        self.set_gradient_checkpointing(False)
+        if not self.is_input_encoding_pass:
+            if self.model.training:
+                # self.reset_memory(input_ids, attention_mask)
+                self.long_inputs_encoded, self.long_inputs_mask = self.chunked_encode_input(input_ids=input_ids, attention_mask=attention_mask)
+                input_ids = input_ids[:, :self.actual_model_window_size]
+                attention_mask = attention_mask[:, :self.actual_model_window_size] if attention_mask is not None else None
+                # input_ids = input_ids[:, :self.model_encoder_max_len]
+                # labels = labels[:, :self.model_encoder_max_len] if labels is not None else None
+            else:
+                if kwargs.get('past_key_values') is None:
+                    self.is_first_test_decoding_step = True
+
+                if input_ids is not None:
+                    # self.input_ids_size += input_ids.shape[-1]
+                    self.input_ids_size += 1
+                if kwargs.get('decoder_input_ids') is not None:
+                    self.generated_input_ids = torch.cat([self.generated_input_ids, kwargs['decoder_input_ids']], axis=-1)
+            
+        result = self.original_forward_func(input_ids=input_ids, labels=labels, attention_mask=attention_mask, **kwargs)
+        self.is_first_test_decoding_step = False
+        return result
+
+    def create_cross_attn_pre_forward_hook(self, original_cross_attn_forward_func, cur_layer_num):
+        def attention_pre_forward_hook(hidden_states, attention_mask=None, *args, **kwargs):
+            self.cur_decoder_layer_index = cur_layer_num
+            if kwargs.get('past_key_value') is not None:
+                # it's a tuple, and we convert it to a list to be able to perform assignment 
+                # and modify its items from our attention_forward_hook
+                self.cur_layer_key_value_placeholder = \
+                    kwargs['past_key_value'] = list(kwargs['past_key_value']) # (batch, head, time, attn_dim)
+
+            batch_size, tgt_len, dim = hidden_states.shape
+            if self.model.training:
+                # from: (batch, tgt_len, dim) to: (batch * tgt_len, 1, dim)
+                hidden_states = hidden_states.reshape(-1, 1, hidden_states.shape[-1])
+                # from: (batch, 1, tgt_len, dim) to: (batch * tgt_len, 1, 1, dim)
+                attention_mask = attention_mask.reshape(-1, 1, 1, attention_mask.shape[-1])
+                
+                attn_output, attn_weights_reshaped, past_key_value = original_cross_attn_forward_func(hidden_states=hidden_states, attention_mask=attention_mask, *args, **kwargs)
+                attn_output = attn_output.reshape(batch_size, tgt_len, dim)
+                result = (attn_output, attn_weights_reshaped, past_key_value)
+            else:
+                result = original_cross_attn_forward_func(hidden_states=hidden_states, attention_mask=attention_mask, *args, **kwargs)
+                # Uri: this part adds the generated tokens to the prompt. 
+                # However it was commented out because currently we always keep the generated tokens in the attention window
+                # if not self.is_encoder_decoder and not self.is_input_encoding_pass and \
+                #         past_key_value[0].shape[2] > self.prompt_keys[self.cur_decoder_layer_index].shape[2]:
+                #     self.prompt_keys[self.cur_decoder_layer_index] = torch.cat([self.prompt_keys[self.cur_decoder_layer_index], past_key_value[0][:,:,-1:]], dim=-2)
+                #     self.prompt_values[self.cur_decoder_layer_index] = torch.cat([self.prompt_values[self.cur_decoder_layer_index], past_key_value[1][:,:,-1:]], dim=-2)
+                #     if self.cur_decoder_layer_index == self.model.config.num_hidden_layers - 1:
+                #         self.prompt_attention_mask = torch.cat([
+                #             self.prompt_attention_mask, 
+                #             torch.ones([self.prompt_attention_mask.shape[0], 1], dtype=self.prompt_attention_mask.dtype).to(self.device)], dim=-1)
+        
+            return result
+        return attention_pre_forward_hook
+
+
+    def attention_forward_hook(self, module, input, output):
+        # output: (batch, time, 3 * heads * attention_dim)
+        if self.is_input_encoding_pass or self.is_first_test_decoding_step:
+            return
+        with torch.no_grad():
+            prompt_size = self.prompt_input_ids.shape[1]
+            generated_size = self.input_ids_size - prompt_size
+            window_size = self.cur_layer_key_value_placeholder[0].shape[-2]
+            # topk = min(self.actual_model_window_size, attn_weights.shape[-1])
+            topk = min(prompt_size, window_size)
+            if not self.is_encoder_decoder:
+                topk = min(topk, window_size - generated_size + 1)
+            if self.gpu_index:
+                topk = min(topk, 2048)
+
+            query = self.process_query(output)[:,-1] # (batch * beam, head, dim)
+            query = query[:, self.head_nums] # (batch * beam, head, dim)
+
+            if self.use_datastore:
+                # query: (batch, beam, head, dim)
+                # need to multiply by key vector
+                # query.view(query.shape[0], query.shape[1] * query.shape[2])
+                # k_proj in attention? 
+                datastore_index = 0 if self.is_encoder_decoder else self.cur_decoder_layer_index
+                attention_layer_list = self.get_kv_projections(self.layer_begin, self.layer_end)
+                k_proj_layer = [layers[0] for layers in attention_layer_list][self.cur_decoder_layer_index]
+                v_proj_layer = [layers[1] for layers in attention_layer_list][self.cur_decoder_layer_index]
+                
+                # modify query by k_projs 
+                k_proj = k_proj_layer.weight
+                datastore_query = self.preprocess_query(query, k_proj) # (batch * beam, num_heads, embed_dim)
+                batch_size = self.datastore[datastore_index].batch_size
+                datastore_query = datastore_query.view((batch_size, -1, datastore_query.shape[2])) # (batch, beam * num_heads, embed_dim)
+                # then search
+                if self.reconstruct_embeddings:
+                    # embeddings: (batch, beam * head, actual_model_window_size, dim)
+                    _, top_search_key_indices, embeddings = self.datastore[datastore_index].search_and_reconstruct(datastore_query, k=topk) 
+                else:
+                    _, top_search_key_indices = self.datastore[datastore_index].search(datastore_query, k=topk)
+                    # self.embeddings: (batch,              src_len, dim)
+                    # indices:         (batch, beam * head, actual_model_window_size)
+                    # embeddings: (batch, beam * head, actual_model_window_size, dim)
+                    embeddings = torch.take_along_dim(input=self.hidden_states[datastore_index].unsqueeze(1), 
+                        indices=top_search_key_indices.unsqueeze(-1).to(self.hidden_states[datastore_index].device), dim=-2)
+                    embeddings = embeddings.to(self.device)
+                # (batch, beam, head, actual_model_window_size)
+                # top_search_key_scores = top_search_key_scores.reshape(batch_size, -1, *top_search_key_scores.shape[1:])
+                top_search_key_indices = top_search_key_indices.reshape(batch_size, -1, *top_search_key_indices.shape[1:])
+                # embeddings: (batch, beam, head, actual_model_window_size, dim)
+                embeddings = embeddings.reshape(batch_size, -1, self.num_heads, *embeddings.shape[2:])
+                                    
+            # raw_values are actually token indices; need to look them up
+            if (not self.use_datastore) or self.test_datastore:
+                this_layer_prompt_keys = self.prompt_keys[self.cur_decoder_layer_index]
+                this_layer_prompt_values = self.prompt_values[self.cur_decoder_layer_index]
+                # query: (batch * beam, head, dim)
+                batch_size = self.prompt_input_ids.shape[0]
+                beam_size = query.shape[0] // batch_size
+                # query: (batch, beam, head, dim)
+                query = query.reshape(batch_size, beam_size, *query.shape[1:])
+                # this_layer_prompt_keys: (batch, head, source_len, dim)
+                # this_layer_prompt_keys.unsqueeze(1):  (batch, 1, head, source_len, dim)
+                # query.unsqueeze(-1):             (batch, beam, head, dim, 1)
+                # attn_weights:  (batch, beam, head, source_len)
+                attn_weights = torch.matmul(this_layer_prompt_keys.unsqueeze(1)[:, :, self.head_nums], query.unsqueeze(-1)).squeeze(-1) 
+                # attn_weights = torch.matmul(query.unsqueeze(-2), this_layer_prompt_keys.unsqueeze(1)[:, :, self.head_nums]).squeeze(-2) 
+                prompt_attention_mask_to_add = (1 - self.prompt_attention_mask) * -1e9 # (batch, source_len)
+                prompt_attention_mask_to_add = prompt_attention_mask_to_add.unsqueeze(1).unsqueeze(1)
+                attn_weights += prompt_attention_mask_to_add # (batch, beam, head, source_len)
+                if self.exclude_attention and attn_weights.shape[-1] > self.actual_model_window_size:
+                    attn_weights[..., :self.actual_model_window_size] -= 1e9
+
+                # target_keys, target_values, topk = self.get_target_slices(output)
+                top_key_scores, top_key_indices = torch.topk(attn_weights, k=topk, dim=-1, sorted=True) # (batch, beam, head, trunc_source)
+                if self.save_heatmap:
+                    # heatrow: (beam, heads, source_len)
+                    heatrow = torch.zeros([top_key_indices.shape[1], top_key_indices.shape[2], this_layer_prompt_keys.shape[-2]], dtype=torch.float)
+                    heatrow = heatrow.scatter(index=top_key_indices[0], src=torch.ones_like(top_key_scores[0]), dim=-1)
+                    # heatrow = torch.nn.functional.softmax(heatrow, dim=-1)
+                    # self.heatmap: (beam, heads, targets, source_len)
+                    self.heatmap = torch.cat([self.heatmap, heatrow.unsqueeze(-2)], axis=-2)
+
+            if self.test_datastore:
+                assert top_key_indices.shape == top_search_key_indices.shape
+                assert torch.mean((top_key_indices == top_search_key_indices).float()) > 0.99
+
+            if self.verbose:
+                if self.is_encoder_decoder:
+                    for i, beam in enumerate(self.generated_input_ids):
+                        print(f'({i}) Generated: {self.tokenizer.decode(beam)}')
+                # else:
+                #     print(f'Generated: {self.tokenizer.decode(self.input_ids)}')
+                print()
+        
+        if self.use_datastore:
+            # k_proj_layer.weight, v_proj_layer.weight: (embed_dim, embed_dim)
+            # embeddings: (batch, beam, head, encoder_len, embed_dim)
+            retrieved_keys, retrieved_values = self.post_process_retrieved(embeddings, k_proj_layer, v_proj_layer, top_search_key_indices)
+        else:
+            # this_layer_prompt_keys:   (batch,       head, source_len, dim)
+            # top_key_indices:          (batch, beam, head, trunc_source)
+            retrieved_keys = torch.take_along_dim(this_layer_prompt_keys.unsqueeze(1), indices=top_key_indices.unsqueeze(-1), 
+                dim=-2) # (batch, head, trunc_source, attn_dim)
+            retrieved_values = torch.take_along_dim(this_layer_prompt_values.unsqueeze(1), indices=top_key_indices.unsqueeze(-1), 
+                dim=-2) # (batch, head, trunc_source, attn_dim)
+
+        if self.test_datastore:
+            correct_keys = torch.take_along_dim(this_layer_prompt_keys.unsqueeze(1), indices=top_key_indices.unsqueeze(-1), 
+                dim=-2) # (batch, head, trunc_source, attn_dim)
+            correct_values = torch.take_along_dim(this_layer_prompt_values.unsqueeze(1), indices=top_key_indices.unsqueeze(-1), 
+                dim=-2) # (batch, head, trunc_source, attn_dim)
+            assert correct_keys.shape == retrieved_keys.shape
+            assert correct_values.shape == retrieved_values.shape
+            assert torch.mean(torch.isclose(correct_keys, retrieved_keys, rtol=1e-3, atol=1e-3).float()) > 0.99
+            assert torch.mean(torch.isclose(correct_values, retrieved_values, rtol=1e-3, atol=1e-3).float()) > 0.99
+
+        # retrieved_keys, retrieved_values: (batch * beam, head, encoder_len, attn_dim)
+        retrieved_keys = retrieved_keys.flatten(0, 1)[:,:,:topk]
+        retrieved_values = retrieved_values.flatten(0, 1)[:,:,:topk]
+        self.cur_layer_key_value_placeholder[0] = torch.cat([retrieved_keys, self.cur_layer_key_value_placeholder[0][:,:,topk:]], dim=-2)
+        self.cur_layer_key_value_placeholder[1] = torch.cat([retrieved_values, self.cur_layer_key_value_placeholder[1][:,:,topk:]], dim=-2)
+        return
+
+    def train_attention_forward_hook(self, module, input, output):
+        # output: (batch, time, 3 * heads * attention_dim)
+        if self.is_input_encoding_pass or self.is_first_test_decoding_step:
+            return
+        this_layer_prompt_keys = self.cur_layer_key_value_placeholder[0]
+        this_layer_prompt_values = self.cur_layer_key_value_placeholder[1]
+        with torch.no_grad():
+            query = self.process_query(output) # (batch * beam, tgt_len, head, dim)
+            # query = query[:, :, self.head_nums] # (batch * beam, head, dim)
+            
+            # query: (batch * beam, tgt_len, head, dim)
+            batch_size = this_layer_prompt_keys.shape[0]
+            tgt_len = query.shape[0] // batch_size
+            # query: (batch, tgt, head, dim)
+            query = query.reshape(batch_size, tgt_len, *query.shape[2:])
+            # this_layer_prompt_keys: (batch, head, source_len, dim)
+            # this_layer_prompt_keys.unsqueeze(1):  (batch, 1, head, source_len, dim)
+            # attn_weights:  (batch, tgt_len, head, 1, source_len)
+            # attn_weights = torch.matmul(query.unsqueeze(-2), this_layer_prompt_keys.unsqueeze(1).permute(0,1,2,4,3))
+            attn_weights = torch.matmul(this_layer_prompt_keys.unsqueeze(1), query.unsqueeze(-1)) \
+                .reshape(batch_size, tgt_len, query.shape[-2], 1, this_layer_prompt_keys.shape[-2])
+            # attn_weights = torch.matmul(query.unsqueeze(-2), this_layer_prompt_keys.unsqueeze(1)[:, :, self.head_nums]).squeeze(-2) 
+            prompt_attention_mask_to_add = (1 - self.long_inputs_mask) * -1e9 # (batch, source_len)
+            prompt_attention_mask_to_add = prompt_attention_mask_to_add.unsqueeze(1).unsqueeze(1).unsqueeze(1)
+            attn_weights += prompt_attention_mask_to_add # (batch, beam, head, source_len)
+
+            # target_keys, target_values, topk = self.get_target_slices(output)
+            topk = min(self.actual_model_window_size, attn_weights.shape[-1])
+            top_key_scores, top_key_indices = torch.topk(attn_weights, k=min(topk, attn_weights.shape[-1]), dim=-1, sorted=True) # (batch, beam, head, tgt, trunc_source)
+
+                   
+        # this_layer_prompt_keys:   (batch,          head,    source_len, dim)
+        # top_key_indices:          (batch, tgt_len, head, 1, trunc_source)
+        new_keys = torch.take_along_dim(this_layer_prompt_keys.unsqueeze(2).unsqueeze(1), indices=top_key_indices.unsqueeze(-1), 
+            dim=-2) # (batch, tgt_len, head, 1, trunc_source, attn_dim)
+        new_values = torch.take_along_dim(this_layer_prompt_values.unsqueeze(2).unsqueeze(1), indices=top_key_indices.unsqueeze(-1), 
+            dim=-2) # (batch, tgt_len, head, 1, trunc_source, attn_dim)
+        
+        # (batch * beam, head, tgt_len, trunc_source, attn_dim)
+        self.cur_layer_key_value_placeholder[0] = new_keys.flatten(0, 1).squeeze(2)
+        self.cur_layer_key_value_placeholder[1] = new_values.flatten(0, 1).squeeze(2)
+        return
+
+    
+    def preprocess_query(self, query, k_proj_weight):
+        k_proj = k_proj_weight.view(1, self.num_heads, query.shape[-1], k_proj_weight.shape[0]) # (1, num_heads, attn_dim, embed_dim)
+        datastore_query = query.unsqueeze(-2) # (batch * beam, num_heads, 1, attn_dim)
+        datastore_query = torch.matmul(datastore_query, k_proj) # (batch * beam, num_heads, 1, embed_dim)
+        datastore_query = datastore_query.squeeze(-2)  # (batch * beam, num_heads, embed_dim)
+        return datastore_query
+
+    def post_process_retrieved(self, embeddings, k_proj_layer, v_proj_layer, top_search_key_indices):
+        embed_dim = embeddings.shape[-1]
+        k_weight = k_proj_layer.weight.view(1, 1, self.num_heads, embed_dim // self.num_heads, embed_dim).transpose(-2,-1) # (1, 1, heads, embed_dim, attn_dim)
+        k_bias = 0
+        if k_proj_layer.bias is not None:
+            k_bias = k_proj_layer.bias.view(1, self.num_heads, embed_dim // self.num_heads).unsqueeze(-2).unsqueeze(0)
+        v_weight = v_proj_layer.weight.view(1, 1, self.num_heads, embed_dim // self.num_heads, embed_dim).transpose(-2,-1)  # (1, heads, embed_dim, attn_dim)
+        v_bias = 0
+        if v_proj_layer.bias is not None:
+            v_bias = v_proj_layer.bias.view(1, self.num_heads, embed_dim // self.num_heads).unsqueeze(-2).unsqueeze(0)
+        # new_keys, new_values: (batch, beam, head, encoder_len, attn_dim)
+        retrieved_keys = torch.matmul(embeddings, k_weight) + k_bias # (beam, head, encoder_len, embed_dim)
+        retrieved_values = torch.matmul(embeddings, v_weight) + v_bias # (beam, head, encoder_len, embed_dim)
+        return retrieved_keys, retrieved_values
+
+    def set_gradient_checkpointing(self, value):
+        self.model.base_model.decoder.gradient_checkpointing = value
+
+    def reorder_cache_hook(self, past, beam_idx):
+        self.last_beam_idx = beam_idx
+        self.generated_input_ids = self.generated_input_ids[beam_idx]
+        for i, layer_prev_tokens in enumerate(self.prev_tokens):
+            if layer_prev_tokens is not None:
+                self.prev_tokens[i] = layer_prev_tokens.flatten(0, 1)[beam_idx].reshape(layer_prev_tokens.shape)
+        if self.save_heatmap and self.heatmap.numel() > 0:
+            self.heatmap = self.heatmap[beam_idx]
+        return self.original_reorder_cache_func(past, beam_idx)
+    
+    @classmethod
+    def convert_model(cls, model, *args, **kwargs):
+        # if type(model.config) in MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING:
+        # elif type(model.config) in MODEL_WITH_LM_HEAD_MAPPING:
+        # else:
+        #     raise ValueError(f'Unsupported model type: {type(model.config)}')
+        # if model.config.is_encoder_decoder:
+        #     model_clone = AutoModelForSeq2SeqLM.from_config(model.config)
+        # else:
+        #     model_clone = AutoModelForCausalLM.from_config(model.config)
+        # model_clone.load_state_dict(model.state_dict()).to(args.device)
+        type_to_class = {
+            BartModel: UnlimiformerBART,
+            BartForConditionalGeneration: UnlimiformerBART,
+            T5Model: UnlimiformerT5,
+            T5ForConditionalGeneration: UnlimiformerT5,
+            LEDModel: UnlimiformerLED,
+            LEDForConditionalGeneration: UnlimiformerLED,
+            # LlamaModel: UnlimiformerLLaMa,
+            # LlamaForCausalLM: UnlimiformerLLaMa,
+        }
+        type_to_class[type(model)](model, *args, **kwargs)
+        return model
+        
+
+    def plot_heatmap(self, data, xticklabels='auto', yticklabels='auto'):
+        # data: (heads, targets, source_len)
+        import seaborn as sb
+        import matplotlib.pyplot as plt
+        # print('gat = np.array([')
+        # for row in data[0]:
+        #     rowstr = ', '.join([f'{x:.2f}' for x in row])
+        #     print(f'    [{rowstr}],')
+        # print(']')
+
+        # sb.set(font_scale=1.5, rc={'text.usetex': True})
+        for i in range(data.shape[0]):
+            fig, axes = plt.subplots(1, 1, figsize=(40, 100))
+            cur_ax = axes
+            axes.set_title(f'Head #{i}, length: {data.shape[2]}, target length: {data.shape[1]}')
+            cur_ax = axes
+            # annot = [[x for x in row] for row in data]
+            ax = sb.heatmap(data[i], annot=False, fmt='.2f',
+                            xticklabels=512, yticklabels=yticklabels, ax=cur_ax)
+            ax.xaxis.tick_top()
+            plt.savefig(f'knns_head{i}.pdf')
+            # plt.savefig('gat_s10_contrast.pdf')
+            plt.show()
+
+
+class UnlimiformerBART(Unlimiformer[BartModel]):
+    def __init__(self, model: BartModel, *args, **kwargs):
+        super().__init__(model, *args, **kwargs)
+
+    def create_key_value(self, encoder_hidden_states, decoder_layer):
+        # (batch, time, hidden_dim)
+        attention = decoder_layer.encoder_attn
+        # key, value: (batch, heads, time, attn_dim)
+        key = attention.k_proj(encoder_hidden_states)
+        key = key.view(key.shape[0], -1, attention.num_heads, attention.head_dim).transpose(1, 2).contiguous()
+        value = attention.v_proj(encoder_hidden_states)
+        value = value.view(value.shape[0], -1, attention.num_heads, attention.head_dim).transpose(1, 2).contiguous()
+        # key, value: (batch, heads, time, attn_dim)
+        return key, value 
+
+    def process_key_value(self, capturers):
+        key_capturer, value_capturer = capturers
+        key, value = key_capturer.captured, value_capturer.captured
+        # (batch, time, heads, attn_dim)
+        attention = self.model.base_model.decoder.layers[-1].encoder_attn
+
+        # query, key, value: (batch, heads, time, attn_dim)
+        # query = query.view(query.shape[0], query.shape[1], attention.num_heads, attention.head_dim).transpose(1, 2).contiguous()
+        key = key.view(key.shape[0], -1, attention.num_heads, attention.head_dim).transpose(1, 2).contiguous()
+        value = value.view(value.shape[0], -1, attention.num_heads, attention.head_dim).transpose(1, 2).contiguous()
+        
+        return key, value
+
+    def process_query(self, output):
+        # (batch, time, heads, attn_dim)
+        attention = self.model.base_model.decoder.layers[-1].encoder_attn
+        # query: (batch, heads, time, attn_dim)
+        # query = output.view(output.shape[0], output.shape[1], attention.num_heads, attention.head_dim).transpose(1, 2).contiguous()
+        query = output.view(output.shape[0], output.shape[1], attention.num_heads, attention.head_dim).contiguous()
+        return query
+
+    def get_kv_projections(self, layer_begin, layer_end): 
+        return [
+            [layer.encoder_attn.k_proj, layer.encoder_attn.v_proj]
+            for layer in self.model.base_model.decoder.layers[layer_begin:layer_end]
+        ]
+
+    def activation_to_capture(self, layer_begin, layer_end): 
+        if self.use_datastore:
+            return [self.model.base_model.encoder.layers[-1]]
+        else:
+            return self.get_kv_projections(layer_begin, layer_end)
+
+    def attention_op_to_run(self, layer_begin, layer_end):
+        return [
+            layer.encoder_attn.q_proj
+                for layer in self.model.base_model.decoder.layers[layer_begin:layer_end]
+        ]
+
+    def attention_layer_to_run(self, layer_begin, layer_end): 
+        return self.model.base_model.decoder.layers[layer_begin:layer_end]
+
+    def self_attention(self, decoder_layer):
+        return decoder_layer.self_attn 
+
+    def cross_attention(self, decoder_layer):
+        return decoder_layer.encoder_attn
+    
+    def window_size(self):
+        return self.model.config.max_position_embeddings
+
+    def create_decoder_layer_args(self, hidden_states, attention_mask, encoder_hidden_states,
+                encoder_attention_mask, layer_head_mask, cross_attn_layer_head_mask,
+                past_key_value, output_attentions, position_bias,
+                encoder_decoder_position_bias, use_cache, key, value):
+        args = {'hidden_states': hidden_states, 
+                'attention_mask': attention_mask, 
+                'encoder_hidden_states': encoder_hidden_states, 
+                'encoder_attention_mask': encoder_attention_mask, 
+                'layer_head_mask': layer_head_mask, 
+                'cross_attn_layer_head_mask': cross_attn_layer_head_mask, 
+                'past_key_value': (None, None, key, value), 
+                'output_attentions': output_attentions, 
+                'use_cache': use_cache,}
+        if key is None and value is None:
+            args['past_key_value'] = None
+        return args
+
+class UnlimiformerT5(Unlimiformer[T5Model]):
+    def __init__(self, model: T5Model, *args, **kwargs):
+        super().__init__(model, *args, **kwargs)
+
+    def create_key_value(self, encoder_hidden_states, decoder_layer):
+        # (batch, time, hidden_dim)
+        attention = decoder_layer.layer[1].EncDecAttention
+        # key, value: (batch, heads, time, attn_dim)
+        key = attention.k(encoder_hidden_states)
+        key = key.view(key.shape[0], -1, attention.n_heads, attention.key_value_proj_dim).transpose(1, 2).contiguous()
+        value = attention.v(encoder_hidden_states)
+        value = value.view(value.shape[0], -1, attention.n_heads, attention.key_value_proj_dim).transpose(1, 2).contiguous()
+        
+        return key, value 
+    
+    def process_key_value(self, capturers):
+        key_capturer, value_capturer = capturers
+        key, value = key_capturer.captured, value_capturer.captured
+        # (batch, time, heads, attn_dim)
+        attention = self.model.base_model.decoder.block[-1].layer[1].EncDecAttention
+
+        # query, key, value: (batch, heads, time, attn_dim)
+        # query = query.view(query.shape[0], query.shape[1], attention.num_heads, attention.head_dim).transpose(1, 2).contiguous()
+        key = key.view(key.shape[0], -1, attention.n_heads, attention.key_value_proj_dim).transpose(1, 2).contiguous()
+        value = value.view(value.shape[0], -1, attention.n_heads, attention.key_value_proj_dim).transpose(1, 2).contiguous()
+        
+        return key, value
+
+    def process_query(self, output):
+        # (batch, time, heads, attn_dim)
+        attention = self.model.base_model.decoder.block[-1].layer[1].EncDecAttention
+        # query: (batch, heads, time, attn_dim)
+        query = output.view(output.shape[0], -1, attention.n_heads, attention.key_value_proj_dim).contiguous()
+        return query
+
+    def get_kv_projections(self, layer_begin, layer_end):
+        return [
+            [layer.layer[1].EncDecAttention.k, layer.layer[1].EncDecAttention.v]
+                for layer in self.model.base_model.decoder.block[layer_begin:layer_end]
+        ]
+    
+    def activation_to_capture(self, layer_begin, layer_end): 
+        if self.use_datastore:
+            return [self.model.base_model.encoder.layers[-1]]
+        else:
+            return self.get_kv_projections(layer_begin, layer_end)
+
+    def attention_op_to_run(self, layer_begin, layer_end):
+        return [
+            layer.layer[1].EncDecAttention.q
+                for layer in self.model.base_model.decoder.block[layer_begin:layer_end]
+        ]
+    
+    def attention_layer_to_run(self, layer_begin, layer_end): 
+        return self.model.base_model.decoder.block[layer_begin:layer_end]
+
+    def self_attention(self, decoder_layer):
+        return decoder_layer.layer[0]
+
+    def cross_attention(self, decoder_layer):
+        return decoder_layer.layer[1]
+
+    def window_size(self):
+        try:
+            size = self.model.config.n_positions
+        except AttributeError:
+            size = 1024
+        return size
+
+    def create_decoder_layer_args(self, hidden_states, attention_mask, encoder_hidden_states,
+            encoder_attention_mask, layer_head_mask, cross_attn_layer_head_mask,
+            past_key_value, output_attentions, position_bias,
+            encoder_decoder_position_bias, use_cache, key, value):
+        args = {'hidden_states': hidden_states,
+            'attention_mask': attention_mask,
+            'position_bias': position_bias,
+            'encoder_hidden_states': encoder_hidden_states,
+            'encoder_attention_mask': encoder_attention_mask,
+            'encoder_decoder_position_bias': encoder_decoder_position_bias,
+            'layer_head_mask': layer_head_mask,
+            'cross_attn_layer_head_mask': cross_attn_layer_head_mask,
+            'past_key_value': (None, None, key, value),
+            'use_cache': use_cache,
+            'output_attentions': output_attentions}
+        if key is None and value is None:
+            args['past_key_value'] = None
+        return args
+
+class UnlimiformerLED(UnlimiformerBART):
+    def __init__(self, model: LEDModel, *args, **kwargs):
+        super().__init__(model, *args, **kwargs)
+
+    def window_size(self):
+        return self.model.config.max_encoder_position_embeddings
+
+# class UnlimiformerLLaMa(Unlimiformer[LlamaModel]):
+#     def __init__(self, model: LlamaModel, *args, **kwargs):
+#         super().__init__(model, *args, **kwargs)
+    
+#     def get_kv_projections(self, layer_begin, layer_end): 
+#         return [
+#             [layer.self_attn.k_proj, layer.self_attn.v_proj]
+#             for layer in self.model.base_model.layers[layer_begin:layer_end]
+#         ]
+
+#     def activation_to_capture(self, layer_begin, layer_end): 
+#         if self.use_datastore:
+#             return [
+#                 layer.input_layernorm
+#                 for layer in self.model.base_model.layers[layer_begin:layer_end]
+#             ]
+#         else:
+#             return self.get_kv_projections(layer_begin, layer_end)
+
+#     def attention_op_to_run(self, layer_begin, layer_end):
+#         return [
+#             layer.self_attn.q_proj
+#                 for layer in self.model.base_model.layers[layer_begin:layer_end]
+#         ]
+
+#     def attention_layer_to_run(self, layer_begin, layer_end): 
+#         return self.model.base_model.layers[layer_begin:layer_end]
+
+#     def self_attention(self, decoder_layer):
+#         return decoder_layer.self_attn 
+
+#     def cross_attention(self, decoder_layer):
+#         return decoder_layer.self_attn
+    
+#     def window_size(self):
+#         return self.model.config.max_position_embeddings
+    
+#     def set_gradient_checkpointing(self, value):
+#         self.model.base_model.gradient_checkpointing = value
+
+#     def process_key_value(self, capturers):
+#         key_capturer, value_capturer = capturers
+#         # (batch, time, heads * attn_dim)
+#         key, value = key_capturer.captured, value_capturer.captured
+#         attention = self.model.base_model.layers[-1].self_attn
+
+#         # (batch, heads, time, attn_dim)
+#         key = key.view(key.shape[0], -1, attention.num_heads, attention.head_dim).transpose(1, 2).contiguous()
+#         value = value.view(value.shape[0], -1, attention.num_heads, attention.head_dim).transpose(1, 2).contiguous()
+        
+#         return key, value
+    
+#     def process_query(self, output):
+#         # output: (batch, time, heads * attn_dim)
+#         attention = self.model.base_model.layers[-1].self_attn
+        
+#         # query: (batch, time, heads, attn_dim)
+#         query = output.view(output.shape[0], output.shape[1], attention.num_heads, attention.head_dim).contiguous()
+#         return query
+
+#     def rotate_half(self, 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)
+
+#     def preprocess_query(self, query, k_proj_weight):
+#         # query: (batch * time, head, dim)
+#         attention = self.model.base_model.layers[-1].self_attn
+#         num_generated = min(self.input_ids_size - self.prompt_input_ids.shape[1], self.actual_model_window_size)
+#         cos, sin = attention.rotary_emb(query, seq_len=num_generated)
+#         cos = cos[:,:,-1]  # [1, 1, dim]
+#         sin = sin[:,:,-1]  # [1, 1, dim]
+#         # cos = cos[-1].unsqueeze(0).unsqueeze(0)  # [bs, 1, seq_len, dim]
+#         # sin = sin[-1].unsqueeze(0)  # [bs, 1, seq_len, dim]
+#         query = (query * cos) + (self.rotate_half(query) * sin)
+        
+#         k_proj = k_proj_weight.view(1, self.num_heads, query.shape[-1], k_proj_weight.shape[0]) # (1, num_heads, attn_dim, embed_dim)
+#         k_proj_l = k_proj[..., :k_proj.shape[-2] // 2, :]
+#         k_proj_r = k_proj[..., k_proj.shape[-2] // 2:, :]
+#         k_proj_rotated = torch.cat([-k_proj_l, k_proj_r], dim=-2)
+
+#         datastore_query = query.unsqueeze(-2) # (batch * beam, num_heads, 1, attn_dim)
+#         datastore_query = torch.matmul(datastore_query, k_proj + k_proj_rotated) # (batch * beam, num_heads, 1, embed_dim)
+#         datastore_query = datastore_query.squeeze(-2)  # (batch * beam, num_heads, embed_dim)
+#         return datastore_query
+
+#     def post_process_retrieved(self, embeddings, k_proj_layer, v_proj_layer, top_search_key_indices):
+#         embed_dim = embeddings.shape[-1]
+#         k_weight = k_proj_layer.weight.view(1, 1, self.num_heads, embed_dim // self.num_heads, embed_dim).transpose(-2,-1) # (1, 1, heads, embed_dim, attn_dim)
+#         k_bias = 0
+#         if k_proj_layer.bias is not None:
+#             k_bias = k_proj_layer.bias.view(1, self.num_heads, embed_dim // self.num_heads).unsqueeze(-2).unsqueeze(0)
+#         v_weight = v_proj_layer.weight.view(1, 1, self.num_heads, embed_dim // self.num_heads, embed_dim).transpose(-2,-1)  # (1, heads, embed_dim, attn_dim)
+#         v_bias = 0
+#         if v_proj_layer.bias is not None:
+#             v_bias = v_proj_layer.bias.view(1, self.num_heads, embed_dim // self.num_heads).unsqueeze(-2).unsqueeze(0)
+#         # new_keys, new_values: (batch, beam, head, encoder_len, attn_dim)
+#         retrieved_keys = torch.matmul(embeddings, k_weight) + k_bias # (beam, head, encoder_len, embed_dim)
+#         retrieved_values = torch.matmul(embeddings, v_weight) + v_bias # (beam, head, encoder_len, embed_dim)
+
+#         attention = self.model.base_model.layers[-1].self_attn
+#         cos, sin = attention.rotary_emb(retrieved_values, seq_len=self.hidden_states[0].shape[1])
+#         cos = cos.squeeze(1).squeeze(0)  # [seq_len, dim]
+#         sin = sin.squeeze(1).squeeze(0)  # [seq_len, dim]
+#         if self.prompt_input_ids.shape[1] > self.actual_model_window_size:
+#             # scale the top key indices to the actual model window size, such that the model will not see
+#             # positional embeddings that did not appear at training time
+#             scaled_key_indices = ((top_search_key_indices / self.prompt_input_ids.shape[1]) * self.actual_model_window_size).int()
+#         else:
+#             scaled_key_indices = top_search_key_indices
+#         # top_search_key_indices = top_search_key_indices.to(cos.device)
+#         scaled_key_indices = scaled_key_indices.to(cos.device)
+#         cos = cos[scaled_key_indices]  # [bs, 1, seq_len, dim]
+#         sin = sin[scaled_key_indices]  # [bs, 1, seq_len, dim]
+#         retrieved_keys = (retrieved_keys * cos) + (self.rotate_half(retrieved_keys) * sin)
+#         return retrieved_keys, retrieved_values
+        
+
+class ActivationCapturer(nn.Module):
+    def __init__(self, layer, capture_input=False):
+        super().__init__()
+        self.layer = layer
+        self.capture_input = capture_input
+
+        self.captured = None
+
+    def unwrap_tuple(self, t):
+        if isinstance(t, tuple) and len(t) == 1:
+            t = t[0]
+        return t
+    
+    def forward(self, module, layer_input, layer_output):
+        if self.capture_input:
+            self.captured = self.unwrap_tuple(layer_input)
+        else:
+            self.captured = self.unwrap_tuple(layer_output)
+    

unlimiformer/random_training_unlimiformer.py

@@ -0,0 +1,224 @@
+import contextlib
+import numpy as np
+import torch
+from torch import nn
+from enum import Enum, auto
+from .model import Unlimiformer, ModelType, UnlimiformerBART, UnlimiformerT5, UnlimiformerLED
+from transformers import BartModel, BartForConditionalGeneration, \
+    T5Model, T5ForConditionalGeneration, \
+    LEDModel, LEDForConditionalGeneration, \
+    AutoModelForSeq2SeqLM
+
+class RandomTrainingUnlimiformer(Unlimiformer[ModelType]):
+    def __init__(self, model: ModelType, *args, **kwargs):
+        super().__init__(model, *args, **kwargs)
+        self.training_hooks_injected = False
+        self.train_step = 0
+
+    @classmethod
+    def convert_model(cls, model, *args, **kwargs):
+        # model_clone = AutoModelForSeq2SeqLM.from_config(model.config)
+        # model_clone.load_state_dict(model.state_dict()).to(args.device)
+        type_to_class = {
+            BartModel: RandomUnlimiformerBART,
+            BartForConditionalGeneration: RandomUnlimiformerBART,
+            T5Model: RandomUnlimiformerT5,
+            T5ForConditionalGeneration: RandomUnlimiformerT5,
+            LEDModel: RandomUnlimiformerLED,
+            LEDForConditionalGeneration: RandomUnlimiformerLED,
+        }
+        type_to_class[type(model)](model, *args, **kwargs)
+        return model
+
+    def pre_eval_hook(self):
+        self.remove_training_hooks(self.model)
+        self.inject_hooks(self.model)   
+        self.original_model_eval_func()
+
+    def pre_train_hook(self, mode=True):
+        # mode=True means model.train() is called
+        # mode=False means model.eval() is called
+        torch.cuda.empty_cache()
+        if mode is True:
+            self.break_out(self.model)
+            self.remove_training_hooks(self.model)
+            if self.unlimiformer_training and self.train_step % 2 == 0:
+                super().inject_training_hooks(self.model)
+            else:
+                self.inject_training_hooks(self.model)
+            self.train_step += 1
+        self.original_model_train_func(mode)
+    
+    def inject_training_hooks(self, model):
+        if self.training_hooks_injected:
+            return
+        # self.original_forward_func = model.forward
+        model.forward = self.random_inputs_forward_hook
+
+        decoder_layers_to_run = self.attention_layer_to_run(self.layer_begin, self.layer_end)
+        
+        self.original_decoder_layer_self_attn_forward_funcs = []
+        for decoder_layer in decoder_layers_to_run:
+            attention = self.self_attention(decoder_layer)
+            self.original_decoder_layer_self_attn_forward_funcs.append(attention.forward)
+            attention.forward = self.create_self_attn_random_pre_forward_hook(attention.forward)
+
+        self.original_decoder_layer_forward_funcs = []
+        for decoder_layer in decoder_layers_to_run:
+            self.original_decoder_layer_forward_funcs.append(decoder_layer.forward)
+            decoder_layer.forward = self.create_decoder_layer_random_func(decoder_layer.forward, decoder_layer)
+
+        self.original_decoder_layer_cross_attn_forward_funcs = []
+        for i, decoder_layer in enumerate(decoder_layers_to_run):
+            attention = self.cross_attention(decoder_layer)
+            self.original_decoder_layer_cross_attn_forward_funcs.append(attention.forward)
+
+        self.inject_hooks_for_unaffected_layers(model, decoder_layers_to_run)
+
+        self.training_hooks_injected = True
+
+    def create_self_attn_random_pre_forward_hook(self, original_self_attn_forward_func):
+        def self_attention_pre_forward_hook(*args, **kwargs):
+            kwargs['past_key_value'] = None
+            return original_self_attn_forward_func(*args, **kwargs)
+        
+        return self_attention_pre_forward_hook
+
+    def create_decoder_layer_random_func(self, decoder_layer_original_forward_func, decoder_layer):
+        def checkpointed_decoder_layer(
+                hidden_states: torch.Tensor,
+                attention_mask=None,
+                encoder_hidden_states=None,
+                encoder_attention_mask=None,
+                layer_head_mask=None,
+                cross_attn_layer_head_mask=None,
+                past_key_value=None,
+                output_attentions=False,
+                position_bias=None,
+                encoder_decoder_position_bias=None,
+                use_cache=True):
+
+            
+            
+            def sample_and_forward(hidden_states, attention_mask, 
+                    encoder_hidden_states, encoder_attention_mask, layer_head_mask, 
+                    cross_attn_layer_head_mask, past_key_value, 
+                    output_attentions, use_cache, long_inputs, long_inputs_mask, rand_indices,
+                    position_bias, encoder_decoder_position_bias):
+                
+                sampled_input, _ = self.sample_long_input(long_inputs, long_inputs_mask, rand_indices)
+                key, value = self.create_key_value(sampled_input, decoder_layer)
+                decoder_layer_args = self.create_decoder_layer_args(
+                    hidden_states=hidden_states,
+                    attention_mask=attention_mask,
+                    encoder_hidden_states=encoder_hidden_states,
+                    encoder_attention_mask=encoder_attention_mask,
+                    layer_head_mask=layer_head_mask,
+                    cross_attn_layer_head_mask=cross_attn_layer_head_mask,
+                    past_key_value=past_key_value,
+                    output_attentions=output_attentions,
+                    position_bias=position_bias,
+                    encoder_decoder_position_bias=encoder_decoder_position_bias,
+                    use_cache=use_cache,
+                    key=key,value=value
+                )
+                return decoder_layer_original_forward_func(**decoder_layer_args)
+                    
+
+            with torch.no_grad():
+                # This sampling must be done outside of the checkpoint, to ensure that the same sampling happens
+                # both in "forward" and "backward" passes
+                rand_indices = self.sample_random_indices()
+
+            return torch.utils.checkpoint.checkpoint(
+                sample_and_forward, hidden_states, attention_mask, 
+                encoder_hidden_states, encoder_attention_mask, layer_head_mask, 
+                cross_attn_layer_head_mask, None, 
+                output_attentions, use_cache, self.long_inputs_encoded, self.long_inputs_mask, rand_indices,
+                position_bias, encoder_decoder_position_bias)
+
+        return checkpointed_decoder_layer
+
+    def sample_random_indices(self):
+        rand_indices_list = []
+        seq_lens = self.long_inputs_mask.sum(-1).tolist()
+        for seq_len in seq_lens:
+            if seq_len < self.actual_model_window_size:
+                rand_indices = torch.arange(self.actual_model_window_size).to(self.device)
+                rand_indices_list.append(rand_indices)
+                continue
+            
+            rand_indices = torch.torch.randperm(seq_len)[:self.actual_model_window_size].to(self.device)
+            if seq_len < self.actual_model_window_size:
+                padding = max(self.actual_model_window_size - seq_len, 0)
+                rand_indices = torch.cat([rand_indices, torch.arange(padding).to(self.device) + seq_len], axis=-1).to(self.device)
+            rand_indices_list.append(rand_indices)
+        rand_indices = torch.stack(rand_indices_list, dim=0)
+        return rand_indices
+
+    def random_inputs_forward_hook(self, input_ids=None, attention_mask=None, labels=None, **kwargs):
+        self.model.base_model.decoder.gradient_checkpointing = False
+        self.long_inputs_encoded, self.long_inputs_mask = self.chunked_encode_input(input_ids=input_ids, attention_mask=attention_mask)
+
+        #  TODO: should the inputs be sampled or the truncated beginning?
+        # if self.random_knn_initial_inputs:
+        #     encoded_inputs, encoded_inputs_mask = self.sample_long_input(self.long_inputs_encoded, self.long_inputs_mask)
+        # else:
+        encoded_inputs = self.long_inputs_encoded[:, :self.actual_model_window_size]
+        encoded_inputs_mask = self.long_inputs_mask[:, :self.actual_model_window_size]
+        return self.original_forward_func(encoder_outputs=(encoded_inputs, ), labels=labels, attention_mask=encoded_inputs_mask, **kwargs)
+
+    def sample_long_input(self, long_inputs_encoded, long_inputs_mask, random_indices=None):
+        if long_inputs_mask.shape[-1] < self.actual_model_window_size:
+            return long_inputs_encoded, long_inputs_mask
+        batch_size = long_inputs_encoded.shape[0]
+        
+        if random_indices is None:
+            random_indices = self.sample_random_indices()
+        random_mask = torch.zeros_like(long_inputs_mask).to(self.device) \
+            .scatter_(dim=-1, index=random_indices, src=torch.ones_like(random_indices)).bool().to(self.device)
+        sampled_input = long_inputs_encoded[random_mask].reshape(batch_size, self.actual_model_window_size, -1).to(self.device)
+        sampled_mask = long_inputs_mask[random_mask].reshape(batch_size, self.actual_model_window_size).to(self.device)
+        return sampled_input, sampled_mask
+
+    def chunked_encode_input(self, input_ids, attention_mask):
+        long_inputs_encoded = []
+        long_inputs_mask = []
+        window_indices = self.window_indices(input_ids.shape[-1])
+
+        self.is_input_encoding_pass = True
+        for context_start_ind, context_end_ind, update_start_ind, update_end_ind in window_indices:
+            chunk = input_ids[:, context_start_ind:context_end_ind]
+            chunk_attention_mask = attention_mask[:, context_start_ind:context_end_ind]
+            output = self.model.base_model.encoder(chunk, attention_mask=chunk_attention_mask, return_dict=True, output_hidden_states=True)
+            encoder_last_hidden_state = output.last_hidden_state # (batch, time, dim)
+            
+            # list of (batch, head, chunked_time, dim)
+            encoder_last_hidden_state = encoder_last_hidden_state[:, update_start_ind:update_end_ind] # (batch, chunked_time, dim)
+            chunk_attention_mask = chunk_attention_mask[:, update_start_ind:update_end_ind] # (batch, chunked_time)
+
+            long_inputs_encoded.append(encoder_last_hidden_state) # (batch, chunked_source_len, dim)
+            long_inputs_mask.append(chunk_attention_mask) # (batch, chunked_source_len)
+        
+        long_inputs_encoded = torch.cat(long_inputs_encoded, dim=1) # (batch, source_len, dim)
+        long_inputs_mask = torch.cat(long_inputs_mask, dim=1) # (batch, source_len)
+
+        self.is_input_encoding_pass = False
+        if self.verbose:
+            print(f'Input: '
+                f'{self.tokenizer.decode(input_ids[0][:self.actual_model_window_size], skip_special_tokens=True)} ||| '
+                f'{self.tokenizer.decode(input_ids[0][self.actual_model_window_size:], skip_special_tokens=True)}')
+            print()
+        return long_inputs_encoded, long_inputs_mask
+
+class RandomUnlimiformerBART(RandomTrainingUnlimiformer[BartModel], UnlimiformerBART):
+    def __init__(self, model: BartModel, *args, **kwargs):
+        super().__init__(model, *args, **kwargs)
+
+class RandomUnlimiformerT5(RandomTrainingUnlimiformer[T5Model], UnlimiformerT5):
+    def __init__(self, model: T5Model, *args, **kwargs):
+        super().__init__(model, *args, **kwargs)
+
+class RandomUnlimiformerLED(RandomTrainingUnlimiformer[LEDModel], UnlimiformerLED):
+    def __init__(self, model: LEDModel, *args, **kwargs):
+        super().__init__(model, *args, **kwargs)

unlimiformer/run.py

@@ -0,0 +1,1180 @@
+#!/usr/bin/env python
+# coding=utf-8
+# Copyright 2021 The HuggingFace 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.
+"""
+Fine-tuning the library models for sequence to sequence.
+"""
+# You can also adapt this script on your own sequence to sequence task. Pointers for this are left as comments.
+import logging
+import os
+import sys
+
+import numpy as np
+from unlimiformer import Unlimiformer
+from random_training_unlimiformer import RandomTrainingUnlimiformer
+
+import nltk
+
+# we import the logging frameworks before any other import to make sure all monkey patching for the logging are active
+# from sled import SledConfig
+
+import wandb
+import torch
+
+sys.path.insert(0, os.path.dirname(__file__))  # seq2seq package path
+sys.path.insert(0, os.getcwd())
+
+from dataclasses import dataclass, field, replace
+from typing import List, Optional
+import json
+from copy import deepcopy
+import torch.nn.functional as F
+
+import datasets
+
+import transformers
+from transformers import (
+    AutoConfig,
+    AutoModelForSeq2SeqLM,
+    AutoTokenizer,
+    EarlyStoppingCallback,
+    set_seed, WEIGHTS_NAME,
+)
+from transformers.trainer_utils import get_last_checkpoint
+from transformers import DataCollatorForSeq2Seq
+
+from datasets import load_dataset
+
+# noinspection PyUnresolvedReferences
+# import sled  # *** required so that SledModels will be registered for the AutoClasses ***
+
+from utils.config import handle_args_to_ignore
+from utils.decoding import decode
+from metrics import load_metric
+from utils.duplicates import drop_duplicates_in_input
+from utils.override_training_args import TrainingOverridesArguments
+from utils.custom_seq2seq_trainer import CustomTrainer
+from utils.custom_hf_argument_parser import CustomHfArgumentParser
+from metrics.metrics import HFMetricWrapper, MetricCollection
+
+logger = logging.getLogger('sled')
+
+PREFIX_DOC_SEP = '\n\n'
+
+DEBUG = os.environ.get('DEBUG', 'false').lower() in {'1', 'true', 'yes'}  # If set, will set some configuration to help debug
+if DEBUG:
+    assert not torch.cuda.is_available() or torch.cuda.device_count() == 1
+
+
+@dataclass
+class ModelArguments:
+    """
+    Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.
+    """
+
+    model_name_or_path: str = field(
+        default=None, metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}
+    )
+    config_name: Optional[str] = field(
+        default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
+    )
+    tokenizer_name: Optional[str] = field(
+        default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
+    )
+    cache_dir: Optional[str] = field(
+        default=None,
+        metadata={"help": "Where to store the pretrained models downloaded from huggingface.co"},
+    )
+    use_fast_tokenizer: bool = field(
+        default=True,
+        metadata={"help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."},
+    )
+    model_revision: str = field(
+        default="main",
+        metadata={"help": "The specific model version to use (can be a branch name, tag name or commit id)."},
+    )
+    drop_duplicates_in_eval: bool = field(
+        default=True,
+    )
+
+    def __post_init__(self):
+        pass
+    
+
+
+@dataclass
+class DataTrainingArguments:
+    """
+    Arguments pertaining to what data we are going to input our model for training and eval.
+    """
+
+    dataset_name: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": "The name of the dataset to use (via the datasets library) or name of the file in src/data."
+        },
+    )
+    dataset_config_name: Optional[str] = field(
+        default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}
+    )
+    metric_names: Optional[List[str]] = field(
+        default=None,
+        metadata={"help": "The name of the metric to use (from src/metrics)."},
+    )
+    input_column: Optional[str] = field(
+        default=None,
+        metadata={"help": "The name of the column in the datasets containing the full texts (for summarization)."},
+    )
+    input_prefix_column: Optional[str] = field(
+        default=None,
+        metadata={"help": "The name of the column in the datasets containing the input prefix (e.g. questions), when those exist."},
+    )
+    output_column: Optional[str] = field(
+        default=None,
+        metadata={"help": "The name of the column in the datasets containing the summaries (for summarization)."},
+    )
+    train_file: Optional[str] = field(
+        default=None, metadata={"help": "The input training data file (a jsonlines or csv file)."}
+    )
+    validation_file: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": "An optional input evaluation data file to evaluate the metrics (rouge) on "
+            "(a jsonlines or csv file)."
+        },
+    )
+    test_file: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": "An optional input test data file to evaluate the metrics (rouge) on " "(a jsonlines or csv file)."
+        },
+    )
+    overwrite_cache: bool = field(
+        default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
+    )
+    preprocessing_num_workers: Optional[int] = field(
+        default=None,
+        metadata={"help": "The number of processes to use for the preprocessing."},
+    )
+    max_source_length: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": "The maximum total input sequence length after tokenization. Sequences longer "
+            "than this will be truncated, sequences shorter will be padded."
+        },
+    )
+    eval_max_source_length: Optional[int] = field(
+        default=None,
+        metadata={"help": "if None, will be same as max_source_length"},
+    )
+    max_prefix_length: Optional[int] = field(
+        default=0,
+        metadata={
+            "help": "The maximum total input_prefix sequence length after tokenization. Sequences longer "
+            "than this will be truncated, sequences shorter will be padded from the left "
+                    "(only used if prefixes are not merged)."
+        },
+    )
+    max_target_length: Optional[int] = field(
+        default=128,
+        metadata={
+            "help": "The maximum total sequence length for target text after tokenization. Sequences longer "
+            "than this will be truncated, sequences shorter will be padded."
+        },
+    )
+    val_max_target_length: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": "The maximum total sequence length for validation target text after tokenization. Sequences longer "
+            "than this will be truncated, sequences shorter will be padded. Will default to `max_target_length`."
+            "This argument is also used to override the ``max_length`` param of ``model.generate``, which is used "
+            "during ``evaluate`` and ``predict``."
+        },
+    )
+    pad_to_max_length: bool = field(
+        default=False,
+        metadata={
+            "help": "Whether to pad all samples to model maximum sentence length. "
+            "If False, will pad the samples dynamically when batching to the maximum length in the batch. More "
+            "efficient on GPU but very bad for TPU."
+        },
+    )
+    max_train_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": "For debugging purposes or quicker training, truncate the number of training examples to this "
+            "value if set."
+        },
+    )
+    max_eval_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": "For debugging purposes or quicker training, truncate the number of evaluation examples to this "
+            "value if set."
+        },
+    )
+    max_predict_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": "For debugging purposes or quicker training, truncate the number of prediction examples to this "
+            "value if set."
+        },
+    )
+    num_beams: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": "Number of beams to use for evaluation. This argument will be passed to ``model.generate``, "
+            "which is used during ``evaluate`` and ``predict``."
+        },
+    )
+    ignore_pad_token_for_loss: bool = field(
+        default=True,
+        metadata={
+            "help": "Whether to ignore the tokens corresponding to padded labels in the loss computation or not."
+        },
+    )
+    source_prefix: Optional[str] = field(
+        default=None, metadata={"help": "A prefix to add before every source text (useful for T5 models)."}
+    )
+    data_dir: Optional[str] = field(
+        default=None,
+        metadata={"help": "Defining the data_dir of the dataset configuration."},
+    )
+    download_mode: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": "Defining the download_mode when loading the dataset. Options are `reuse_dataset_if_exists` (default), `reuse_cache_if_exists` and `force_redownload`."
+        },
+    )
+    evaluate_on_training_data: bool = field(
+        default=False,
+        metadata={"help": "Whether to evaluate on training data or not, to make sure the model can overfit."},
+    )
+    folder_suffix: str = field(
+        default="",
+        metadata={"help": "args to be suffixes for the output folder of the run"},
+    )
+    preprocess_only: bool = field(
+        default=False,
+        metadata={"help": "Preprocess only: Don't start training, just do the things before"},
+    )
+    assign_zero_to_too_long_val_examples: bool = field(
+        default=False,
+        metadata={
+            "help": "If true, all sequences longer then max_source_length will be assign a score of 0 in the metric evaluation"
+        },
+    )
+    shared_storage: bool = field(
+        default=True,
+        metadata={"help": "Whether nodes share the same storage"},
+    )
+    trim_very_long_strings: bool = field(
+        default=False,
+        metadata={"help": "Whether to trim very long strings before tokenizing them"},
+    )
+    pad_prefix: bool = field(
+        default=False,
+        metadata={
+            "help": "Whether to pad the prefix if it exists to max_prefix_length. "
+                    "Note - important if you are using a SLED model on an input that contains an input_prefix"
+        },
+    )
+    test_start_ind: Optional[int] = field(
+        default=None,
+        metadata={"help": "if given, uses the test set starting from this index"},
+    )
+    test_end_ind: Optional[int] = field(
+        default=None,
+        metadata={"help": "if given, uses the test set ending at this index"},
+    )
+    # Uri:
+    patience: Optional[int] = field(
+        default=None,
+    )
+    length_penalty: Optional[float] = field(
+        default=1.0,
+    )
+    extra_metrics: Optional[List[str]] = field(
+        default=None,
+        metadata={"help": "The name of the metric to use (from src/metrics)."},
+    )
+    chunked_training_size: Optional[int] = field(
+        default=None,
+    )
+    oracle_training: Optional[bool] = field(
+        default=False,
+        metadata={"help": "If True, train on the input sentences that provide the highest ROUGE score with the labels"}
+    )
+    oracle_merge: Optional[bool] = field(
+        default=False,
+        metadata={"help": "If True, merge the oracle dataset and the standard training dataset"}
+    )
+    def __post_init__(self):
+        if self.val_max_target_length is None:
+            self.val_max_target_length = self.max_target_length
+        if self.pad_prefix and self.max_prefix_length == 0:
+            raise ValueError('When padding prefix, you must set a max_prefix_length')
+        assert self.max_prefix_length == 0 or self.max_prefix_length <= 0.5*self.max_source_length,\
+            'If max_prefix_length is given, it must be much shorter than the total input'
+        # Uri: 
+        if self.eval_max_source_length is None:
+            self.eval_max_source_length = self.max_source_length
+
+
+@dataclass
+class UnlimiformerArguments:
+    """
+    Arguments pertaining to what data we are going to input our model for training and eval.
+    """
+    test_unlimiformer: Optional[bool] = field(
+        default=False,
+        metadata={
+            "help": "whether to use KNN."
+        },
+    )
+    unlimiformer_verbose: Optional[bool] = field(
+        default=False,
+        metadata={
+            "help": "whether to print KNN intermediate predictions (mostly for debugging)."
+        },
+    )
+    layer_begin: Optional[int] = field(
+        default=0,
+        metadata={"help": "The layer to begin applying KNN to. KNN will be applied to layers[knn_layer_begin:layer_end]. "
+                          "By default, it will be applied to all layers: [0:None]]"}, 
+    )
+    layer_end: Optional[int] = field(
+        default=None,
+        metadata={"help": "The layer to end applying KNN to. KNN will be applied to layers[knn_layer_begin:layer_end]. "
+                          "By default, it will be applied to all layers: [0:None]]"}, 
+    )
+    unlimiformer_chunk_overlap: Optional[float] = field(
+        default=0.5,
+        metadata={"help": "The fraction of overlap between input chunks"},
+    )
+    unlimiformer_chunk_size: Optional[int] = field(
+        default=None,
+        metadata={"help": "The size of each input chunk"},
+    )
+    unlimiformer_head_num: Optional[int] = field(
+        default=None,
+        metadata={"help": "The head to apply KNN to (if None, apply to all heads)"},
+    )
+    unlimiformer_exclude: Optional[bool] = field(
+        default=False,
+        metadata={
+            "help": "If True, prioritize the inputs that are **not** in the standard attention window."
+        },
+    )
+    random_unlimiformer_training: Optional[bool] = field(
+        default=False,
+    )
+    unlimiformer_training: Optional[bool] = field(
+        default=False,
+    )
+    use_datastore: Optional[bool] = field(default=False)
+    flat_index: Optional[bool] = field(default=False)
+    test_datastore: Optional[bool] = field(default=False)
+    reconstruct_embeddings: Optional[bool] = field(default=False)
+    gpu_datastore: Optional[bool] = field(default=True)
+    gpu_index: Optional[bool] = field(default=True)
+
+
+def main():
+    handle_args_to_ignore(sys.argv)  # Just for sweeps
+
+    # See all possible arguments in src/transformers/training_args.py
+    # or by passing the --help flag to this script.
+    # We now keep distinct sets of args, for a cleaner separation of concerns.
+
+    parser = CustomHfArgumentParser((ModelArguments, DataTrainingArguments, TrainingOverridesArguments, UnlimiformerArguments))
+    model_args, data_args, training_args, unlimiformer_args = parser.parse_dictionary_and_args()
+    
+    set_up_logging(training_args)
+    logger.info(f"Training Arguments: {training_args}")
+    logger.info(f"Data Arguments: {data_args}")
+    logger.info(f"Model Arguments: {model_args}")
+    logger.info(f"Unlimiformer Arguments: {unlimiformer_args}")
+
+
+    # Added to avoid wandb.errors.UsageError: Error communicating with wandb process
+    wandb.init(settings=wandb.Settings(start_method="fork"), name=training_args.output_dir)
+
+    # Used to find missing dependencies early on
+    load_metric(data_args.metric_names, **locals())
+    load_extra_metrics(data_args.extra_metrics)
+
+    if data_args.source_prefix is None and model_args.model_name_or_path in [
+        "t5-small",
+        "t5-base",
+        "t5-large",
+        "t5-3b",
+        "t5-11b",
+    ]:
+        logger.warning(
+            "You're running a t5 model but didn't provide a source prefix, which is the expected, e.g. with "
+            "`--source_prefix 'summarize: ' `"
+        )
+
+    # Detecting last checkpoint.
+    last_checkpoint = _detect_last_checkpoint(training_args)
+
+    # Set seed before initializing model.
+    set_seed(training_args.seed)
+
+    seq2seq_dataset = _get_dataset(data_args, model_args, training_args)
+
+    # Load pretrained model and tokenizer
+    #
+    # Distributed training:
+    # The .from_pretrained methods guarantee that only one local process can concurrently
+    # download model & vocab.
+    config_name = None
+    if model_args.config_name:
+        config_name = model_args.config_name
+    else:
+        if os.path.isfile(model_args.model_name_or_path):
+            config_name = os.path.dirname(model_args.model_name_or_path)
+        else:
+            config_name = model_args.model_name_or_path
+
+    config_overrides = {}
+    if training_args.gradient_checkpointing is not None:
+        config_overrides["gradient_checkpointing"] = training_args.gradient_checkpointing
+
+    config = AutoConfig.from_pretrained(
+        config_name,
+        cache_dir=model_args.cache_dir,
+        revision=model_args.model_revision,
+        use_auth_token=training_args.use_auth_token,
+        **config_overrides
+    )
+    # override for sled models to make sure we are explicit in our request
+    # if isinstance(config, SledConfig) and (not data_args.pad_prefix or data_args.max_prefix_length == 0):
+    #     logger.warning('Setting prepend_prefix to False if using a SLED model, as the input does not have a prefix or '
+    #                    'pad_prefix is False (all prefixes must be of the same length for SLED). If you do not use SLED '
+    #                    'or finetune on a dataset with no prefixes, ignore this warning')
+    #     config.prepend_prefix = False
+
+    if model_args.model_name_or_path is None:
+        # Padding for divisibility by 8
+        if config.vocab_size % 8 != 0 and training_args.fp16_padding:
+            config.vocab_size += 8 - (config.vocab_size % 8)
+
+    tokenizer_name = None
+    if model_args.tokenizer_name:
+        tokenizer_name = model_args.tokenizer_name
+    else:
+        if os.path.isfile(model_args.model_name_or_path):
+            tokenizer_name = os.path.dirname(model_args.model_name_or_path)
+        else:
+            tokenizer_name = model_args.model_name_or_path
+    tokenizer = AutoTokenizer.from_pretrained(
+        tokenizer_name,
+        cache_dir=model_args.cache_dir,
+        use_fast=model_args.use_fast_tokenizer,
+        revision=model_args.model_revision,
+        use_auth_token=training_args.use_auth_token,
+    )
+    if model_args.model_name_or_path is not None:
+        model = AutoModelForSeq2SeqLM.from_pretrained(
+            model_args.model_name_or_path,
+            from_tf=bool(".ckpt" in model_args.model_name_or_path),
+            config=config,
+            cache_dir=model_args.cache_dir,
+            revision=model_args.model_revision,
+            use_auth_token=training_args.use_auth_token,
+        )
+    else:
+        model = AutoModelForSeq2SeqLM.from_config(
+            config,
+        )
+    if unlimiformer_args.test_unlimiformer:
+        unlimiformer_kwargs = {
+            'layer_begin': unlimiformer_args.layer_begin, 
+            'layer_end': unlimiformer_args.layer_end,
+            'unlimiformer_head_num': unlimiformer_args.unlimiformer_head_num, 
+            'exclude_attention': unlimiformer_args.unlimiformer_exclude, 
+            'chunk_overlap': unlimiformer_args.unlimiformer_chunk_overlap,
+            'model_encoder_max_len': unlimiformer_args.unlimiformer_chunk_size,
+            'verbose': unlimiformer_args.unlimiformer_verbose, 'tokenizer': tokenizer,
+            'unlimiformer_training': unlimiformer_args.unlimiformer_training,
+            'use_datastore': unlimiformer_args.use_datastore,
+            'flat_index': unlimiformer_args.flat_index,
+            'test_datastore': unlimiformer_args.test_datastore,
+            'reconstruct_embeddings': unlimiformer_args.reconstruct_embeddings,
+            'gpu_datastore': unlimiformer_args.gpu_datastore,
+            'gpu_index': unlimiformer_args.gpu_index
+        }
+        if unlimiformer_args.random_unlimiformer_training:
+            model = RandomTrainingUnlimiformer.convert_model(model, **unlimiformer_kwargs)
+        else:
+            model = Unlimiformer.convert_model(model, **unlimiformer_kwargs)
+
+    model.config.use_cache = True
+    if training_args.gradient_checkpointing and getattr(model.config, 'use_cache', False) and training_args.do_train:
+        logger.warning('Cannot use cache in models when using gradient checkpointing. turning it off')
+        model.config.use_cache = False
+
+    model.resize_token_embeddings(len(tokenizer))
+
+    if model.config.decoder_start_token_id is None:
+        raise ValueError("Make sure that `config.decoder_start_token_id` is correctly defined")
+
+    prefix = data_args.source_prefix if data_args.source_prefix is not None else ""
+
+    # Preprocessing the datasets.
+    # We need to tokenize inputs and targets.
+    if training_args.do_train:
+        column_names = seq2seq_dataset["train"].column_names
+    elif training_args.do_eval:
+        column_names = seq2seq_dataset["validation"].column_names
+    elif training_args.do_predict:
+        column_names = seq2seq_dataset["test"].column_names
+    else:
+        logger.info("There is nothing to do. Please pass `do_train`, `do_eval` and/or `do_predict`.")
+        return
+
+    # Get the column names for input/target.
+    if data_args.input_column is None:
+        input_column = "input"
+    else:
+        input_column = data_args.input_column
+        if input_column not in column_names:
+            raise ValueError(
+                f"--input_column' value '{data_args.input_column}' needs to be one of: {', '.join(column_names)}"
+            )
+    if data_args.input_prefix_column is None:
+        input_prefix_column = "input_prefix"
+    else:
+        input_prefix_column = data_args.input_prefix_column
+        if input_prefix_column not in column_names:
+            raise ValueError(
+                f"--input_prefix_column' value '{data_args.input_prefix_column}' needs to be one of: {', '.join(column_names)}"
+            )
+    if data_args.output_column is None:
+        output_column = "output"
+    else:
+        output_column = data_args.output_column
+        if output_column not in column_names:
+            raise ValueError(
+                f"--output_column' value '{data_args.output_column}' needs to be one of: {', '.join(column_names)}"
+            )
+
+    # Temporarily set max_target_length for training.
+    max_target_length = data_args.max_target_length
+    padding = "max_length" if data_args.pad_to_max_length else False
+
+    if training_args.label_smoothing_factor > 0 and not hasattr(model, "prepare_decoder_input_ids_from_labels"):
+        logger.warning(
+            "label_smoothing is enabled but the `prepare_decoder_input_ids_from_labels` method is not defined for"
+            f"`{model.__class__.__name__}`. This will lead to loss being calculated twice and will take up more memory"
+        )
+
+    def preprocess_function_kwargs_fn():
+        return {
+            "tokenizer": deepcopy(tokenizer),
+            "prefix": prefix,
+            "input_column": input_column,
+            "input_prefix_column": input_prefix_column,
+            "output_column": output_column,
+            "max_source_length": data_args.max_source_length,
+            "max_prefix_length": data_args.max_prefix_length,
+            "max_target_length": max_target_length,
+            "prefix_sep": PREFIX_DOC_SEP,
+            "padding": padding,
+            "ignore_pad_token_for_loss": data_args.ignore_pad_token_for_loss,
+            "assign_zero_to_too_long_val_examples": data_args.assign_zero_to_too_long_val_examples,
+            "trim_very_long_strings": data_args.trim_very_long_strings,
+            "pad_prefix": data_args.pad_prefix
+        }
+
+    if training_args.do_train:
+        if "train" not in seq2seq_dataset:
+            raise ValueError("--do_train requires a train dataset")
+        logger.info("")
+        logger.info("Training examples before tokenization:")
+        if input_prefix_column in column_names:
+            logger.info(f"input_prefix #0: {seq2seq_dataset['train'][0][input_prefix_column]}")
+        # logger.info(f"input #0: {seq2seq_dataset['train'][0]['input']}")
+        # logger.info(f"output #0: {seq2seq_dataset['train'][0]['output']}")
+        if input_prefix_column in column_names:
+            logger.info(f"input_prefix #1: {seq2seq_dataset['train'][1][input_prefix_column]}")
+        # logger.info(f"input #1: {seq2seq_dataset['train'][1]['input']}")
+        # logger.info(f"output #1: {seq2seq_dataset['train'][1]['output']}")
+        logger.info("")
+        untokenized_train_dataset = seq2seq_dataset["train"]
+        if data_args.max_train_samples is not None:
+            untokenized_train_dataset = untokenized_train_dataset.select(range(data_args.max_train_samples))
+
+        if DEBUG:
+            # In debug mode, we want to recreate the data
+            data_args.shared_storage = False
+            data_args.overwrite_cache = True
+        with training_args.main_process_first(
+            local=not data_args.shared_storage, desc="train dataset map pre-processing"
+            ):
+
+            if data_args.oracle_training:
+                logger.info("Using oracle training")
+                oracle_processed_dir = f'oracle_input_{data_args.dataset_config_name}'
+                if os.path.isdir(oracle_processed_dir):
+                    logger.info(f"Using oracle training from {oracle_processed_dir}")
+                    oracle_training_set = datasets.load_from_disk(oracle_processed_dir)
+                else:
+                    rouge_scorer = datasets.load_metric('rouge')
+                    oracle_training_set = untokenized_train_dataset.map(
+                        extract_oracle_sent_batch,
+                        fn_kwargs={'max_length': data_args.max_source_length,
+                                'tokenizer': tokenizer,
+                                'rouge_scorer': rouge_scorer},
+                        batched=True,
+                        batch_size=1,
+                        num_proc=data_args.preprocessing_num_workers,
+                        load_from_cache_file=not data_args.overwrite_cache,
+                        desc="Extracting oracle sentences from every training example",
+                    )
+                    oracle_training_set.save_to_disk(oracle_processed_dir)
+                
+                
+                if data_args.oracle_merge:
+                    untokenized_train_dataset = datasets.concatenate_datasets([untokenized_train_dataset, oracle_training_set])
+                    untokenized_train_dataset = untokenized_train_dataset.shuffle(seed=training_args.seed)
+                else:
+                    untokenized_train_dataset = oracle_training_set
+
+            train_dataset = untokenized_train_dataset.map(
+                preprocess_function,
+                fn_kwargs=preprocess_function_kwargs_fn(),
+                batched=True,
+                num_proc=data_args.preprocessing_num_workers,
+                remove_columns=untokenized_train_dataset.column_names,
+                load_from_cache_file=not data_args.overwrite_cache,
+                desc="Running tokenizer on train dataset",
+            )
+
+            if data_args.chunked_training_size is not None:
+                train_dataset = train_dataset.map(
+                    chunk_dataset_function,
+                    fn_kwargs={'chunk_size': data_args.chunked_training_size},
+                    batched=True,
+                    num_proc=data_args.preprocessing_num_workers,
+                    load_from_cache_file=not data_args.overwrite_cache,
+                    desc="Chunking train dataset source",
+                )
+                train_dataset = train_dataset.shuffle(seed=training_args.seed)
+
+    if training_args.do_eval:
+        max_target_length = data_args.val_max_target_length
+        preprocess_function_kwargs = preprocess_function_kwargs_fn()
+        preprocess_function_kwargs["max_target_length"] = max_target_length
+        preprocess_function_kwargs['max_source_length'] = data_args.eval_max_source_length
+        if "validation" not in seq2seq_dataset:
+            raise ValueError("--do_eval requires a validation dataset")
+        logger.info("")
+        logger.info("Validation examples before tokenization:")
+        if input_prefix_column in column_names:
+            logger.info(f"input_prefix #0: {seq2seq_dataset['validation'][0][input_prefix_column]}")
+        # logger.info(f"input #0: {seq2seq_dataset['validation'][0]['input']}")
+        # logger.info(f"output #0: {seq2seq_dataset['validation'][0]['output']}")
+        if input_prefix_column in column_names:
+            logger.info(f"input_prefix #1: {seq2seq_dataset['validation'][1][input_prefix_column]}")
+        # logger.info(f"input #1: {seq2seq_dataset['validation'][1]['input']}")
+        # logger.info(f"output #1: {seq2seq_dataset['validation'][1]['output']}")
+        logger.info("")
+        untokenized_eval_dataset = seq2seq_dataset["validation"]
+        if data_args.max_eval_samples is not None:
+            untokenized_eval_dataset = untokenized_eval_dataset.select(range(data_args.max_eval_samples))
+        if model_args.drop_duplicates_in_eval is True:
+            untokenized_eval_dataset = drop_duplicates_in_input(untokenized_eval_dataset)
+        untokenized_eval_dataset_orig = untokenized_eval_dataset
+        assert training_args.eval_fraction > 0
+        n = len(untokenized_eval_dataset)
+        training_args = replace(training_args, eval_fraction = min(training_args.eval_fraction, n))
+        if training_args.eval_fraction != 1:
+            if training_args.eval_fraction > 1:
+                assert training_args.eval_fraction == int(training_args.eval_fraction)
+                logger.info(f'using predetermined absolute samples from eval set ({training_args.eval_fraction} )')
+                training_args = replace(training_args, eval_fraction = training_args.eval_fraction / n)
+            indices = np.random.permutation(n)[:int(np.ceil(max(1, training_args.eval_fraction * n)))]
+            untokenized_eval_dataset = type(untokenized_eval_dataset).from_dict(untokenized_eval_dataset[indices])
+            logger.info(f'During training, will only use {training_args.eval_fraction:.3%} samples of the eval set '
+                        f'which amounts to {len(untokenized_eval_dataset)} out of {n} samples')
+
+        eval_dataset = process_eval_set(data_args, preprocess_function_kwargs, training_args, untokenized_eval_dataset)
+        eval_dataset_orig = eval_dataset
+        if training_args.eval_fraction < 1:
+            eval_dataset_orig = process_eval_set(data_args, preprocess_function_kwargs, training_args,
+                                                 untokenized_eval_dataset_orig)
+
+    if training_args.do_predict:
+        max_target_length = data_args.val_max_target_length
+        preprocess_function_kwargs = preprocess_function_kwargs_fn()
+        preprocess_function_kwargs["max_target_length"] = max_target_length
+        preprocess_function_kwargs['max_source_length'] = data_args.eval_max_source_length
+        if "test" not in seq2seq_dataset:
+            raise ValueError("--do_predict requires a test dataset")
+        untokenized_predict_dataset = seq2seq_dataset["test"]
+        if data_args.max_predict_samples is not None:
+            untokenized_predict_dataset = untokenized_predict_dataset.select(range(data_args.max_predict_samples))
+        if model_args.drop_duplicates_in_eval is True:
+            untokenized_predict_dataset = drop_duplicates_in_input(untokenized_predict_dataset)
+
+        if output_column in untokenized_predict_dataset.column_names:
+            untokenized_predict_dataset = untokenized_predict_dataset.remove_columns(output_column)
+
+        if data_args.test_start_ind is not None:
+            sind =  data_args.test_start_ind
+            eind = -1 if data_args.test_end_ind is None else data_args.test_end_ind
+            logger.info(f'Using only a subset of the test dataset [{sind}, {eind}]')
+            untokenized_predict_dataset = type(untokenized_predict_dataset).from_dict(untokenized_predict_dataset[sind:eind])
+
+        with training_args.main_process_first(
+            local=not data_args.shared_storage, desc="prediction dataset map pre-processing"
+        ):
+            predict_dataset = untokenized_predict_dataset.map(
+                preprocess_function,
+                fn_kwargs=preprocess_function_kwargs,
+                batched=True,
+                num_proc=data_args.preprocessing_num_workers,
+                remove_columns=untokenized_predict_dataset.column_names,
+                load_from_cache_file=not data_args.overwrite_cache,
+                desc="Running tokenizer on prediction dataset",
+            )
+
+    if data_args.preprocess_only:
+        logger.info(f"With --preprocess_only, exiting after preprocess_on the data")
+        exit()
+
+    # Data collator
+    label_pad_token_id = -100 if data_args.ignore_pad_token_for_loss else tokenizer.pad_token_id
+    pad_to = 8 if training_args.fp16 and training_args.fp16_padding else None
+
+
+    data_collator = DataCollatorForSeq2Seq(
+        tokenizer,
+        model=model,
+        label_pad_token_id=label_pad_token_id,
+        pad_to_multiple_of=pad_to,
+    )
+
+    # Metric
+    compute_metrics = load_metric(data_args.metric_names, **locals())
+    compute_metrics = load_extra_metrics(data_args.extra_metrics, compute_metrics)
+
+    # Initialize our Trainer
+    trainer = CustomTrainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset if training_args.do_train else None,
+        eval_dataset=eval_dataset if training_args.do_eval else None,
+        untokenized_eval_dataset=untokenized_eval_dataset if training_args.do_eval else None,
+        tokenizer=tokenizer,
+        data_collator=data_collator,
+        compute_metrics=compute_metrics if training_args.predict_with_generate else None,
+        output_dir=training_args.output_dir,
+        data_args=data_args,
+        callbacks=[EarlyStoppingCallback(early_stopping_patience=data_args.patience)] if data_args.patience is not None else None,
+    )
+
+    # setup_cometml_trainer_callback(trainer)
+
+    # Training
+    if training_args.do_train:
+        checkpoint = None
+        if training_args.resume_from_checkpoint is not None:
+            checkpoint = training_args.resume_from_checkpoint
+        elif last_checkpoint is not None:
+            checkpoint = last_checkpoint  # look for checkpoints in the outdir
+
+        train_result = trainer.train(resume_from_checkpoint=checkpoint)
+        logger.info('Done training')
+        trainer.save_model()  # Saves the tokenizer too for easy upload
+
+        metrics = train_result.metrics
+        max_train_samples = (
+            data_args.max_train_samples if data_args.max_train_samples is not None else len(train_dataset)
+        )
+        metrics["train_samples"] = min(max_train_samples, len(train_dataset))
+
+        trainer.log_metrics("train", metrics)
+        trainer.save_metrics("train", metrics)
+        trainer.save_state()
+
+    # Evaluation
+    results = {}
+    if training_args.do_eval:
+        logger.info("*** Evaluate ***")
+
+        if training_args.eval_fraction < 1:
+            logger.info('setting the eval set back to the full one')
+            trainer.eval_dataset = eval_dataset_orig
+            trainer._untokenized_eval_dataset = untokenized_eval_dataset_orig
+
+        metrics = trainer.evaluate(metric_key_prefix="eval", use_cache=True, length_penalty=data_args.length_penalty)
+        logger.info('Done evaluating')
+
+        max_eval_samples = data_args.max_eval_samples if data_args.max_eval_samples is not None else len(eval_dataset)
+        metrics["eval_samples"] = min(max_eval_samples, len(eval_dataset))
+
+        trainer.log_metrics("eval", metrics)
+        trainer.save_metrics("eval", metrics)
+
+    if training_args.do_predict:
+        logger.info("*** Predict ***")
+        trainer.args.predict_with_generate = True # during prediction, we don't have labels
+
+        # load last (and best) model, or the one specified if any
+        logger.info("*** Loading model weights before the prediction ***")
+        last_checkpoint = model_args.model_name_or_path if os.path.isdir(model_args.model_name_or_path) else _detect_last_checkpoint(training_args)
+        if last_checkpoint is not None and os.path.isdir(last_checkpoint):
+            logger.info(f'Loading weights from {last_checkpoint} for the prediction')
+            state_dict = torch.load(os.path.join(last_checkpoint, WEIGHTS_NAME), map_location="cpu")
+            # If the model is on the GPU, it still works!
+            # trainer._load_state_dict_in_model(state_dict)
+            # release memory
+            del state_dict
+            logger.info("*** Done loading weights ***")
+        elif training_args.do_train:
+            raise ValueError('Could not find a model to load for prediction')
+        else:
+            logger.info(f'Using {model_args.model_name_or_path} as the model for the prediction')
+
+        predict_results = trainer.predict(predict_dataset, metric_key_prefix="predict", use_cache=True)
+        logger.info('Done predicting')
+
+        metrics = predict_results.metrics
+        max_predict_samples = (
+            data_args.max_predict_samples if data_args.max_predict_samples is not None else len(predict_dataset)
+        )
+        metrics["predict_samples"] = min(max_predict_samples, len(predict_dataset))
+
+        trainer.log_metrics("predict", metrics)
+        trainer.save_metrics("predict", metrics)
+
+        if trainer.is_world_process_zero():
+            if training_args.predict_with_generate:
+                id_to_prediction = {}
+                for i, instance in enumerate(untokenized_predict_dataset):
+                    id_to_prediction[instance["id"]] = predict_results.predictions[i]
+                predictions = decode(id_to_prediction, tokenizer, data_args)
+                output_name = "generated_predictions.json"
+                if data_args.test_start_ind is not None:
+                    output_name = f"generated_predictions_{data_args.test_start_ind}_{data_args.test_end_ind}.json"
+                output_prediction_file = os.path.join(training_args.output_dir, output_name)
+                with open(output_prediction_file, "w") as writer:
+                    json.dump(predictions, writer, indent=4)
+
+    if training_args.push_to_hub:
+        kwargs = {"finetuned_from": model_args.model_name_or_path, "tasks": "summarization"}
+        if data_args.dataset_name is not None:
+            kwargs["dataset_tags"] = data_args.dataset_name
+            if data_args.dataset_config_name is not None:
+                kwargs["dataset_args"] = data_args.dataset_config_name
+                kwargs["dataset"] = f"{data_args.dataset_name} {data_args.dataset_config_name}"
+            else:
+                kwargs["dataset"] = data_args.dataset_name
+
+        trainer.push_to_hub(**kwargs)
+
+    return results
+
+def _detect_last_checkpoint(training_args):
+    last_checkpoint = None
+    if os.path.isdir(training_args.output_dir) and training_args.do_train:
+        if not training_args.overwrite_output_dir:
+            last_checkpoint = get_last_checkpoint(training_args.output_dir)
+
+            if last_checkpoint is not None and training_args.resume_from_checkpoint is None:
+                logger.info(
+                    f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
+                    "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
+                )
+    return last_checkpoint
+
+def process_eval_set(data_args, preprocess_function_kwargs, training_args, untokenized_eval_dataset):
+    with training_args.main_process_first(
+            local=not data_args.shared_storage, desc="validation dataset map pre-processing"
+    ):
+        eval_dataset = untokenized_eval_dataset.map(
+            preprocess_function,
+            fn_kwargs=preprocess_function_kwargs,
+            batched=True,
+            num_proc=data_args.preprocessing_num_workers,
+            remove_columns=untokenized_eval_dataset.column_names,
+            load_from_cache_file=not data_args.overwrite_cache,
+            desc="Running tokenizer on validation dataset",
+        )
+    return eval_dataset
+
+
+def _get_dataset(data_args, model_args, training_args):
+    # Get the datasets: you can either provide your own CSV/JSON training and evaluation files (see below)
+    # or just provide the name of one of the public datasets available on the hub at https://huggingface.co/datasets/
+    # (the dataset will be downloaded automatically from the datasets Hub).
+    #
+    # For CSV/JSON files this script will use the first column for the full texts and the second column for the
+    # summaries (unless you specify column names for this with the `input_column` and `output_column` arguments).
+    #
+    # In distributed training, the load_dataset function guarantee that only one local process can concurrently
+    # download the dataset.
+    data_files = None
+    if data_args.train_file is not None or data_args.validation_file is not None or data_args.test_file is not None:
+        data_files = {}
+        if data_args.train_file is not None:
+            data_files["train"] = data_args.train_file
+        if data_args.validation_file is not None:
+            data_files["validation"] = data_args.validation_file
+        if data_args.test_file is not None:
+            data_files["test"] = data_args.test_file
+    # Downloading and loading a dataset from the hub/local script.
+    seq2seq_dataset = load_dataset(
+        data_args.dataset_name,
+        data_args.dataset_config_name,
+        verification_mode='no_checks',
+        cache_dir=model_args.cache_dir,
+        data_dir=data_args.data_dir,
+        data_files=data_files,
+        download_mode=data_args.download_mode,
+        use_auth_token=training_args.use_auth_token
+    )
+    if training_args.do_train:
+        training_args.apply_overrides(len(seq2seq_dataset['train']))
+    if data_args.evaluate_on_training_data:
+        seq2seq_dataset["validation"] = seq2seq_dataset["train"]
+
+    # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at
+    # https://huggingface.co/docs/datasets/loading_datasets.html.
+
+    return seq2seq_dataset
+
+
+def set_up_logging(training_args):
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+        datefmt="%m/%d/%Y %H:%M:%S",
+        handlers=[logging.StreamHandler(sys.stdout)],
+    )
+    log_level = training_args.get_process_log_level()
+    logger.setLevel(log_level)
+    datasets.utils.logging.set_verbosity(log_level)
+    transformers.utils.logging.set_verbosity(log_level)
+    transformers.utils.logging.enable_default_handler()
+    transformers.utils.logging.enable_explicit_format()
+    # Log on each process the small summary:
+    logger.warning(
+        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
+        + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
+    )
+    logger.info(f"Training/evaluation parameters {training_args}")
+
+def extract_oracle_sent_batch(examples, max_length, tokenizer, rouge_scorer):
+    items = examples.data.items()
+    keys = [item[0] for item in items]
+    values = [item[1] for item in items]
+    extracted = {k: [] for k in keys}
+    input_str = 'input'
+
+    for ex in zip(*values):
+        ex = dict(zip(keys, ex))
+        ex_input = ex[input_str]
+        extracted_input = extract_oracle_sentences(ex_input, ex['output'], max_length, tokenizer, rouge_scorer)
+        extracted[input_str].append(extracted_input)
+        for k in set(keys) - {input_str}:
+            extracted[k].append(ex[k])
+    return extracted
+
+def extract_oracle_sentences(input_sequence, output, max_length, tokenizer, rouge_scorer, criterion='rouge/geometric_mean'):
+    sentences = nltk.sent_tokenize(input_sequence)
+    selected_mask = [False for _ in sentences]
+
+    max_rouge = 0.0
+    joined_selection = ''
+    counter = 0
+    while len(tokenizer(joined_selection)) < max_length and counter < 100:
+        cur_max_rouge = max_rouge
+        max_index = -1
+        
+        cur_candidate_indices = []
+        cur_candidates = []
+        for i in range(len(sentences)):
+            if selected_mask[i]:
+                # We already selected this sentence
+                continue
+            candidate_mask = list(selected_mask)
+            candidate_mask[i] = True
+            candidate_prediction = ' '.join(sent for sent, mask in zip(sentences, candidate_mask) if mask)
+            cur_candidates.append(candidate_prediction)
+            cur_candidate_indices.append(i)
+        
+        rouge = rouge_scorer.compute(predictions=cur_candidates, references=[[output]] * len(cur_candidates), use_aggregator=False)
+        aggregated_rouge_types = [s1.fmeasure * s2.fmeasure * sL.fmeasure for s1, s2, sL in zip(rouge['rouge1'], rouge['rouge2'], rouge['rougeLsum'])]
+        max_index = np.argmax(aggregated_rouge_types)
+        cur_max_rouge = aggregated_rouge_types[max_index]
+        
+        if max_rouge >= cur_max_rouge:
+            # No sentence improves the score
+            break
+        
+        selected_mask[cur_candidate_indices[max_index]] = True
+        max_rouge = cur_max_rouge
+        joined_selection = ' '.join(sent for sent, mask in zip(sentences, selected_mask) if mask)
+        counter += 1
+    
+    return joined_selection        
+    
+
+def chunk_dataset_function(examples, chunk_size):
+    input_ids_str = 'input_ids'
+    attention_mask_str = 'attention_mask'
+    items = examples.data.items()
+    keys = [item[0] for item in items]
+    values = [item[1] for item in items]
+    chunked = {k: [] for k in keys}
+    for ex in zip(*values):
+        ex = dict(zip(keys, ex))
+        for i in range(0, len(ex[input_ids_str]), chunk_size):
+            chunked_input_ids_st = ex[input_ids_str][i:i + chunk_size]
+            chunked_attention_mask = ex[attention_mask_str][i:i + chunk_size]
+
+            if sum(chunked_attention_mask) < 10:
+                continue
+            chunked[input_ids_str].append(chunked_input_ids_st)
+            chunked[attention_mask_str].append(chunked_attention_mask)
+            for k in set(keys) - {input_ids_str, attention_mask_str}:
+                chunked[k].append(ex[k])
+    return chunked
+
+    
+
+def preprocess_function(
+    examples,
+    tokenizer,
+    prefix,
+    input_column,
+    input_prefix_column,
+    output_column,
+    max_source_length,
+    max_prefix_length,
+    max_target_length,
+    prefix_sep,
+    padding,
+    ignore_pad_token_for_loss,
+    assign_zero_to_too_long_val_examples,
+    trim_very_long_strings,
+    pad_prefix
+):
+    if not isinstance(examples[input_column][0], str):
+        model_inputs = _preprocess_tokenized_inputs()
+    else:
+        model_inputs = _preprocess_raw_inputs(assign_zero_to_too_long_val_examples, examples, input_column, input_prefix_column,
+                                              max_source_length, padding, prefix, tokenizer, trim_very_long_strings, max_prefix_length,
+                                              prefix_sep, pad_prefix)
+
+    _preprocess_targets(examples, ignore_pad_token_for_loss, max_target_length, model_inputs, output_column, padding, tokenizer)
+    model_inputs["length"] = [len(x) for x in model_inputs["input_ids"]]
+    return model_inputs
+
+
+def _preprocess_raw_inputs(assign_zero_to_too_long_val_examples, examples, input_column, input_prefix_column,
+                           max_source_length, padding, prefix, tokenizer, trim_very_long_strings, max_prefix_length,
+                           prefix_sep, pad_prefix):
+    inputs = examples[input_column]
+
+    # the given prefix is what used in models like T5 (e.g. "summarize: ")
+    # if prefix exists, it is added to the input_prefixes
+    if input_prefix_column in examples.keys():
+        input_prefixes = [inp + prefix_sep for inp in examples[input_prefix_column]]
+        if prefix != "":
+            input_prefixes = [prefix + inp for inp in input_prefixes]
+    elif prefix != "":
+        inputs = [prefix + inp for inp in inputs]
+
+    # tokenize the input prefix if it exists
+    model_prefix_inputs = None
+    if input_prefix_column in examples.keys():
+        if trim_very_long_strings:
+            input_prefixes = [inp[: max_prefix_length * 7] for inp in input_prefixes]
+        if pad_prefix:
+            model_prefix_inputs = tokenizer(input_prefixes, max_length=max_prefix_length, padding='max_length', truncation=True)
+        else:
+            # for led, we do not pad the prefix
+            model_prefix_inputs = tokenizer(input_prefixes, max_length=max_source_length, padding='do_not_pad', truncation=True)
+
+    if trim_very_long_strings:
+        inputs = [inp[: max_source_length * 7] for inp in inputs]
+    model_inputs = tokenizer(inputs, max_length=max_source_length, padding=padding, truncation=True)
+
+    if max_source_length is not None and assign_zero_to_too_long_val_examples:
+        model_inputs_untrimmed = tokenizer(inputs)
+        model_inputs["not_valid_for_eval"] = [
+            len(token_ids) > max_source_length for token_ids in model_inputs_untrimmed["input_ids"]
+        ]
+    else:
+        model_inputs["not_valid_for_eval"] = [False] * len(model_inputs["input_ids"])
+
+    # now, combine the concat prefix to the input, trimming it to max_source_length if given
+    if model_prefix_inputs is not None:
+        max_source_length = max_source_length or -1
+        model_inputs['input_ids'] = [(inp1+inp2)[:max_source_length] for inp1, inp2
+                                     in zip(model_prefix_inputs['input_ids'], model_inputs['input_ids'])]
+        model_inputs['attention_mask'] = [(inp1+inp2)[:max_source_length] for inp1, inp2
+                                          in zip(model_prefix_inputs['attention_mask'], model_inputs['attention_mask'])]
+        # add prefix_length
+        if pad_prefix:
+            # no need to go over them as they will all be of the same length
+            model_inputs['prefix_length'] = [max_prefix_length] * len(model_inputs['input_ids'])
+        else:
+            model_inputs['prefix_length'] = [len(inp) for inp in model_prefix_inputs['input_ids']]
+
+    return model_inputs
+
+def _preprocess_targets(examples, ignore_pad_token_for_loss, max_target_length, model_inputs, output_column, padding, tokenizer):
+    targets = examples[output_column] if output_column in examples else None
+    if targets is not None:
+        if not isinstance(targets[0], str):
+            if max_target_length is not None:
+                targets = [target[:max_target_length] for target in targets]
+            model_inputs["labels"] = targets
+        else:
+            # Setup the tokenizer for targets
+            with tokenizer.as_target_tokenizer():
+                labels = tokenizer(targets, max_length=max_target_length, padding=padding, truncation=True)
+
+            # If we are padding here, replace all tokenizer.pad_token_id in the labels by -100 when we want to ignore
+            # padding in the loss.
+            if padding == "max_length" and ignore_pad_token_for_loss:
+                labels["input_ids"] = [
+                    [(l if l != tokenizer.pad_token_id else -100) for l in label] for label in labels["input_ids"]
+                ]
+
+            model_inputs["labels"] = labels["input_ids"]
+
+def load_extra_metrics(metric_names, loaded_metrics=None):
+    if loaded_metrics is None:
+        loaded_metrics = MetricCollection([])
+    if metric_names is not None:
+        for metric_name in metric_names:
+            if len(metric_name) > 0:
+                loaded_metrics._metrics.append(HFMetricWrapper(metric_name))
+    return loaded_metrics
+
+def _mp_fn(index):
+    # For xla_spawn (TPUs)
+    main()
+
+
+if __name__ == "__main__":
+    main()

unlimiformer/run_generation.py

@@ -0,0 +1,577 @@
+#!/usr/bin/env python
+# coding=utf-8
+# Copyright 2018 Google AI, Google Brain and Carnegie Mellon University Authors and the HuggingFace Inc. team.
+# Copyright (c) 2018, NVIDIA CORPORATION.  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.
+""" Conditional text generation with the auto-regressive models of the library (GPT/GPT-2/CTRL/Transformer-XL/XLNet)
+"""
+
+
+import argparse
+import inspect
+import logging
+
+from dataclasses import dataclass, field
+from typing import  Tuple, List, Optional, Union
+
+import numpy as np
+import torch
+import os
+
+normal_repr = torch.Tensor.__repr__
+torch.Tensor.__repr__ = lambda self: f"{self.shape}_{normal_repr(self)}"
+
+from transformers import (
+    AutoTokenizer,
+    BloomForCausalLM,
+    BloomTokenizerFast,
+    CTRLLMHeadModel,
+    CTRLTokenizer,
+    GenerationMixin,
+    GPT2LMHeadModel,
+    GPT2Tokenizer,
+    GPTJForCausalLM,
+    HfArgumentParser,
+    LlamaForCausalLM,
+    LlamaTokenizer,
+    OpenAIGPTLMHeadModel,
+    OpenAIGPTTokenizer,
+    OPTForCausalLM,
+    TransfoXLLMHeadModel,
+    TransfoXLTokenizer,
+    XLMTokenizer,
+    XLMWithLMHeadModel,
+    XLNetLMHeadModel,
+    XLNetTokenizer,
+    TextStreamer,
+)
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from unlimiformer import Unlimiformer
+from random_training_unlimiformer import RandomTrainingUnlimiformer
+
+@dataclass
+class UnlimiformerArguments:
+    """
+    Arguments pertaining to what data we are going to input our model for training and eval.
+    """
+    test_unlimiformer: Optional[bool] = field(
+        default=False,
+        metadata={
+            "help": "whether to use KNN."
+        },
+    )
+    unlimiformer_verbose: Optional[bool] = field(
+        default=False,
+        metadata={
+            "help": "whether to print KNN intermediate predictions (mostly for debugging)."
+        },
+    )
+    layer_begin: Optional[int] = field(
+        default=0,
+        metadata={"help": "The layer to begin applying KNN to. KNN will be applied to layers[knn_layer_begin:layer_end]. "
+                          "By default, it will be applied to all layers: [0:None]]"}, 
+    )
+    layer_end: Optional[int] = field(
+        default=None,
+        metadata={"help": "The layer to end applying KNN to. KNN will be applied to layers[knn_layer_begin:layer_end]. "
+                          "By default, it will be applied to all layers: [0:None]]"}, 
+    )
+    unlimiformer_chunk_overlap: Optional[float] = field(
+        default=0.5,
+        metadata={"help": "The fraction of overlap between input chunks"},
+    )
+    unlimiformer_chunk_size: Optional[int] = field(
+        default=None,
+        metadata={"help": "The size of each input chunk"},
+    )
+    unlimiformer_head_num: Optional[int] = field(
+        default=None,
+        metadata={"help": "The head to apply KNN to (if None, apply to all heads)"},
+    )
+    unlimiformer_exclude: Optional[bool] = field(
+        default=False,
+        metadata={
+            "help": "If True, prioritize the inputs that are **not** in the standard attention window."
+        },
+    )
+    random_unlimiformer_training: Optional[bool] = field(
+        default=False,
+    )
+    unlimiformer_training: Optional[bool] = field(
+        default=False,
+    )
+    index_devices: Optional[List[int]] = field(
+        default_factory=lambda: (0,),
+    )
+    datastore_device: Optional[int] = field(
+        default=0,
+    )
+    use_datastore: Optional[bool] = field(default=True)
+    flat_index: Optional[bool] = field(default=True)
+    test_datastore: Optional[bool] = field(default=False)
+    reconstruct_embeddings: Optional[bool] = field(default=False)
+    gpu_datastore: Optional[bool] = field(default=True)
+    gpu_index: Optional[bool] = field(default=True)
+
+
+logging.basicConfig(
+    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+    datefmt="%m/%d/%Y %H:%M:%S",
+    level=logging.INFO,
+)
+logger = logging.getLogger(__name__)
+
+MAX_LENGTH = int(10000)  # Hardcoded max length to avoid infinite loop
+
+MODEL_CLASSES = {
+    "gpt2": (GPT2LMHeadModel, GPT2Tokenizer),
+    "ctrl": (CTRLLMHeadModel, CTRLTokenizer),
+    "openai-gpt": (OpenAIGPTLMHeadModel, OpenAIGPTTokenizer),
+    "xlnet": (XLNetLMHeadModel, XLNetTokenizer),
+    "transfo-xl": (TransfoXLLMHeadModel, TransfoXLTokenizer),
+    "xlm": (XLMWithLMHeadModel, XLMTokenizer),
+    "gptj": (GPTJForCausalLM, AutoTokenizer),
+    "bloom": (BloomForCausalLM, BloomTokenizerFast),
+    "llama": (LlamaForCausalLM, LlamaTokenizer),
+    "opt": (OPTForCausalLM, GPT2Tokenizer),
+}
+
+# Padding text to help Transformer-XL and XLNet with short prompts as proposed by Aman Rusia
+# in https://github.com/rusiaaman/XLNet-gen#methodology
+# and https://medium.com/@amanrusia/xlnet-speaks-comparison-to-gpt-2-ea1a4e9ba39e
+PREFIX = """In 1991, the remains of Russian Tsar Nicholas II and his family
+(except for Alexei and Maria) are discovered.
+The voice of Nicholas's young son, Tsarevich Alexei Nikolaevich, narrates the
+remainder of the story. 1883 Western Siberia,
+a young Grigori Rasputin is asked by his father and a group of men to perform magic.
+Rasputin has a vision and denounces one of the men as a horse thief. Although his
+father initially slaps him for making such an accusation, Rasputin watches as the
+man is chased outside and beaten. Twenty years later, Rasputin sees a vision of
+the Virgin Mary, prompting him to become a priest. Rasputin quickly becomes famous,
+with people, even a bishop, begging for his blessing. <eod> </s> <eos>"""
+
+
+def set_seed(args):
+    np.random.seed(args.seed)
+    torch.manual_seed(args.seed)
+    if args.n_gpu > 0:
+        torch.cuda.manual_seed_all(args.seed)
+
+
+#
+# Functions to prepare models' input
+#
+
+
+def prepare_ctrl_input(args, _, tokenizer, prompt_text):
+    if args.temperature > 0.7:
+        logger.info("CTRL typically works better with lower temperatures (and lower top_k).")
+
+    encoded_prompt = tokenizer.encode(prompt_text, add_special_tokens=False)
+    if not any(encoded_prompt[0] == x for x in tokenizer.control_codes.values()):
+        logger.info("WARNING! You are not starting your generation from a control code so you won't get good results")
+    return prompt_text
+
+
+def prepare_xlm_input(args, model, tokenizer, prompt_text):
+    # kwargs = {"language": None, "mask_token_id": None}
+
+    # Set the language
+    use_lang_emb = hasattr(model.config, "use_lang_emb") and model.config.use_lang_emb
+    if hasattr(model.config, "lang2id") and use_lang_emb:
+        available_languages = model.config.lang2id.keys()
+        if args.xlm_language in available_languages:
+            language = args.xlm_language
+        else:
+            language = None
+            while language not in available_languages:
+                language = input("Using XLM. Select language in " + str(list(available_languages)) + " >>> ")
+
+        model.config.lang_id = model.config.lang2id[language]
+        # kwargs["language"] = tokenizer.lang2id[language]
+
+    # TODO fix mask_token_id setup when configurations will be synchronized between models and tokenizers
+    # XLM masked-language modeling (MLM) models need masked token
+    # is_xlm_mlm = "mlm" in args.model_name_or_path
+    # if is_xlm_mlm:
+    #     kwargs["mask_token_id"] = tokenizer.mask_token_id
+
+    return prompt_text
+
+
+def prepare_xlnet_input(args, _, tokenizer, prompt_text):
+    prefix = args.prefix if args.prefix else args.padding_text if args.padding_text else PREFIX
+    prompt_text = prefix + prompt_text
+    return prompt_text
+
+
+def prepare_transfoxl_input(args, _, tokenizer, prompt_text):
+    prefix = args.prefix if args.prefix else args.padding_text if args.padding_text else PREFIX
+    prompt_text = prefix + prompt_text
+    return prompt_text
+
+
+PREPROCESSING_FUNCTIONS = {
+    "ctrl": prepare_ctrl_input,
+    "xlm": prepare_xlm_input,
+    "xlnet": prepare_xlnet_input,
+    "transfo-xl": prepare_transfoxl_input,
+}
+
+
+def adjust_length_to_model(length, max_sequence_length):
+    if length < 0 and max_sequence_length > 0:
+        length = max_sequence_length
+    elif 0 < max_sequence_length < length:
+        length = max_sequence_length  # No generation bigger than model size
+    elif length < 0:
+        length = MAX_LENGTH  # avoid infinite loop
+    return length
+
+
+def sparse_model_config(model_config):
+    embedding_size = None
+    if hasattr(model_config, "hidden_size"):
+        embedding_size = model_config.hidden_size
+    elif hasattr(model_config, "n_embed"):
+        embedding_size = model_config.n_embed
+    elif hasattr(model_config, "n_embd"):
+        embedding_size = model_config.n_embd
+
+    num_head = None
+    if hasattr(model_config, "num_attention_heads"):
+        num_head = model_config.num_attention_heads
+    elif hasattr(model_config, "n_head"):
+        num_head = model_config.n_head
+
+    if embedding_size is None or num_head is None or num_head == 0:
+        raise ValueError("Check the model config")
+
+    num_embedding_size_per_head = int(embedding_size / num_head)
+    if hasattr(model_config, "n_layer"):
+        num_layer = model_config.n_layer
+    elif hasattr(model_config, "num_hidden_layers"):
+        num_layer = model_config.num_hidden_layers
+    else:
+        raise ValueError("Number of hidden layers couldn't be determined from the model config")
+
+    return num_layer, num_head, num_embedding_size_per_head
+
+
+def generate_past_key_values(model, batch_size, seq_len):
+    num_block_layers, num_attention_heads, num_embedding_size_per_head = sparse_model_config(model.config)
+    if model.config.model_type == "bloom":
+        past_key_values = tuple(
+            (
+                torch.empty(int(num_attention_heads * batch_size), num_embedding_size_per_head, seq_len)
+                .to(model.dtype)
+                .to(model.device),
+                torch.empty(int(num_attention_heads * batch_size), seq_len, num_embedding_size_per_head)
+                .to(model.dtype)
+                .to(model.device),
+            )
+            for _ in range(num_block_layers)
+        )
+    else:
+        past_key_values = tuple(
+            (
+                torch.empty(batch_size, num_attention_heads, seq_len, num_embedding_size_per_head)
+                .to(model.dtype)
+                .to(model.device),
+                torch.empty(batch_size, num_attention_heads, seq_len, num_embedding_size_per_head)
+                .to(model.dtype)
+                .to(model.device),
+            )
+            for _ in range(num_block_layers)
+        )
+    return past_key_values
+
+
+def prepare_jit_inputs(inputs, model, tokenizer):
+    batch_size = len(inputs)
+    dummy_input = tokenizer.batch_encode_plus(inputs, return_tensors="pt")
+    dummy_input = dummy_input.to(model.device)
+    if model.config.use_cache:
+        dummy_input["past_key_values"] = generate_past_key_values(model, batch_size, 1)
+    dummy_input["attention_mask"] = torch.cat(
+        [
+            torch.zeros(dummy_input["attention_mask"].shape[0], 1)
+            .to(dummy_input["attention_mask"].dtype)
+            .to(model.device),
+            dummy_input["attention_mask"],
+        ],
+        -1,
+    )
+    return dummy_input
+
+
+class _ModelFallbackWrapper(GenerationMixin):
+    __slots__ = ("_optimized", "_default")
+
+    def __init__(self, optimized, default):
+        self._optimized = optimized
+        self._default = default
+
+    def __call__(self, *args, **kwargs):
+        if kwargs["past_key_values"] is None and self._default.config.use_cache:
+            kwargs["past_key_values"] = generate_past_key_values(self._default, kwargs["input_ids"].shape[0], 0)
+        kwargs.pop("position_ids", None)
+        for k in list(kwargs.keys()):
+            if kwargs[k] is None or isinstance(kwargs[k], bool):
+                kwargs.pop(k)
+        outputs = self._optimized(**kwargs)
+        lm_logits = outputs[0]
+        past_key_values = outputs[1]
+        fixed_output = CausalLMOutputWithPast(
+            loss=None,
+            logits=lm_logits,
+            past_key_values=past_key_values,
+            hidden_states=None,
+            attentions=None,
+        )
+        return fixed_output
+
+    def __getattr__(self, item):
+        return getattr(self._default, item)
+
+    def prepare_inputs_for_generation(
+        self, input_ids, past_key_values=None, inputs_embeds=None, use_cache=None, **kwargs
+    ):
+        return self._default.prepare_inputs_for_generation(
+            input_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, **kwargs
+        )
+
+    def _reorder_cache(
+        self, past_key_values: Tuple[Tuple[torch.Tensor]], beam_idx: torch.Tensor
+    ) -> Tuple[Tuple[torch.Tensor]]:
+        """
+        This function is used to re-order the `past_key_values` cache if [`~PretrainedModel.beam_search`] or
+        [`~PretrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct
+        beam_idx at every generation step.
+        """
+        return self._default._reorder_cache(past_key_values, beam_idx)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model_type",
+        default=None,
+        type=str,
+        required=True,
+        help="Model type selected in the list: " + ", ".join(MODEL_CLASSES.keys()),
+    )
+    parser.add_argument(
+        "--model_name_or_path",
+        default=None,
+        type=str,
+        required=True,
+        help="Path to pre-trained model or shortcut name selected in the list: " + ", ".join(MODEL_CLASSES.keys()),
+    )
+
+    parser.add_argument("--prompt", type=str, default="")
+    parser.add_argument("--length", type=int, default=100)
+    parser.add_argument("--num_hidden_layers", type=int, default=None)
+    parser.add_argument("--stop_token", type=str, default=None, help="Token at which text generation is stopped")
+
+    parser.add_argument(
+        "--temperature",
+        type=float,
+        default=1.0,
+        help="temperature of 1.0 has no effect, lower tend toward greedy sampling",
+    )
+    parser.add_argument(
+        "--repetition_penalty", type=float, default=1.0, help="primarily useful for CTRL model; in that case, use 1.2"
+    )
+    parser.add_argument("--k", type=int, default=0)
+    parser.add_argument("--p", type=float, default=0.9)
+
+    parser.add_argument("--prefix", type=str, default="", help="Text added prior to input.")
+    parser.add_argument("--suffix", type=str, default="", help="Text added after the input.")
+    parser.add_argument("--padding_text", type=str, default="", help="Deprecated, the use of `--prefix` is preferred.")
+    parser.add_argument("--xlm_language", type=str, default="", help="Optional language when used with the XLM model.")
+
+    parser.add_argument("--seed", type=int, default=42, help="random seed for initialization")
+    parser.add_argument("--no_cuda", action="store_true", help="Avoid using CUDA when available")
+    parser.add_argument("--stream_output", action="store_true")
+    parser.add_argument("--num_return_sequences", type=int, default=1, help="The number of samples to generate.")
+    parser.add_argument(
+        "--fp16",
+        action="store_true",
+        help="Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit",
+    )
+    parser.add_argument("--jit", action="store_true", help="Whether or not to use jit trace to accelerate inference")
+
+    # args = parser.parse_args()
+    args, unknown_args = parser.parse_known_args()
+
+    hf_parser = HfArgumentParser(UnlimiformerArguments)
+    unlimiformer_args, unknown_unlimiformer_args = hf_parser.parse_known_args()
+    
+    if len(set(unknown_args) & set(unknown_unlimiformer_args)) > 0:
+        raise ValueError(f"Unknown arguments detected: {set(unknown_args) & set(unknown_unlimiformer_args)}")
+
+    args.device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu")
+    args.n_gpu = 0 if args.no_cuda else torch.cuda.device_count()
+
+    logger.warning(f"device: {args.device}, n_gpu: {args.n_gpu}, 16-bits training: {args.fp16}")
+
+    set_seed(args)
+
+    # Initialize the model and tokenizer
+    try:
+        args.model_type = args.model_type.lower()
+        model_class, tokenizer_class = MODEL_CLASSES[args.model_type]
+    except KeyError:
+        raise KeyError("the model {} you specified is not supported. You are welcome to add it and open a PR :)")
+
+    tokenizer = tokenizer_class.from_pretrained(args.model_name_or_path)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    model_kwargs = {}
+    if args.num_hidden_layers is not None:
+        model_kwargs["num_hidden_layers"] = args.num_hidden_layers
+    model = model_class.from_pretrained(args.model_name_or_path, **model_kwargs)
+
+    if args.fp16:
+        model.half()
+    model.to(args.device)
+
+    max_seq_length = getattr(model.config, "max_position_embeddings", 0)
+    args.length = adjust_length_to_model(args.length, max_sequence_length=max_seq_length)
+    logger.info(args)
+
+    if unlimiformer_args.test_unlimiformer:
+        unlimiformer_kwargs = {
+            'layer_begin': unlimiformer_args.layer_begin, 
+            'layer_end': unlimiformer_args.layer_end,
+            'unlimiformer_head_num': unlimiformer_args.unlimiformer_head_num, 
+            'exclude_attention': unlimiformer_args.unlimiformer_exclude, 
+            'chunk_overlap': unlimiformer_args.unlimiformer_chunk_overlap,
+            'model_encoder_max_len': unlimiformer_args.unlimiformer_chunk_size,
+            'verbose': unlimiformer_args.unlimiformer_verbose, 'tokenizer': tokenizer,
+            'unlimiformer_training': unlimiformer_args.unlimiformer_training,
+            'use_datastore': unlimiformer_args.use_datastore,
+            'flat_index': unlimiformer_args.flat_index,
+            'test_datastore': unlimiformer_args.test_datastore,
+            'reconstruct_embeddings': unlimiformer_args.reconstruct_embeddings,
+            'gpu_datastore': unlimiformer_args.gpu_datastore,
+            'gpu_index': unlimiformer_args.gpu_index,
+            'index_devices': unlimiformer_args.index_devices,
+            'datastore_device': unlimiformer_args.datastore_device,
+        }
+        if unlimiformer_args.random_unlimiformer_training:
+            model = RandomTrainingUnlimiformer.convert_model(model, **unlimiformer_kwargs)
+        else:
+            model = Unlimiformer.convert_model(model, **unlimiformer_kwargs)
+
+    prompt_text = args.prompt if args.prompt else input("Model prompt >>> ")
+    # Check if prompt_text is a valid file name:
+    if os.path.exists(prompt_text):
+        with open(prompt_text, "r") as f:
+            prompt_text = f.read()
+
+    # Different models need different input formatting and/or extra arguments
+    requires_preprocessing = args.model_type in PREPROCESSING_FUNCTIONS.keys()
+    if requires_preprocessing:
+        prepare_input = PREPROCESSING_FUNCTIONS.get(args.model_type)
+        preprocessed_prompt_text = prepare_input(args, model, tokenizer, prompt_text)
+
+        if model.__class__.__name__ in ["TransfoXLLMHeadModel"]:
+            tokenizer_kwargs = {"add_space_before_punct_symbol": True}
+        else:
+            tokenizer_kwargs = {}
+
+        encoded_prompt = tokenizer.encode(
+            preprocessed_prompt_text, add_special_tokens=False, return_tensors="pt", **tokenizer_kwargs
+        )
+    else:
+        # prefix = args.prefix if args.prefix else args.padding_text
+        prompt_text = f'{args.prefix}{prompt_text}{args.suffix}'
+        encoded_prompt = tokenizer.encode(prompt_text, add_special_tokens=False, return_tensors="pt")
+    
+    if not unlimiformer_args.test_unlimiformer:
+        encoded_prompt = encoded_prompt[:, -2048:]
+        encoded_prompt = encoded_prompt.to(args.device)
+
+    if encoded_prompt.size()[-1] == 0:
+        input_ids = None
+    else:
+        input_ids = encoded_prompt
+
+    if args.jit:
+        jit_input_texts = ["enable jit"]
+        jit_inputs = prepare_jit_inputs(jit_input_texts, model, tokenizer)
+        torch._C._jit_set_texpr_fuser_enabled(False)
+        model.config.return_dict = False
+        if hasattr(model, "forward"):
+            sig = inspect.signature(model.forward)
+        else:
+            sig = inspect.signature(model.__call__)
+        jit_inputs = tuple(jit_inputs[key] for key in sig.parameters if jit_inputs.get(key, None) is not None)
+        traced_model = torch.jit.trace(model, jit_inputs, strict=False)
+        traced_model = torch.jit.freeze(traced_model.eval())
+        traced_model(*jit_inputs)
+        traced_model(*jit_inputs)
+
+        model = _ModelFallbackWrapper(traced_model, model)
+
+    model.eval()
+    output_sequences = model.generate(
+        input_ids=input_ids,
+        # max_length=args.length + len(encoded_prompt[0]),
+        max_new_tokens=args.length,
+        temperature=args.temperature,
+        top_k=args.k,
+        top_p=args.p,
+        repetition_penalty=args.repetition_penalty,
+        do_sample=True,
+        num_return_sequences=args.num_return_sequences,
+        streamer=TextStreamer(tokenizer, skip_prompt=True) if args.stream_output else None,
+    )
+
+    # Remove the batch dimension when returning multiple sequences
+    if len(output_sequences.shape) > 2:
+        output_sequences.squeeze_()
+
+    generated_sequences = []
+
+    for generated_sequence_idx, generated_sequence in enumerate(output_sequences):
+        print(f"=== GENERATED SEQUENCE {generated_sequence_idx + 1} (input length: {input_ids.shape[-1]}) ===")
+        generated_sequence = generated_sequence.tolist()
+        # generated_sequence = generated_sequence[len(encoded_prompt[0]):] + tokenizer.encode(' <end_of_prompt> ') + generated_sequence[:len(encoded_prompt[0])]
+
+        # Decode text
+        # text = tokenizer.decode(generated_sequence, clean_up_tokenization_spaces=True)
+        prompt_length = min(input_ids.shape[-1], model.unlimiformer.window_size()) if unlimiformer_args.test_unlimiformer else input_ids.shape[-1]
+        completion = tokenizer.decode(generated_sequence[prompt_length:])
+
+        # Remove all text after the stop token
+        # text = text[: text.find(args.stop_token) if args.stop_token else None]
+
+        # Add the prompt at the beginning of the sequence. Remove the excess text that was used for pre-processing
+        total_sequence = (
+            # prompt_text + 
+            '|||' + completion
+        )
+
+        generated_sequences.append(total_sequence)
+        print(total_sequence)
+
+    return generated_sequences
+
+
+if __name__ == "__main__":
+    main()

unlimiformer/usage.py

@@ -0,0 +1,91 @@
+from .model import Unlimiformer
+from .random_training_unlimiformer import RandomTrainingUnlimiformer
+
+from dataclasses import dataclass, field
+from typing import List, Optional
+
+
+@dataclass
+class UnlimiformerArguments:
+    """
+    Arguments pertaining to what data we are going to input our model for training and eval.
+    """
+    test_unlimiformer: Optional[bool] = field(
+        default=True,
+        metadata={
+            "help": "whether to use KNN."
+        },
+    )
+    unlimiformer_verbose: Optional[bool] = field(
+        default=False,
+        metadata={
+            "help": "whether to print KNN intermediate predictions (mostly for debugging)."
+        },
+    )
+    layer_begin: Optional[int] = field(
+        default=0,
+        metadata={"help": "The layer to begin applying KNN to. KNN will be applied to layers[knn_layer_begin:layer_end]. "
+                          "By default, it will be applied to all layers: [0:None]]"}, 
+    )
+    layer_end: Optional[int] = field(
+        default=None,
+        metadata={"help": "The layer to end applying KNN to. KNN will be applied to layers[knn_layer_begin:layer_end]. "
+                          "By default, it will be applied to all layers: [0:None]]"}, 
+    )
+    unlimiformer_chunk_overlap: Optional[float] = field(
+        default=0.5,
+        metadata={"help": "The fraction of overlap between input chunks"},
+    )
+    unlimiformer_chunk_size: Optional[int] = field(
+        default=None,
+        metadata={"help": "The size of each input chunk"},
+    )
+    unlimiformer_head_num: Optional[int] = field(
+        default=None,
+        metadata={"help": "The head to apply KNN to (if None, apply to all heads)"},
+    )
+    unlimiformer_exclude: Optional[bool] = field(
+        default=False,
+        metadata={
+            "help": "If True, prioritize the inputs that are **not** in the standard attention window."
+        },
+    )
+    random_unlimiformer_training: Optional[bool] = field(
+        default=False,
+    )
+    unlimiformer_training: Optional[bool] = field(
+        default=False,
+    )
+    use_datastore: Optional[bool] = field(default=False)
+    flat_index: Optional[bool] = field(default=False)
+    test_datastore: Optional[bool] = field(default=False)
+    reconstruct_embeddings: Optional[bool] = field(default=False)
+    gpu_datastore: Optional[bool] = field(default=True)
+    gpu_index: Optional[bool] = field(default=True)
+
+
+
+# include these lines in your code somewhere before model training
+def training_addin():
+    if unlimiformer_args.test_unlimiformer:
+        unlimiformer_kwargs = {
+            'layer_begin': unlimiformer_args.layer_begin, 
+            'layer_end': unlimiformer_args.layer_end,
+            'unlimiformer_head_num': unlimiformer_args.unlimiformer_head_num, 
+            'exclude_attention': unlimiformer_args.unlimiformer_exclude, 
+            'chunk_overlap': unlimiformer_args.unlimiformer_chunk_overlap,
+            'model_encoder_max_len': unlimiformer_args.unlimiformer_chunk_size,
+            'verbose': unlimiformer_args.unlimiformer_verbose, 'tokenizer': tokenizer,
+            'unlimiformer_training': unlimiformer_args.unlimiformer_training,
+            'use_datastore': unlimiformer_args.use_datastore,
+            'flat_index': unlimiformer_args.flat_index,
+            'test_datastore': unlimiformer_args.test_datastore,
+            'reconstruct_embeddings': unlimiformer_args.reconstruct_embeddings,
+            'gpu_datastore': unlimiformer_args.gpu_datastore,
+            'gpu_index': unlimiformer_args.gpu_index
+        }
+        if unlimiformer_args.random_unlimiformer_training:
+            model = RandomTrainingUnlimiformer.convert_model(model, **unlimiformer_kwargs)
+        else:
+            model = Unlimiformer.convert_model(model, **unlimiformer_kwargs)
+

unlimiformer/utils/config.py

@@ -0,0 +1,13 @@
+from typing import List
+
+
+def handle_args_to_ignore(args: List[str]):
+    indices_to_remove = []
+    for i, arg in enumerate(args):
+        if "_ignore_" in arg:
+            indices_to_remove.append(i)
+            if not arg.startswith("-"):
+                indices_to_remove.append(i - 1)
+
+    for i in sorted(indices_to_remove, reverse=True):
+        del args[i]

unlimiformer/utils/custom_hf_argument_parser.py

@@ -0,0 +1,39 @@
+import json
+import logging
+import sys
+from typing import Tuple
+
+from transformers import HfArgumentParser
+from transformers.hf_argparser import DataClass
+
+
+class CustomHfArgumentParser(HfArgumentParser):
+    def parse_dictionary_and_args(self) -> Tuple[DataClass, ...]:
+        """
+        Alternative helper method that does not use `argparse` at all, instead loading a json file and populating the
+        dataclass types.
+        """
+        args = []
+        data = {}
+        for i in range(1, len(sys.argv)):
+            if not sys.argv[i].endswith('.json'):
+                break
+
+            with open(sys.argv[i]) as f:
+                new_data = json.load(f)
+            conflicting_keys = set(new_data.keys()).intersection(data.keys())
+            if len(conflicting_keys) > 0:
+                raise ValueError(f'There are conflicting keys in the config files: {conflicting_keys}')
+            data.update(new_data)
+
+        for k, v in data.items():
+            # if any options were given explicitly through the CLA then they override anything defined in the config files
+            if f'--{k}' in sys.argv:
+                logging.info(f'While {k}={v} was given in a config file, a manual override was set through the CLA')
+                continue
+            args.extend(
+                ["--" + k, *(v if isinstance(v, list) else [str(v)])]
+            )  # add the file arguments first so command line args has precedence
+        args += sys.argv[i:]
+
+        return self.parse_args_into_dataclasses(args=args, look_for_args_file=False)

unlimiformer/utils/custom_seq2seq_trainer.py

@@ -0,0 +1,328 @@
+import json
+import math
+import os
+import time
+from collections import defaultdict
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+from datasets import Dataset
+from torch import nn
+from transformers.debug_utils import DebugOption
+from transformers.deepspeed import is_deepspeed_zero3_enabled
+from transformers.trainer_utils import speed_metrics
+
+from transformers.utils import logging
+from transformers import Seq2SeqTrainer, is_torch_tpu_available
+
+import gc
+
+if is_torch_tpu_available(check_device=False):
+    import torch_xla.core.xla_model as xm
+    import torch_xla.debug.metrics as met
+
+
+from utils.decoding import decode
+
+logger = logging.get_logger(__name__)
+
+
+def _clean_memory():
+    gc.collect()
+    torch.cuda.empty_cache()
+
+# This custom trainer is based on the trainer defined in https://github.com/huggingface/transformers/compare/main...eladsegal:public-transformers:scrolls
+class CustomTrainer(Seq2SeqTrainer):
+    def __init__(
+        self, *args, untokenized_eval_dataset=None, data_args=None, output_dir: Optional[str] = None, **kwargs
+    ):
+        super().__init__(*args, **kwargs)
+        self._untokenized_eval_dataset = untokenized_eval_dataset
+        self._max_length = data_args.val_max_target_length
+        self._num_beams = data_args.num_beams
+        self._output_dir = output_dir
+        self._data_args = data_args
+        self.mock_predictions_to_assign_zero_metric_score = self.tokenizer.encode("TOO_MANY_INPUT_TOKENS",return_tensors="np")[0]
+
+    def prediction_step(
+            self,
+            model: nn.Module,
+            inputs: Dict[str, Union[torch.Tensor, Any]],
+            prediction_loss_only: bool,
+            ignore_keys: Optional[List[str]] = None,
+    ) -> Tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]:
+        """
+        Perform an evaluation step on `model` using `inputs`.
+
+        Subclass and override to inject custom behavior.
+
+        Args:
+            model (`nn.Module`):
+                The model to evaluate.
+            inputs (`Dict[str, Union[torch.Tensor, Any]]`):
+                The inputs and targets of the model.
+
+                The dictionary will be unpacked before being fed to the model. Most models expect the targets under the
+                argument `labels`. Check your model's documentation for all accepted arguments.
+            prediction_loss_only (`bool`):
+                Whether or not to ret`urn the loss only.
+
+        Return:
+            Tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]: A tuple with the loss, logits and
+            labels (each being optional).
+        """
+        if not ("labels" in inputs or 'decoder_input_ids' in inputs):
+            if model.training:
+                logger.warning('When computing loss, must give labels or decoder_input_ids. '
+                           'If you only perform prediction, you can safely ignore this message')
+            # This is an issue here because the input may be longer than the max-output length of the model,
+            # and if nothing was given it will shift the input and use it to compute loss (and later discard it).
+            # This may cause an indexing error when absolute embeddings are used (CUDA device side assert)
+            inputs['decoder_input_ids'] = inputs['input_ids'][:,:2].clone()  # dummy outputs
+
+        if not self.args.predict_with_generate or prediction_loss_only:
+            return super().prediction_step(
+                model, inputs, prediction_loss_only=prediction_loss_only, ignore_keys=ignore_keys
+            )
+
+        has_labels = "labels" in inputs
+        inputs = self._prepare_inputs(inputs)
+
+        # XXX: adapt synced_gpus for fairscale as well
+        gen_kwargs = self._gen_kwargs.copy()
+        gen_kwargs["max_length"] = (
+            gen_kwargs["max_length"] if gen_kwargs.get("max_length") is not None else self.model.config.max_length
+        )
+        gen_kwargs["num_beams"] = (
+            gen_kwargs["num_beams"] if gen_kwargs.get("num_beams") is not None else self.model.config.num_beams
+        )
+        default_synced_gpus = True if is_deepspeed_zero3_enabled() else False
+        gen_kwargs["synced_gpus"] = (
+            gen_kwargs["synced_gpus"] if gen_kwargs.get("synced_gpus") is not None else default_synced_gpus
+        )
+
+        if "attention_mask" in inputs:
+            gen_kwargs["attention_mask"] = inputs.get("attention_mask", None)
+        if "global_attention_mask" in inputs:
+            gen_kwargs["global_attention_mask"] = inputs.get("global_attention_mask", None)
+
+        # --------------------- addition compared to the source file --------------------
+        if 'prefix_length' in inputs:
+            gen_kwargs['prefix_length'] = inputs['prefix_length']
+        _clean_memory()
+        # ------------------------------------------------------------------------------
+
+        # prepare generation inputs
+        # some encoder-decoder models can have varying encoder's and thus
+        # varying model input names
+        if hasattr(self.model, "encoder") and self.model.encoder.main_input_name != self.model.main_input_name:
+            generation_inputs = inputs[self.model.encoder.main_input_name]
+        else:
+            generation_inputs = inputs[self.model.main_input_name]
+
+         # Uri: to make sure we use cache even during mid-training evaluation, where this is disabled in general:
+        gen_kwargs['use_cache'] = True
+        
+        generated_tokens = self.model.generate(
+            generation_inputs,
+            **gen_kwargs,
+        )
+        # --------------------- addition compared to the source file --------------------
+        _clean_memory()
+        # ------------------------------------------------------------------------------
+        # in case the batch is shorter than max length, the output should be padded
+        if generated_tokens.shape[-1] < gen_kwargs["max_length"]:
+            generated_tokens = self._pad_tensors_to_max_len(generated_tokens, gen_kwargs["max_length"])
+
+        if has_labels:  # changed the order of the if's here because there is no point going through the model if there are no labels to compute the loss on..
+            with torch.no_grad():
+                with self.compute_loss_context_manager():
+                    outputs = model(**inputs)
+                    if self.label_smoother is not None:
+                        loss = self.label_smoother(outputs, inputs["labels"]).mean().detach()
+                    else:
+                        loss = (outputs["loss"] if isinstance(outputs, dict) else outputs[0]).mean().detach()
+        else:
+            loss = None
+
+        if self.args.prediction_loss_only:
+            return (loss, None, None)
+
+        if has_labels:
+            labels = inputs["labels"]
+            if labels.shape[-1] < gen_kwargs["max_length"]:
+                labels = self._pad_tensors_to_max_len(labels, gen_kwargs["max_length"])
+        else:
+            labels = None
+
+        return (loss, generated_tokens, labels)
+
+    @property
+    def _restart_generator(self):
+        if getattr(self, '_is_restart_generator', False):
+            self._is_restart_generator = False
+            return True
+        return False
+
+    def set_restart_generator(self):
+        self._is_restart_generator = True
+
+    def _get_train_sampler(self) -> Optional[torch.utils.data.Sampler]:
+        sampler = super()._get_train_sampler()
+        try:
+            if self._restart_generator:
+                sampler.generator.manual_seed(self._initial_seed)
+            else:
+                self._initial_seed = sampler.generator.initial_seed()
+        except Exception as e:
+            logger.warning(f'Cannot save or set the seed of the generator: {e}')
+        return sampler
+
+    def _post_process_function(self, untokenized_eval_dataset, predictions):
+        id_to_prediction = {}
+        id_to_label_ids = defaultdict(list)
+
+        assert len(untokenized_eval_dataset) == len(self.eval_dataset)
+
+        for i, (instance, not_valid_for_eval) in enumerate(zip(untokenized_eval_dataset, self.eval_dataset["not_valid_for_eval"])):
+            if not_valid_for_eval:
+                id_to_prediction[instance["id"]] = self.mock_predictions_to_assign_zero_metric_score
+            else:
+                id_to_prediction[instance["id"]] = predictions[i]
+
+            if "outputs" in instance:
+                id_to_label_ids[instance["id"]] = instance["outputs"]
+            else:
+                id_to_label_ids[instance["id"]].append(instance["output"])
+
+        return id_to_prediction, id_to_label_ids
+
+    def evaluate(
+        self,
+        eval_dataset: Optional[Dataset] = None,
+        ignore_keys: Optional[List[str]] = None,
+        metric_key_prefix: str = "eval",
+        untokenized_eval_dataset: Optional[Dataset] = None,
+            **gen_kwargs
+    ) -> Dict[str, float]:
+        """
+        Run evaluation and returns metrics.
+
+        The calling script will be responsible for providing a method to compute metrics, as they are task-dependent
+        (pass it to the init `compute_metrics` argument).
+
+        You can also subclass and override this method to inject custom behavior.
+
+        Args:
+            eval_dataset (`Dataset`, *optional*):
+                Pass a dataset if you wish to override `self.eval_dataset`. If it is an [`~datasets.Dataset`], columns
+                not accepted by the `model.forward()` method are automatically removed. It must implement the `__len__`
+                method.
+            ignore_keys (`List[str]`, *optional*):
+                A list of keys in the output of your model (if it is a dictionary) that should be ignored when
+                gathering predictions.
+            metric_key_prefix (`str`, *optional*, defaults to `"eval"`):
+                An optional prefix to be used as the metrics key prefix. For example the metrics "bleu" will be named
+                "eval_bleu" if the prefix is `"eval"` (default)
+            max_length (`int`, *optional*):
+                The maximum target length to use when predicting with the generate method.
+            num_beams (`int`, *optional*):
+                Number of beams for beam search that will be used when predicting with the generate method. 1 means no
+                beam search.
+            gen_kwargs:
+                Additional `generate` specific kwargs.
+
+        Returns:
+            A dictionary containing the evaluation loss and the potential metrics computed from the predictions. The
+            dictionary also contains the epoch number which comes from the training state.
+        """
+
+        gen_kwargs = gen_kwargs.copy()
+        gen_kwargs["max_length"] = (
+            gen_kwargs["max_length"] if gen_kwargs.get("max_length") is not None else self.args.generation_max_length
+        )
+        gen_kwargs["num_beams"] = (
+            gen_kwargs["num_beams"] if gen_kwargs.get("num_beams") is not None else self.args.generation_num_beams
+        )
+        self._gen_kwargs = gen_kwargs
+
+        self._memory_tracker.start()
+
+        eval_dataloader = self.get_eval_dataloader(eval_dataset)
+        # ----------------------------------- Added -----------------------------------
+        untokenized_eval_dataset = (
+            self._untokenized_eval_dataset if untokenized_eval_dataset is None else untokenized_eval_dataset
+        )
+        compute_metrics = self.compute_metrics
+        self.compute_metrics = None
+        # -----------------------------------------------------------------------------
+
+        start_time = time.time()
+
+        eval_loop = self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
+        try:
+            output = eval_loop(
+                eval_dataloader,
+                description="Evaluation",
+                # No point gathering the predictions if there are no metrics, otherwise we defer to
+                # self.args.prediction_loss_only
+                prediction_loss_only=None,  # MODIFIED since we need the predictions
+                ignore_keys=ignore_keys,
+                metric_key_prefix=metric_key_prefix,
+            )
+        finally:
+            # ----------------------------------- Added -----------------------------------
+            # revert the compute metrics back
+            self.compute_metrics = compute_metrics
+            # -----------------------------------------------------------------------------
+
+        # ----------------------------------- Added -----------------------------------
+        # compute our metrics
+        if output.predictions is not None:
+            eval_preds = self._post_process_function(untokenized_eval_dataset, output.predictions)
+
+            if self._output_dir is not None and self.is_world_process_zero():
+                predictions = decode(eval_preds[0], self.tokenizer, self._data_args)
+                output_prediction_file = os.path.join(
+                    self._output_dir, f"generated_predictions_eval_{self.state.global_step}.json"
+                )
+                with open(output_prediction_file, "w") as writer:
+                    json.dump(predictions, writer, indent=4)
+
+                output_labels_file = os.path.join(
+                    self._output_dir, f"eval_labels.json"
+                )
+                if not os.path.isfile(output_labels_file):
+                    with open(output_labels_file, "w") as writer:
+                        json.dump(eval_preds[1], writer, indent=4)
+
+            if self.compute_metrics is not None:
+                output.metrics.update(self.compute_metrics(*eval_preds))
+
+            # Prefix all keys with metric_key_prefix + '_'
+        for key in list(output.metrics.keys()):
+            if not key.startswith(f"{metric_key_prefix}_"):
+                output.metrics[f"{metric_key_prefix}_{key}"] = output.metrics.pop(key)
+        # -----------------------------------------------------------------------------
+
+        total_batch_size = self.args.eval_batch_size * self.args.world_size
+        output.metrics.update(
+            speed_metrics(
+                metric_key_prefix,
+                start_time,
+                num_samples=output.num_samples,
+                num_steps=math.ceil(output.num_samples / total_batch_size),
+            )
+        )
+
+        self.log(output.metrics)
+
+        if DebugOption.TPU_METRICS_DEBUG in self.args.debug:
+            # tpu-comment: Logging debug metrics for PyTorch/XLA (compile, execute times, ops, etc.)
+            xm.master_print(met.metrics_report())
+
+        self.control = self.callback_handler.on_evaluate(self.args, self.state, self.control, output.metrics)
+
+        self._memory_tracker.stop_and_update_metrics(output.metrics)
+
+        return output.metrics

unlimiformer/utils/decoding.py

@@ -0,0 +1,59 @@
+from functools import partial
+import numpy as np
+
+
+def decode(id_to_something, tokenizer=None, data_args=None):
+    decode_fn = None
+    switch_case = None
+    elem = next(iter(id_to_something.values()))
+    if isinstance(elem, str):
+        switch_case = -1
+        decode_fn = lambda text: text.strip()
+    elif isinstance(elem, list) and not isinstance(elem[0], int):
+        if isinstance(elem[0], str):
+            switch_case = 0
+            decode_fn = lambda texts: [text.strip() for text in texts]
+        else:
+            switch_case = 1
+            decode_fn = lambda token_ids_list: [
+                text.strip()
+                for text in partial(
+                    tokenizer.batch_decode, skip_special_tokens=True, clean_up_tokenization_spaces=True
+                )(token_ids_list)
+            ]
+    else:
+        switch_case = 2
+        decode_fn = lambda token_ids: partial(
+            tokenizer.decode, skip_special_tokens=True, clean_up_tokenization_spaces=True
+        )(token_ids).strip()
+
+    id_to_text = {}
+    for id_, something in id_to_something.items():
+        if switch_case == -1 or switch_case == 0:
+            obj_to_decode = something
+        else:
+            if data_args is None:
+                data_args = {}
+            if not isinstance(data_args, dict):
+                data_args = vars(data_args)
+            if data_args.get("ignore_pad_token_for_loss", True):
+                # Replace -100 in the token_ids as we can't decode them.
+                if switch_case == 1:
+                    token_ids_list = something
+                    for i in range(len(token_ids_list)):
+                        token_ids_list[i] = _replace_padding(token_ids_list[i], tokenizer.pad_token_id)
+                    obj_to_decode = token_ids_list
+                elif switch_case == 2:
+                    token_ids = something
+                    token_ids = _replace_padding(token_ids, tokenizer.pad_token_id)
+                    obj_to_decode = token_ids
+            else:
+                obj_to_decode = something
+
+        id_to_text[id_] = decode_fn(obj_to_decode)
+
+    return id_to_text
+
+
+def _replace_padding(token_ids: np.array, pad_token_id):
+    return np.where(token_ids != -100, token_ids, pad_token_id)

unlimiformer/utils/duplicates.py

@@ -0,0 +1,15 @@
+def drop_duplicates_in_input(untokenized_dataset):
+    indices_to_keep = []
+    id_to_idx = {}
+    outputs = []
+    for i, (id_, output) in enumerate(zip(untokenized_dataset["id"], untokenized_dataset["output"])):
+        if id_ in id_to_idx:
+            outputs[id_to_idx[id_]].append(output)
+            continue
+        indices_to_keep.append(i)
+        id_to_idx[id_] = len(outputs)
+        outputs.append([output])
+    untokenized_dataset = untokenized_dataset.select(indices_to_keep).flatten_indices()
+    untokenized_dataset = untokenized_dataset.remove_columns("output")
+    untokenized_dataset = untokenized_dataset.add_column("outputs", outputs)
+    return untokenized_dataset

unlimiformer/utils/override_training_args.py

@@ -0,0 +1,106 @@
+
+import os
+import sys
+
+import torch.cuda
+from transformers.utils import logging
+
+sys.path.insert(0, os.getcwd())
+
+from dataclasses import dataclass, field
+
+from transformers.trainer_utils import IntervalStrategy
+from transformers import Seq2SeqTrainingArguments
+
+logger = logging.get_logger('swed_logger')
+
+@dataclass
+class TrainingOverridesArguments(Seq2SeqTrainingArguments):
+    """
+    To use if, it requires evaluation_strategy == IntervalStrategy.STEPS
+    """
+    eval_steps_override: float = field(default=0, metadata={"help": "a fraction, to set the the save_steps w.r.t to number of steps in "
+                                                                    "a single epoch. changes eval_steps. 0 to disable (default)"})
+    save_steps_override: float = field(default=0, metadata={"help": "a fraction, to set the the save_steps w.r.t to number of steps in "
+                                                                    "a single epoch. changes save_steps. must be a multiple of eval_steps"
+                                                                    " (or eval_steps_override if given). 0 to disable (default)"})
+
+    eval_fraction: float = field(default=1, metadata={
+        "help": "A float in (0,1] that corresponds to how much of the eval set to use during evaluations "
+                "(same subset all the time) or an integer >= 2 which amounts to the absolute number of training "
+                "samples to use. 1. to disable it and use the entire eval set "})
+
+    use_auth_token: bool = field(
+        default=False,
+        metadata={
+            "help": "Will use the token generated when running `transformers-cli login` (necessary to use this script "
+            "with private models). If AUTH_TOKEN is set as an environment variable, would use that"
+        },
+    )
+
+    fp16_padding: bool = field(
+        default=False,
+        metadata={"help": "Whether to use padding for fp16"},
+    )
+
+
+    def __post_init__(self):
+        super(TrainingOverridesArguments, self).__post_init__()
+        if self.eval_steps_override > 0 or self.save_steps_override > 0:
+            if self.evaluation_strategy != IntervalStrategy.STEPS:
+                raise ValueError(
+                    f"using eval/save steps override requires evaluation strategy to be  {IntervalStrategy.STEPS}"
+                )
+            if self.save_steps_override == 0 or self.eval_steps_override == 0:
+                raise ValueError(
+                    f"using eval/save steps override requires both overrides to be non zero"
+                )
+            diff = (self.save_steps_override / self.eval_steps_override) % 1
+            if min(1-diff, diff) > 1e-5:  # we do it like that to support fractions modulo as well, with loss of precision
+                raise ValueError(
+                    f"using eval/save steps override requires save steps override to be a multiple of eval_steps_override"
+                )
+        if self.use_auth_token and 'AUTH_TOKEN' in os.environ:
+            self.use_auth_token = os.getenv('AUTH_TOKEN')
+
+    @property
+    def effective_batch_size(self):
+        if not hasattr(self, '_ebs'):
+            n_gpu = self.n_gpu if torch.cuda.is_available() else 1  # may be on cpu
+            self._ebs = self.per_device_train_batch_size * self.gradient_accumulation_steps * n_gpu
+            logger.warning(f'Training with {self.per_device_train_batch_size} per_device_train_size, {self.n_gpu} gpus and '
+                        f'{self.gradient_accumulation_steps} gradient accumulation steps, resulting in {self._ebs} effective batch size')
+        return self._ebs
+
+    def apply_overrides(self, dataset_size):
+        # Uri:
+        return
+        
+        if self.eval_steps_override == 0:
+            return
+        es, ss = self.eval_steps, self.save_steps
+        total_steps_per_epoch = dataset_size / self.effective_batch_size  # note that this may not be an  integer
+        eval_steps = int(total_steps_per_epoch * self.eval_steps_override)
+        if eval_steps >= self.logging_steps:
+            if eval_steps % self.logging_steps != 0:
+                logger.warning(f'Eval steps override would result in eval every {eval_steps} steps, but it is not a '
+                            f'multiple of logging steps ({self.logging_steps}) so changing to '
+                            f'{eval_steps + self.logging_steps - eval_steps % self.logging_steps}')
+                eval_steps = eval_steps + self.logging_steps - eval_steps % self.logging_steps
+        elif eval_steps < self.logging_steps:
+            logger.warning(f'Eval steps override would result in eval every {eval_steps} steps, but it is not a '
+                        f'multiple of logging steps ({self.logging_steps}) so changing to {self.logging_steps}')
+            eval_steps = self.logging_steps
+        self.eval_steps = eval_steps
+
+        save_steps = int(total_steps_per_epoch * self.save_steps_override)
+        if save_steps < eval_steps or save_steps % eval_steps != 0:
+            logger.warning(f'Save steps override would result in eval every {save_steps} steps, but it is not a '
+                        f'multiple of eval steps ({eval_steps}) so changing to '
+                        f'{save_steps + eval_steps - save_steps % self.eval_steps}')
+            save_steps = save_steps + eval_steps - save_steps % self.eval_steps
+        self.save_steps = save_steps
+
+        logger.warning(f'Using overrides with dataset of size {dataset_size} and effective batch size of '
+                    f'{self.effective_batch_size}, moving from (eval_steps, save_steps) '
+                    f'of {(es, ss)} to {(self.eval_steps, self.save_steps)}')

utils.py

@@ -2,6 +2,7 @@ import re
 import spacy
 import json
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForCausalLM, AutoModel
+from unlimiformer import Unlimiformer, UnlimiformerArguments
 import streamlit as st
 from urllib.request import Request, urlopen, HTTPError
 from bs4 import BeautifulSoup
@@ -16,8 +17,29 @@ def hide_footer():
     st.markdown(hide_st_style, unsafe_allow_html=True)
 
 @st.cache_resource
-def get_seq2seq_model(model_id):
-    return AutoModelForSeq2SeqLM.from_pretrained(model_id)
+def get_seq2seq_model(model_id, use_unlimiformer=True, _tokenizer=None):
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
+    if use_unlimiformer:
+        defaults = UnlimiformerArguments()
+        unlimiformer_kwargs = {
+            'layer_begin': defaults.layer_begin, 
+            'layer_end': defaults.layer_end,
+            'unlimiformer_head_num': defaults.unlimiformer_head_num, 
+            'exclude_attention': defaults.unlimiformer_exclude, 
+            'chunk_overlap': defaults.unlimiformer_chunk_overlap,
+            'model_encoder_max_len': defaults.unlimiformer_chunk_size,
+            'verbose': defaults.unlimiformer_verbose, 'tokenizer': _tokenizer,
+            'unlimiformer_training': defaults.unlimiformer_training,
+            'use_datastore': defaults.use_datastore,
+            'flat_index': defaults.flat_index,
+            'test_datastore': defaults.test_datastore,
+            'reconstruct_embeddings': defaults.reconstruct_embeddings,
+            'gpu_datastore': defaults.gpu_datastore,
+            'gpu_index': defaults.gpu_index
+        }
+        return Unlimiformer.convert_model(model, **unlimiformer_kwargs)
+    else:
+        return model
 
 @st.cache_resource
 def get_causal_model(model_id):