Finally

.gitignore

@@ -0,0 +1,2 @@
+__pycache__/
+flagged/

app.py

@@ -0,0 +1,70 @@
+import gradio as gr
+import torch
+from transformers import GPT2TokenizerFast, GPT2LMHeadModel
+
+from gpt2_knn_attention import GPT2KNNAttention
+from knn_memory import KNNLayer, ClearMemoryLayer
+
+
+def inject_knn_in_gpt2(model, knn_memory, bos_token_id, eos_token_id, device, layer_ind=8):
+    layer = model.transformer.h[layer_ind].attn
+    state = layer.state_dict()
+    knn_layer = GPT2KNNAttention(
+        config, knn_memory, device, is_cross_attention=False, layer_idx=layer.layer_idx)
+    knn_state = knn_layer.state_dict()
+
+    for k, v in state.items():
+        knn_state[k] = v
+
+    knn_layer.load_state_dict(knn_state)
+
+    model.transformer.h[8].attn = knn_layer
+    model.transformer = ClearMemoryLayer(
+        knn_memory, bos_token_id, eos_token_id, model.transformer)
+    model.eval()
+
+
+model_name = "gpt2"
+tokenizer = GPT2TokenizerFast.from_pretrained(model_name)
+model = GPT2LMHeadModel.from_pretrained(model_name)
+config = model.config
+model.eval()
+
+knn_memory = KNNLayer(config, share_memory=False, batch_size=1)
+bos_token_id, eos_token_id = tokenizer.bos_token_id, tokenizer.eos_token_id
+bos_token, eos_token = tokenizer.bos_token, tokenizer.eos_token
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+knn_model = inject_knn_in_gpt2(
+    model, knn_memory, bos_token_id, eos_token_id, device, layer_ind=8)
+knn_model.load_state_dict(torch.load('gpt2_knn_attention.pt'))
+
+
+def generate(text, temperature, max_new_tokens, top_p):
+    encoded_input = tokenizer(text, return_tensors='pt')
+    output = model.generate(**encoded_input, do_sample=True,
+                            max_new_tokens=int(max_new_tokens), temperature=temperature, top_p=top_p)
+    return tokenizer.decode(output[0])
+
+
+desc = "Попытка повторить статью от Google (Memorizing Transformers)[https://arxiv.org/abs/2203.08913]. "\
+       "В ней вводиться новый слой **KNNAttention**, который использует approximate kNN в базе с (key, value), чтобы делать attention по большому контексту. Это позволяет расширить контекст трансформера до размера книг и статей, несильно замедляя его.\n\n"\
+       "Я написал свои **KNNAttention**, переписал слой **GPT2Attention**, чтобы он использовал KNNAttention, а также написал несколько вспомогательный классов для всего этого.\n\n"\
+       "Я сел писать это за **3 недели** до дедлайна, но все равно не довел до результата, которого изначально хотел. Но я доволен проделанной работой :)"
+
+
+demo = gr.Interface(
+    fn=generate,
+    inputs=[gr.inputs.Textbox(lines=5, label="Input Text"),
+            gr.Slider(0.001, 2.0, step=0.05, value=0.8, label='temperature'),
+            gr.Slider(1, 512, step=1, value=32, label='max_new_tokens'),
+            gr.Slider(0.1, 1.0, step=0.02, value=0.92, label='top_p')],
+    outputs=gr.outputs.Textbox(label="Generated Text"),
+    description=desc,
+    title="Memorizing Transformers",
+    examples=[
+        ["My name is Lewis and I like to", 0.8, 32, 0.92]
+    ]
+)
+
+demo.launch()

batched_dataloader.py

@@ -0,0 +1,48 @@
+from torch.utils.data import DataLoader
+from transformers import DataCollatorWithPadding
+
+# some utils for training
+
+
+class BooksBatcherIter:
+    def __init__(self, data_iter, batch_size, tokenizer, chunk_size=1024):
+        self.data_iter = data_iter
+        self.batch_size = batch_size
+        self.chunk_size = chunk_size
+        self.batch_fns = [self._batch_fn()]
+        self.collate_fn = DataCollatorWithPadding(tokenizer)
+
+    def _batch_fn(self):
+        for book in self.data_iter:
+            for i in range(0, len(book), self.chunk_size):
+                yield book[i:i+self.chunk_size]
+
+    def __iter__(self) -> 'BooksBatcherIter':
+        return self
+
+    def __next__(self) -> Any:
+        batch = []
+
+        try:
+            for b in self.batch_fns:
+                batch.append(next(b))
+        except StopIteration:
+            raise StopIteration
+
+        return self.collate_fn(batch)
+
+
+class BooksBatcher:
+    def __init__(self, dataset, batch_size, tokenizer) -> None:
+        self.batch_size = batch_size
+        self.tokenizer = tokenizer
+        self.dataloader = DataLoader(
+            dataset=dataset,
+            batch_size=None,  # return raw samples
+            shuffle=True,
+            num_workers=2,
+            prefetch_factor=4
+        )
+
+    def __iter__(self) -> 'BooksBatcherIter':
+        return BooksBatcherIter(iter(self.dataloader), self.batch_size, self.tokenizer)

gpt2_knn_attention.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:610319abc523c64f595d4cb8fec2ff7faeab331b05722206c6c42656eae0bdff
+size 510408492

gpt2_knn_attention.py

@@ -0,0 +1,151 @@
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from transformers.models.gpt2.modeling_gpt2 import GPT2Attention
+
+
+class GPT2KNNAttention(GPT2Attention):
+    def __init__(self, config, knn_memory, device, is_cross_attention=False, layer_idx=None, num_retrieve_memories=32):
+        super().__init__(config, is_cross_attention, layer_idx)
+
+        self.knn_memory = knn_memory
+        self.device = device
+        self.num_retrieve_memories = num_retrieve_memories
+        self.knn_attn_dropout = nn.Dropout(config.attn_pdrop)
+        self.attn_comb_bias = nn.Parameter(torch.empty(self.num_heads,))
+        nn.init.normal_(self.attn_comb_bias, mean=0.0, std=1.0)
+        # self.attn_comb_bias = nn.Parameter(torch.full((self.num_heads,), 1.0))
+
+    def _knn_attn(self, query, key, value, mask, head_mask=None):
+        query = query.unsqueeze(-2)
+        attn_weights = torch.matmul(query, key.transpose(-1, -2))
+
+        if self.scale_attn_weights:
+            attn_weights = attn_weights / torch.full(
+                [], value.size(-1) ** 0.5, dtype=attn_weights.dtype, device=attn_weights.device
+            )
+
+        # Layer-wise attention scaling
+        if self.scale_attn_by_inverse_layer_idx:
+            attn_weights = attn_weights / float(self.layer_idx + 1)
+
+        # if not self.is_cross_attention:
+        #     raise RuntimeError("KNN attention is not yet implemented for !cross_attention")
+        #     # if only "normal" attention layer implements causal mask
+        #     query_length, key_length = query.size(-3), key.size(-3)
+        #     causal_mask = self.bias[:, :, key_length - query_length : key_length, :key_length]
+        #     mask_value = torch.finfo(attn_weights.dtype).min
+        #     # Need to be a tensor, otherwise we get error: `RuntimeError: expected scalar type float but found double`.
+        #     # Need to be on the same device, otherwise `RuntimeError: ..., x and y to be on the same device`
+        #     mask_value = torch.full([], mask_value, dtype=attn_weights.dtype).to(attn_weights.device)
+        #     attn_weights = torch.where(causal_mask, attn_weights.to(attn_weights.dtype), mask_value)
+
+        attn_weights = nn.functional.softmax(attn_weights, dim=-1)
+
+        # Downcast (if necessary) back to V's dtype (if in mixed-precision) -- No-Op otherwise
+        attn_weights = attn_weights.type(value.dtype)
+        attn_weights = self.knn_attn_dropout(attn_weights)
+
+        # masking missing keys
+        sh = mask.size()
+        attn_weights = attn_weights * mask.view((sh[0], 1, 1, 1, sh[1]))
+
+        # Mask heads if we want to
+        if head_mask is not None:
+            attn_weights = attn_weights * head_mask
+
+        attn_output = torch.matmul(attn_weights, value)
+        attn_output.squeeze_(dim=-2)
+
+        return attn_output
+
+    def _attn(self, query, key, value, attention_mask=None, head_mask=None):
+        attn_output, attn_weights = super()._attn(
+            query, key, value, attention_mask, head_mask)
+        knn_key, knn_value, knn_mask = self.knn_memory.search(
+            query, self.num_retrieve_memories)
+        g = torch.sigmoid(self.attn_comb_bias)[:, None, None]
+
+        if knn_key.numel() == 0:
+            return attn_output * (1 - g), attn_weights
+
+        knn_key, knn_value, knn_mask = knn_key.to(
+            self.device), knn_value.to(self.device), knn_mask.to(self.device)
+        knn_attn_output = self._knn_attn(
+            query, knn_key, knn_value, knn_mask, head_mask)
+
+        # combining two attentions
+        attn = knn_attn_output * g + attn_output * (1 - g)
+
+        return attn, attn_weights
+
+    def _upcast_and_reordered_attn(self, query, key, value, attention_mask=None, head_mask=None):
+        raise RuntimeError(
+            "KNN attention is not yet implemented for _upcast_and_reordered_attn")
+
+    def forward(
+        self,
+        hidden_states: Optional[Tuple[torch.FloatTensor]],
+        layer_past: Optional[Tuple[torch.Tensor]] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        head_mask: Optional[torch.FloatTensor] = None,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        encoder_attention_mask: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = False,
+        output_attentions: Optional[bool] = False,
+    ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]], ...]:
+        if encoder_hidden_states is not None:
+            if not hasattr(self, "q_attn"):
+                raise ValueError(
+                    "If class is used as cross attention, the weights `q_attn` have to be defined. "
+                    "Please make sure to instantiate class with `GPT2Attention(..., is_cross_attention=True)`."
+                )
+
+            query = self.q_attn(hidden_states)
+            key, value = self.c_attn(encoder_hidden_states).split(
+                self.split_size, dim=2)
+            attention_mask = encoder_attention_mask
+        else:
+            query, key, value = self.c_attn(
+                hidden_states).split(self.split_size, dim=2)
+
+        query = self._split_heads(query, self.num_heads, self.head_dim)
+        key = self._split_heads(key, self.num_heads, self.head_dim)
+        value = self._split_heads(value, self.num_heads, self.head_dim)
+
+        # normalization of queries and keys reduces the effect of staleness
+        query, key = F.normalize(query, dim=-1), F.normalize(key, dim=-1)
+        new_memories = (key, value)
+
+        if layer_past is not None:
+            past_key, past_value = layer_past
+            key = torch.cat((past_key, key), dim=-2)
+            value = torch.cat((past_value, value), dim=-2)
+
+        if use_cache is True:
+            present = (key, value)
+        else:
+            present = None
+
+        if self.reorder_and_upcast_attn:
+            raise RuntimeError("Not implemented")
+            attn_output, attn_weights = self._upcast_and_reordered_attn(
+                query, key, value, attention_mask, head_mask)
+        else:
+            attn_output, attn_weights = self._attn(
+                query, key, value, attention_mask, head_mask)
+
+        attn_output = self._merge_heads(
+            attn_output, self.num_heads, self.head_dim)
+        attn_output = self.c_proj(attn_output)
+        attn_output = self.resid_dropout(attn_output)
+
+        outputs = (attn_output, present)
+        if output_attentions:
+            outputs += (attn_weights,)
+
+        self.knn_memory.add(*new_memories)
+
+        return outputs  # a, present, (attentions)

knn_memory.py

@@ -0,0 +1,220 @@
+from collections import deque
+
+import faiss
+import torch
+import torch.nn.functional as F
+import numpy as np
+from torch import nn
+
+
+class KNN:
+    """
+    KNN for one element in batch. Handles all heads
+    """
+
+    def __init__(self, num_heads, head_dim, memories_size=16000, shrink_size=None, cache=None):
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        self.memories_size = memories_size
+        self.shrink_size = shrink_size or memories_size * 1.1
+        self.indexes = [faiss.IndexFlat(
+            self.head_dim, faiss.METRIC_INNER_PRODUCT) for _ in range(self.num_heads)]
+        self.values = [deque([]) for _ in range(self.num_heads)]
+        self.cache = cache
+
+    def __del__(self):
+        if hasattr(self, 'indexes'):
+            del self.indexes
+            del self.values
+
+    def clear(self):
+        for index in self.indexes:
+            index.reset()
+
+        for value in self.values:
+            value.clear()
+
+    def shrink(self):
+        """Shrinks index to memories_size"""
+
+        for i, index in enumerate(self.indexes):
+            if index.ntotal > self.shrink_size:
+                to_delete = index.ntotal - self.memories_size
+                index.remove_ids(np.arange(0, to_delete))
+
+                for _ in range(to_delete):
+                    self.values[i].popleft()
+
+    def add(self, key, value):
+        for i, k in enumerate(key):
+            self.indexes[i].add(k)
+        for i, v in enumerate(value):
+            self.values[i].extend(v)
+
+        if self.cache is not None:
+            raise RuntimeError("Cache for KNN not implemented")
+            # self.cache.add(key)
+
+        self.shrink()
+
+    def search(self, query, k=32):
+        """
+        Searchs for query in keys' index.
+        Returns k most relevant keys and corresponding values
+        """
+
+        k = min(k, len(self.values[0]))
+
+        if k <= 0:
+            return torch.empty((query.shape[0], query.shape[1], 0, query.shape[2])),\
+                torch.empty(
+                    (query.shape[0], query.shape[1], 0, query.shape[2]))
+
+        Ks, Vs = [], []
+
+        for i, q in enumerate(query):
+            D, I, K = self.indexes[i].search_and_reconstruct(q, k=k)
+            V = np.take(self.values[i], indices=I, axis=0)
+            Ks.append(K)
+            Vs.append(V)
+
+        return np.stack(Ks, axis=0), np.stack(Vs, axis=0)
+
+
+class KNNLayer:
+    """
+    KNN Attention layer. Handles KNN's for batch (every elemnt separately)
+    """
+
+    def __init__(self, config, share_memory=True, batch_size=None, memory_size=16000, shrink_size=None, n_jobs=4, cache=None):
+        if not share_memory and batch_size is None:
+            raise RuntimeError(
+                "If share_memory is False, batch_size should be passed")
+
+        self.embed_dim = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = self.embed_dim // self.num_heads
+
+        self.share_memory = share_memory
+        self.batch_size = batch_size
+        self.memory_size = memory_size
+        self.shrink_size = shrink_size or self.memory_size * 1.1
+        self.closed = False
+
+        if not share_memory:
+            self.knns = [KNN(self.num_heads, self.head_dim, memory_size,
+                             self.shrink_size, cache=cache) for _ in range(self.batch_size)]
+        else:
+            self.knn = KNN(self.num_heads, self.head_dim,
+                           memory_size, self.shrink_size, cache=cache)
+
+        faiss.omp_set_num_threads(n_jobs)
+
+    def clear_batches(self, batch_indexes):
+        if self.closed:
+            return
+
+        if not self.share_memory:
+            for idx in batch_indexes:
+                self.knns[idx].clear()
+
+    def clear(self):
+        if self.closed:
+            return
+
+        if self.share_memory:
+            self.knn.clear()
+        else:
+            for idx in range(len(self.knns)):
+                self.knns[idx].clear()
+
+    def add(self, keys, values):
+        if self.closed:
+            return
+
+        keys, values = keys.numpy(force=True), values.numpy(force=True)
+        if not self.share_memory:
+            for i, (key, value) in enumerate(zip(keys, values)):
+                self.knns[i].add(key, value)
+        else:
+            for key, value in zip(keys, values):
+                self.knn.add(key, value)
+
+    def search(self, queries, k=32):
+        queries = queries.numpy(force=True)
+        keys, values = [], []
+        max_len = 0
+
+        if self.share_memory:
+            for query in queries:
+                key, value = self.knn.search(query, k)
+                keys.append(key)
+                values.append(value)
+                max_len = max(max_len, key.shape[2])
+        else:
+            for i, query in enumerate(queries):
+                key, value = self.knns[i].search(query, k)
+                keys.append(key)
+                values.append(value)
+                max_len = max(max_len, key.shape[2])
+
+        masks = np.ones((len(keys), max_len), dtype=np.float32)
+
+        for i, (key, value) in enumerate(zip(keys, values)):
+            l = key.shape[2]
+
+            if l == max_len:
+                continue
+            elif l > max_len:
+                raise RuntimeError("What? max_len is not max")
+
+            sh = list(key.shape)
+            sh[2] = max_len - sh[2]
+            keys[i] = np.concatenate(
+                (key, np.zeros(sh, dtype=np.float32)), axis=2)
+            values[i] = np.concatenate(
+                (value, np.zeros(sh, dtype=np.float32)), axis=2)
+            masks[i, l:] = 0
+
+        return torch.from_numpy(np.stack(keys, axis=0)),\
+            torch.from_numpy(np.stack(values, axis=0)),\
+            torch.from_numpy(masks)
+
+    def close(self):
+        self.closed = True
+
+    def open(self):
+        self.closed = False
+
+    def reset(self):
+        self.open()
+        self.clear()
+
+
+class ClearMemoryLayer(nn.Module):
+    def __init__(self, knn_memory, bos_token, eos_token, next_layer):
+        super().__init__()
+
+        self.knn_memory = knn_memory
+        self.bos_token = bos_token
+        self.eos_token = eos_token
+        self.next_layer = next_layer
+
+    def _clear_if_token(self, tokens, token):
+        batches_to_clear = (tokens == token).any(dim=-1).nonzero()
+
+        if len(batches_to_clear) > 0:
+            self.knn_memory.clear_batches(batches_to_clear[0])
+
+    def forward(self, tokens, *args, **kwargs):
+        # self._clear_if_token(tokens, self.bos_token)
+
+        batches_to_clear = (tokens[:, 0] == self.bos_token).nonzero()
+
+        if len(batches_to_clear) > 0:
+            self.knn_memory.clear_batches(batches_to_clear[:, 0])
+
+        res = self.next_layer(tokens, *args, **kwargs)
+        # self._clear_if_token(tokens, self.eos_token)
+
+        return res

requirements.txt

@@ -0,0 +1,5 @@
+transformers==4.28.1
+torch==2.0.0
+faiss-cpu==1.7.4
+gradio==3.27.0
+numpy==1.21.2

vector_cache.py

@@ -0,0 +1,57 @@
+import json
+import os
+
+import numpy as np
+
+
+class VectorCache:
+    """
+    Caches vectors on disk so one can later build an index on them (indexes like IVF requires big amount of vetores for building)
+    """
+
+    def __init__(self, filename='vector_cache.memmap', d=768, size=7000000):
+        self.filename = filename
+        self.offset_file = filename + '.offset'
+        self.d = d
+        self.size = size
+
+        if os.path.isfile(filename):
+            mode = 'r+'
+            self.f = open(self.offset_file, mode)
+            data = json.load(self.f)
+            self.offset = data[0]
+            self.length = data[1]
+        else:
+            mode = 'w+'
+            self.f = open(self.offset_file, mode)
+            self.offset = 0
+            self.length = 0
+
+        self.db = np.memmap(filename, dtype=np.float32, mode='w+',
+                            shape=(size, d), order='C')
+
+    def sync_offset(self):
+        self.f.seek(0)
+        self.f.truncate(0)
+        self.f.write(json.dumps([self.offset, self.length]))
+
+    def close(self):
+        self.db.flush()
+        self.db.close()
+
+        self.sync_offset()
+        self.f.flush()
+        self.f.close()
+
+    def add(self, vs):
+        l = len(vs)
+        to_end = self.size - self.offset
+
+        if to_end < l:
+            self.add(vs[:to_end])
+            self.add(vs[to_end:])
+            return
+
+        self.db[self.offset:self.offset+l+1, :] = vs
+        self.offset = (self.offset + l + 1) % self.size
+        self.length = min(self.length + l, self.size)