modeling.py

# Copyright 2023 HuggingFace Inc. team and MosaicML NLP team.
#
# 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.

# This code has been adapted from Mosaic ML and Huggingface and inherits the above lisence.
# The original code can be found here:
# https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py
# We annotate the edited code below with 'EM' comments to indicate where we have made changes.
"""PyTorch MPT model."""

import math
from typing import Optional, Tuple, Union

import faiss
import numpy as np
import torch
import torch.utils.checkpoint
from einops import rearrange
from torch import nn
from torch.linalg import vector_norm
from torch.nn import CrossEntropyLoss, LayerNorm
from torch.nn import functional as F
from transformers.file_utils import (
    add_code_sample_docstrings,
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
)
from transformers.modeling_outputs import (
    BaseModelOutputWithPastAndCrossAttentions,
    CausalLMOutputWithCrossAttentions,
)
from transformers.modeling_utils import PreTrainedModel
from transformers.utils import logging

from .configuration import ExtendedMptConfig

logger = logging.get_logger(__name__)

_CHECKPOINT_FOR_DOC = "mosaicml/mpt-7b"
_CONFIG_FOR_DOC = "MptConfig"


# Copied from transformers.models.bloom.modeling_bloom._make_causal_mask
def _make_causal_mask(
    input_ids_shape: torch.Size, device: torch.device, past_key_values_length: int
) -> torch.BoolTensor:
    """
    Make causal mask used for self-attention.
    """
    batch_size, target_length = input_ids_shape
    mask = torch.empty(
        (target_length, target_length + past_key_values_length),
        dtype=torch.bool,
        device=device,
    )
    # ONNX doesn't support `torch.Tensor.triu` properly, thus we use this workaround
    seq_ids = torch.arange(target_length, device=device)
    mask[:, past_key_values_length:] = seq_ids[:, None] < seq_ids[None, :]

    if past_key_values_length > 0:
        mask[:, :past_key_values_length] = False

    expanded_mask = mask[None, None, :, :].expand(
        batch_size, 1, target_length, target_length + past_key_values_length
    )
    return expanded_mask


# Copied from transformers.models.bloom.modeling_bloom._expand_mask
def _expand_mask(mask: torch.Tensor, tgt_length: int) -> torch.BoolTensor:
    """
    Expands attention_mask from `[batch_size, src_length]` to `[batch_size, 1, tgt_length, src_length]`.
    """
    batch_size, src_length = mask.shape
    tgt_length = tgt_length if tgt_length is not None else src_length

    expanded_mask = ~(mask[:, None, None, :].to(torch.bool))
    return expanded_mask.expand(batch_size, 1, tgt_length, src_length)


def build_mpt_alibi_tensor(
    num_heads,
    sequence_length,
    sequence_length_with_past,
    alibi_bias_max=8,
    device=None,
    for_ae=False,
    topk=None,
):
    r"""
    Link to paper: https://arxiv.org/abs/2108.12409 - Alibi tensor is not causal as the original paper mentions, it
    relies on a translation invariance of softmax for quick implementation. This implementation has been copied from
    the alibi implementation of MPT source code that led to slightly different results than the Bloom alibi:
    https://huggingface.co/mosaicml/mpt-7b/blob/main/attention.py#L292
    """
    if not for_ae:
        alibi = torch.arange(
            1 - sequence_length, 1, dtype=torch.int32, device=device
        ).view(1, 1, 1, sequence_length)
    else:  # EM: All memory tokens get same bias
        alibi = (
            torch.tensor(-sequence_length_with_past, dtype=torch.int32, device=device)
            .repeat(sequence_length * topk)
            .view(1, 1, 1, sequence_length * topk)
        )
    num_heads_power_of_2 = 2 ** math.ceil(math.log2(num_heads))

    base = torch.arange(1, num_heads_power_of_2 + 1, dtype=torch.float32, device=device)
    base = base * (alibi_bias_max / num_heads_power_of_2)

    slopes = 1.0 / torch.pow(2, base)
    slopes = slopes.view(1, num_heads, 1, 1)

    if num_heads_power_of_2 != num_heads:
        slopes = torch.concat([slopes[1::2], slopes[::2]])[:num_heads]

    alibi = alibi * slopes
    return alibi.squeeze(0)


class ExtendedMptAttention(nn.Module):
    """Multi-head self attention.
    Using torch or triton attention implemetation enables user to also use additive bias.
    """

    def __init__(self, config: ExtendedMptConfig):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.n_heads = config.n_heads
        self.n_layers = config.n_layers
        self.head_dim = self.hidden_size // self.n_heads
        self.softmax_scale = config.attn_config.softmax_scale
        if self.softmax_scale is None:
            self.softmax_scale = 1 / math.sqrt(self.hidden_size / self.n_heads)

        self.attn_dropout_p = config.attn_config.attn_pdrop
        self.Wqkv = nn.Linear(self.hidden_size, 3 * self.hidden_size, bias=False)
        self.out_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_bias: torch.Tensor,
        past_key_value: Optional[Tuple[torch.Tensor]] = None,
        attention_mask: Optional[torch.Tensor] = None,
        long_range_past_key_value=None,
        topk=None,
        faiss_indexes=None,
        mask_by_sim=None,
        sim_threshold=None,
        position_bias_ae=None,
        current_layer=None,
        output_retrieved_memory_idx=False,
    ):
        batch_size, seq_length = hidden_states.shape[:2]

        mixed_qkv = self.Wqkv(hidden_states)
        query_states, key_states, value_states = mixed_qkv.chunk(3, dim=2)
        query_states = query_states.reshape(
            batch_size, seq_length, self.n_heads, self.head_dim
        ).transpose(1, 2)
        key_states = key_states.reshape(
            batch_size, seq_length, self.n_heads, self.head_dim
        ).transpose(1, 2)
        value_states = value_states.reshape(
            batch_size, seq_length, self.n_heads, self.head_dim
        ).transpose(1, 2)

        if past_key_value is not None:
            if len(past_key_value) != 0:
                key_states = torch.cat([past_key_value[0], key_states], dim=2)
                value_states = torch.cat([past_key_value[1], value_states], dim=2)
        past_key_value = (key_states, value_states)
        bsz, nh, s_q, d = query_states.shape

        attention_scores = (
            torch.matmul(query_states, key_states.transpose(-1, -2))
            * self.softmax_scale
        )
        key_length = key_states.shape[-2]
        query_length = (
            seq_length
            if past_key_value is None
            else seq_length + past_key_value[0].shape[2]
        )
        if position_bias is not None:
            if len(position_bias.shape) != 3:
                raise ValueError(
                    f"Expecting position_bias shape to be 3 dimensions, got {len(position_bias.shape)}"
                )

            position_bias_query_index = max(0, position_bias.size(1) - query_length)
            position_bias_key_index = max(0, position_bias.size(2) - key_length)

            position_bias = position_bias[
                :, position_bias_query_index:, position_bias_key_index:
            ]

            attention_scores = attention_scores + position_bias

        # EM: Retrieve memories from cache or faiss indexes
        if long_range_past_key_value is not None or faiss_indexes is not None:
            if long_range_past_key_value is not None:  # Manual store
                k_cache, v_cache = long_range_past_key_value
                s_cache = k_cache.size(-2)

                k_cache = k_cache.to(key_states.device)
                v_cache = v_cache.to(key_states.device)

                # Normalize query and key vectors
                q_n = query_states / vector_norm(
                    query_states, ord=2, dim=-1, keepdim=True
                ) 
                k_n = k_cache / vector_norm(k_cache, ord=2, dim=-1, keepdim=True)
                sim = q_n.matmul(k_n.transpose(-1, -2))
                if s_cache < topk:   # number of tokens in cache < topk
                    topk = s_cache
                val, idx = torch.topk(sim, k=topk, dim=-1)  # Retrieve topk memories

                reshaped_idx = idx.reshape(bsz, nh, s_q * topk)
                selected_k = k_cache.gather(
                    dim=-2, index=reshaped_idx.unsqueeze(-1).expand(-1, -1, -1, d)
                )
                selected_v = v_cache.gather(
                    dim=-2, index=reshaped_idx.unsqueeze(-1).expand(-1, -1, -1, d)
                )

            elif faiss_indexes is not None:  # FAISS indexes
                kn_index, kv_index = faiss_indexes
                q_n = query_states / vector_norm(
                    query_states, ord=2, dim=-1, keepdim=True
                )
                # One-hot encoding for layer, head to only retrieve memories from the same layer, head
                one_hot_encodings = (
                    F.one_hot(
                        torch.arange(0, nh * self.n_layers, device=query_states.device)
                    )
                    * 10
                )
                q_n = torch.concat(
                    [
                        rearrange(q_n, "b h s d -> b (h s) d", h=nh),
                        one_hot_encodings[nh * current_layer : nh * (current_layer + 1)]
                        .unsqueeze(0)
                        .repeat_interleave(repeats=query_states.size(-2), dim=-2),
                    ],
                    dim=-1,
                ).squeeze()

                if kn_index.ntotal / (nh * self.n_layers) < topk:
                    topk = int(kn_index.ntotal / (nh * self.n_layers))

                val, idx = kn_index.search(q_n.to("cpu").detach().numpy(), k=topk)
                val = torch.tensor(val - 100).reshape(bsz, nh, s_q, topk)  #Similarity includes scale factor from one-hot encoding
                reshaped_idx = torch.tensor(
                    idx % (kn_index.ntotal / (nh * self.n_layers))
                ).reshape(bsz, nh, s_q * topk)

                # Retrieve tensors
                selected_k = rearrange(
                    torch.tensor(kv_index.reconstruct_batch(idx.flatten()))[:, :d],
                    "(h s) d -> 1 h s d",
                    h=nh,
                ).to(query_states.device)
                selected_v = rearrange(
                    torch.tensor(kv_index.reconstruct_batch(idx.flatten()))[:, d:],
                    "(h s) d -> 1 h s d",
                    h=nh,
                ).to(query_states.device)

            selected_key_length = selected_k.size(-2)
            key_length += selected_key_length
            attention_scores_cache = (
                query_states.matmul(selected_k.transpose(-1, -2)) * self.softmax_scale
            )
            # EM: Mask by similarity
            if mask_by_sim:
                sim_mask = (
                    rearrange(~(val > sim_threshold).bool(), "b h s i -> b h (s i)")
                    .unsqueeze(-2)
                    .expand(-1, -1, s_q, -1)
                ).to(query_states.device)

                attention_scores_cache = attention_scores_cache.masked_fill(
                    sim_mask, torch.finfo(query_states.dtype).min
                )

            # EM: Add position bias to cache
            if position_bias_ae is not None:
                if len(position_bias_ae.shape) != 3:
                    raise ValueError(
                        f"Expecting position_bias shape to be 3 dimensions, got {len(position_bias_ae.shape)}"
                    )

                position_bias_query_index = max(
                    0, position_bias_ae.size(1) - query_length
                )
                position_bias_key_index = max(
                    0, position_bias_ae.size(2) - selected_key_length
                )

                position_bias_ae = position_bias_ae[
                    :, position_bias_query_index:, position_bias_key_index:
                ]

                attention_scores_cache = attention_scores_cache + position_bias_ae

            # EM: Concatenate cache and current attention weights, values
            attention_scores = torch.cat(
                [attention_scores_cache, attention_scores], dim=-1
            )  # Concat attention scores, values
            value_states = torch.cat([selected_v, value_states], dim=-2)

        # EM: Create mask for external memories, queries only attend to their own memories
        def _create_external_memories_mask(k, s_q, device):
            mask = torch.zeros(s_q, s_q * k, device=device, dtype=torch.bool)
            for i in range(s_q):
                mask[i, i * k : (i + 1) * k] = 1
            return ~mask

        if attention_mask is not None:
            # EM: Concatenate attention mask with external memories mask
            if long_range_past_key_value is not None or faiss_indexes is not None:
                mask = _create_external_memories_mask(
                    k=topk, s_q=s_q, device=attention_scores.device
                )
                attention_mask = attention_mask.squeeze(dim=0).squeeze(dim=0)
                attention_mask = torch.cat([mask, attention_mask], dim=1)
            attention_scores = attention_scores.masked_fill(
                attention_mask, torch.finfo(query_states.dtype).min
            )

        # (batch_size, n_heads, seq_length, key_length)
        attn_weights = nn.functional.softmax(attention_scores.float(), dim=-1).to(
            value_states.dtype
        )
        attn_weights = nn.functional.dropout(
            attn_weights, p=self.attn_dropout_p, training=self.training
        )

        context_states = torch.matmul(attn_weights, value_states)
        context_states = (
            context_states.permute(0, 2, 1, 3)
            .contiguous()
            .view(batch_size, seq_length, -1)
        )
        attn_output = self.out_proj(context_states)

        if not output_retrieved_memory_idx:
            reshaped_idx = None

        return attn_output, attn_weights, past_key_value, reshaped_idx


class MptMLP(nn.Module):
    def __init__(self, config: ExtendedMptConfig):
        super().__init__()
        hidden_size = config.hidden_size

        self.up_proj = nn.Linear(hidden_size, 4 * hidden_size, bias=False)
        self.act = nn.GELU(approximate="none")
        self.down_proj = nn.Linear(4 * hidden_size, hidden_size, bias=False)
        self.hidden_dropout = config.attn_config.attn_pdrop

    def forward(
        self, hidden_states: torch.Tensor, residual: torch.Tensor
    ) -> torch.Tensor:
        hidden_states = self.act(self.up_proj(hidden_states))

        intermediate_output = self.down_proj(hidden_states)

        output = F.dropout(
            intermediate_output, p=self.hidden_dropout, training=self.training
        )
        output = output + residual

        return output


class MptBlock(nn.Module):
    """MPTBlock"""

    def __init__(self, config: ExtendedMptConfig):
        super().__init__()
        hidden_size = config.hidden_size

        self.norm_1 = LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
        # backward compatibility with weights on the Hub
        self.norm_1.bias = None

        self.num_heads = config.n_heads
        self.attn = ExtendedMptAttention(config)

        self.norm_2 = LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
        # backward compatibility with weights on the Hub
        self.norm_2.bias = None

        self.ffn = MptMLP(config)

        self.dropout_rate = config.attn_config.attn_pdrop
        self.resid_attn_dropout = nn.Dropout(self.dropout_rate)

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_bias: torch.Tensor,
        attention_mask: torch.Tensor,
        layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
        use_cache: bool = False,
        output_attentions: bool = False,
        output_retrieved_memory_idx: bool = False,
        topk: int = None,
        long_range_past_key_value: Optional[Tuple[torch.Tensor]] = None,
        faiss_indexes: Tuple = None,
        position_bias_ae=None,
        current_layer: int = None,
        mask_by_sim: bool = False,
        sim_threshold: float = None,
    ):
        # hidden_states: [batch_size, seq_length, hidden_size]
        # Layer norm at the beginning of the transformer layer.
        layernorm_output = self.norm_1(hidden_states)

        residual = hidden_states

        # Self attention.
        attn_outputs, attn_weights, past_key_value, reshaped_idx = self.attn(
            layernorm_output,
            position_bias=position_bias,
            attention_mask=attention_mask,
            past_key_value=layer_past,
            long_range_past_key_value=long_range_past_key_value,
            topk=topk,
            faiss_indexes=faiss_indexes,
            position_bias_ae=position_bias_ae,
            current_layer=current_layer,
            mask_by_sim=mask_by_sim,
            sim_threshold=sim_threshold,
            output_retrieved_memory_idx=output_retrieved_memory_idx,
        )

        hidden_states = self.resid_attn_dropout(attn_outputs) + residual

        layernorm_output = self.norm_2(hidden_states)

        # Get residual
        residual = hidden_states

        # MLP.
        output = self.ffn(layernorm_output, residual)
        outputs = (output,)

        if use_cache:
            outputs += (past_key_value,)

        if output_attentions:
            outputs += (attn_weights,)
        if output_retrieved_memory_idx:
            outputs += (reshaped_idx,)

        return outputs  # hidden_states, present, attentions


class MptPreTrainedModel(PreTrainedModel):
    """MPT Pretrained Model"""

    config_class = ExtendedMptConfig
    base_model_prefix = "transformer"
    supports_gradient_checkpointing = True
    _no_split_modules = ["MptBlock"]
    _keys_to_ignore_on_load_missing = [r"lm_head.*."]

    def __init__(self, *inputs, **kwargs):
        super().__init__(*inputs, **kwargs)

    def _init_weights(self, module: nn.Module):
        """Initialize the weights."""
        if isinstance(module, nn.Linear):
            # Slightly different from the TF version which uses truncated_normal for initialization
            # cf https://github.com/pytorch/pytorch/pull/5617
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()
        elif isinstance(module, LayerNorm):
            if module.bias is not None:
                module.bias.data.zero_()
            module.weight.data.fill_(1.0)

    def _set_gradient_checkpointing(self, module: nn.Module, value: bool = False):
        if isinstance(module, ExtendedMptConfig):
            module.gradient_checkpointing = value

    @staticmethod
    def _convert_to_mpt_cache(
        past_key_value: Tuple[Tuple[torch.Tensor, torch.Tensor]]
    ) -> Tuple[Tuple[torch.Tensor, torch.Tensor]]:
        """
        Converts the cache to the format expected by Mpt, i.e. to tuple(tuple([batch_size * num_heads, ...]))
        """
        batch_size, num_heads, head_dim, seq_length = past_key_value[0][0].shape
        batch_size_times_num_heads = batch_size * num_heads
        # key:  [batch_size, num_heads, head_dim, seq_length] -> [batch_size * num_heads, head_dim, seq_length]
        # value: [batch_size, num_heads, seq_length, head_dim] -> [batch_size * num_heads, seq_length, head_dim]
        return tuple(
            (
                layer_past[0].reshape(batch_size_times_num_heads, head_dim, seq_length),
                layer_past[1].reshape(batch_size_times_num_heads, seq_length, head_dim),
            )
            for layer_past in past_key_value
        )


MPT_START_DOCSTRING = r"""

    This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
    library implements for all its model (such as downloading or saving, resizing the input embeddings etc.)

    This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
    and behavior.

    Parameters:
        config ([`ExtendedMptConfig`]): Model configuration class with all the parameters of the model.
            Initializing with a config file does not load the weights associated with the model, only the
            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
"""

MPT_INPUTS_DOCSTRING = r"""
    Args:
        input_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`):
            `input_ids_length` = `sequence_length` if `past_key_values` is `None` else `past_key_values[0][0].shape[2]`
            (`sequence_length` of input past key value states). Indices of input sequence tokens in the vocabulary.

            If `past_key_values` is used, only `input_ids` that do not have their past calculated should be passed as
            `input_ids`.

            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
            [`PreTrainedTokenizer.__call__`] for details.

            [What are input IDs?](../glossary#input-ids)
        past_key_values (`Tuple[Tuple[torch.Tensor]]` of length `config.n_layers`):
            Contains precomputed hidden-states (key and values in the attention blocks) as computed by the model (see
            `past_key_values` output below). Can be used to speed up sequential decoding. The `input_ids` which have
            their past given to this model should not be passed as `input_ids` as they have already been computed.

            Each element of `past_key_values` is a tuple (past_key, past_value):
            - past_key: [batch_size * num_heads, head_dim, kv_length]
            - past_value: [batch_size * num_heads, kv_length, head_dim]
        attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
            Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:

            - 1 for tokens that are **not masked**,
            - 0 for tokens that are **masked**.

            [What are attention masks?](../glossary#attention-mask)

        inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
            Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
            is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
            model's internal embedding lookup matrix.

            If `past_key_values` is used, optionally only the last `inputs_embeds` have to be input (see
            `past_key_values`).
        use_cache (`bool`, *optional*):
            If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
            `past_key_values`).
        output_attentions (`bool`, *optional*):
            Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
            tensors for more detail.
        output_hidden_states (`bool`, *optional*):
            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
            more detail.
        return_dict (`bool`, *optional*):
            Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple.
        use_external_mind (`bool`, *optional*, defaults to `True`):
            Whether to attend to external memories.
        long_range_past_key_values (`List[Tuple[torch.FloatTensor]]`, *optional*, defaults to None):
            Manual store for memories.
        faiss_indexes (`Tuple[faiss.swigfaiss_avx2.IndexFlatIP]`, *optional*, defaults to None):
            Vector store for memories.
        topk (`int`, *optional*, defaults to `10`):
            Number of external memories for each query token to retrieve and attend to.
"""


@add_start_docstrings(
    "The bare Mpt Model transformer outputting raw hidden-states without any specific head on top.",
    MPT_START_DOCSTRING,
)
class ExtendedMptModel(MptPreTrainedModel):
    """Extended MPT Model"""

    def __init__(self, config: ExtendedMptConfig):
        super().__init__(config)

        self.hidden_size = config.hidden_size
        self.num_heads = config.n_heads

        # Embedding + LN Embedding
        self.wte = nn.Embedding(config.vocab_size, self.hidden_size)

        # Transformer blocks
        self.blocks = nn.ModuleList([MptBlock(config) for _ in range(config.n_layers)])

        # Final Layer Norm
        self.norm_f = LayerNorm(self.hidden_size, eps=config.layer_norm_epsilon)
        # backward compatibility with weights on the Hub
        self.norm_f.bias = None

        self.gradient_checkpointing = False

        # Initialize weights and apply final processing
        self.post_init()

        self.mask_by_sim = config.attn_config.mask_by_sim
        self.sim_threshold = config.attn_config.sim_threshold
        self.topk = config.attn_config.topk
        self.use_external_mind = config.use_external_mind
        self.use_external_mind_by_layer = config.use_external_mind_by_layer

    def get_input_embeddings(self):
        return self.wte

    def build_mpt_alibi_tensor(
        self,
        num_heads,
        sequence_length,
        sequence_length_with_past,
        alibi_bias_max=8,
        device=None,
        for_ae=None,
        topk=None,
    ):
        return build_mpt_alibi_tensor(
            num_heads,
            sequence_length,
            sequence_length_with_past,
            alibi_bias_max,
            device,
            for_ae=for_ae,
            topk=topk,
        )

    def _prepare_attn_mask(
        self,
        attention_mask: torch.Tensor,
        input_shape: Tuple[int, int],
        past_key_values_length: int,
    ) -> torch.BoolTensor:
        # create causal mask
        # [batch_size, seq_length] -> [batch_size, 1, tgt_length, src_length]
        if input_shape[1] + past_key_values_length != attention_mask.shape[1]:
            raise ValueError(
                "Attention mask shape should be (batch_size, seq_length + past_key_values_length)"
                f" but is {attention_mask.shape} with input_ids shape {input_shape} and past length"
                f" {past_key_values_length}."
            )
        combined_attention_mask = None
        device = attention_mask.device
        _, src_length = input_shape

        if src_length > 1:
            combined_attention_mask = _make_causal_mask(
                input_shape,
                device=device,
                past_key_values_length=past_key_values_length,
            )

        # [batch_size, seq_length] -> [batch_size, 1, tgt_length, src_length]
        expanded_attn_mask = _expand_mask(attention_mask, tgt_length=src_length)
        combined_attention_mask = (
            expanded_attn_mask
            if combined_attention_mask is None
            else expanded_attn_mask | combined_attention_mask
        )

        return combined_attention_mask

    def set_input_embeddings(self, new_embeddings: torch.Tensor):
        self.wte = new_embeddings

    @add_start_docstrings_to_model_forward(MPT_INPUTS_DOCSTRING)
    @add_code_sample_docstrings(
        checkpoint=_CHECKPOINT_FOR_DOC,
        output_type=BaseModelOutputWithPastAndCrossAttentions,
        config_class=_CONFIG_FOR_DOC,
    )
    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,
        attention_mask: Optional[torch.Tensor] = None,
        inputs_embeds: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        output_retrieved_memory_idx: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        use_external_mind: Optional[bool] = None,
        long_range_past_key_values: Optional[list[Tuple[torch.FloatTensor]]] = None,
        faiss_indexes: Tuple = None,
        topk: int = None,
    ) -> Union[Tuple[torch.Tensor, ...], BaseModelOutputWithPastAndCrossAttentions]:
        output_attentions = (
            output_attentions
            if output_attentions is not None
            else self.config.output_attentions
        )
        output_retrieved_memory_idx = (
            output_retrieved_memory_idx
            if output_retrieved_memory_idx is not None
            else False
        )
        output_hidden_states = (
            output_hidden_states
            if output_hidden_states is not None
            else self.config.output_hidden_states
        )
        use_cache = use_cache if use_cache is not None else self.config.use_cache
        return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )
        use_external_mind = (
            use_external_mind
            if use_external_mind is not None
            else self.use_external_mind
        )
        topk = topk if topk is not None else self.topk

        if input_ids is not None and inputs_embeds is not None:
            raise ValueError(
                "You cannot specify both input_ids and inputs_embeds at the same time"
            )
        elif input_ids is not None:
            batch_size, seq_length = input_ids.shape
        elif inputs_embeds is not None:
            batch_size, seq_length, _ = inputs_embeds.shape
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds")

        if past_key_values is None:
            past_key_values = tuple([None] * len(self.blocks))

        if inputs_embeds is None:
            inputs_embeds = self.wte(input_ids)

        hidden_states = inputs_embeds

        presents = () if use_cache else None
        all_self_attentions = () if output_attentions else None
        all_hidden_states = () if output_hidden_states else None
        all_idx = () if output_retrieved_memory_idx else None

        if self.gradient_checkpointing and self.training:
            if use_cache:
                logger.warning_once(
                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
                )
                use_cache = False

        # Compute alibi tensor: check build_alibi_tensor documentation
        seq_length_with_past = seq_length
        past_key_values_length = 0
        if past_key_values[0] is not None:
            past_key_values_length = past_key_values[0][0].shape[2]
            seq_length_with_past = seq_length_with_past + past_key_values_length
        if attention_mask is None:
            attention_mask = torch.ones(
                (batch_size, seq_length_with_past), device=hidden_states.device
            )
        else:
            attention_mask = attention_mask.to(hidden_states.device)

        alibi = self.build_mpt_alibi_tensor(
            self.num_heads,
            self.config.max_seq_len,
            seq_length_with_past,
            device=hidden_states.device,
        )
        # EM: Alibi tensor for retrieved kvs
        alibi_ae = self.build_mpt_alibi_tensor(
            self.num_heads,
            seq_length,
            seq_length_with_past,
            device=hidden_states.device,
            for_ae=True,
            topk=topk,
        )

        causal_mask = self._prepare_attn_mask(
            attention_mask,
            input_shape=(batch_size, seq_length),
            past_key_values_length=past_key_values_length,
        )

        for i, (block, layer_past) in enumerate(zip(self.blocks, past_key_values)):
            if output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)

            long_range_past_key_value = (
                long_range_past_key_values[i]
                if (
                    long_range_past_key_values is not None
                    and self.use_external_mind_by_layer[i]
                    and use_external_mind is True
                )
                else None
            )
            if long_range_past_key_value is not None and faiss_indexes is not None:
                raise NotImplementedError(
                    """Using faiss and passing key value pairs
                    manually are mutually exclusive right now."""
                )
            if self.gradient_checkpointing and self.training:

                def create_custom_forward(module):
                    def custom_forward(*inputs):
                        # None for past_key_value
                        return module(
                            *inputs,
                            use_cache=use_cache,
                            output_attentions=output_attentions,
                        )

                    return custom_forward

                outputs = torch.utils.checkpoint.checkpoint(
                    create_custom_forward(block),
                    hidden_states,
                    alibi,
                    causal_mask,
                    layer_past,
                )
            else:
                outputs = block(
                    hidden_states,
                    layer_past=layer_past,
                    attention_mask=causal_mask,
                    use_cache=use_cache,
                    output_attentions=output_attentions,
                    output_retrieved_memory_idx=output_retrieved_memory_idx,
                    position_bias=alibi,
                    position_bias_ae=alibi_ae,
                    topk=topk,
                    long_range_past_key_value=long_range_past_key_value,
                    faiss_indexes=faiss_indexes,
                    mask_by_sim=self.mask_by_sim,
                    sim_threshold=self.sim_threshold,
                    current_layer=i,
                )

            hidden_states = outputs[0]
            if use_cache is True:
                presents = presents + (outputs[1],)

            if output_attentions:
                all_self_attentions = all_self_attentions + (
                    outputs[2 if use_cache else 1],
                )
            if output_retrieved_memory_idx:
                idx = (
                    3
                    if (use_cache & output_attentions)
                    else 2
                    if (use_cache or output_attentions)
                    else 1
                )
                all_idx = all_idx + (outputs[idx],)

        # Add last hidden state
        hidden_states = self.norm_f(hidden_states)

        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

        if not return_dict:
            return tuple(
                v
                for v in [
                    hidden_states,
                    presents,
                    all_hidden_states,
                    all_self_attentions,
                    all_idx,
                ]
                if v is not None
            )

        return BaseModelOutputWithPastAndCrossAttentions(
            last_hidden_state=hidden_states,
            past_key_values=presents,
            hidden_states=all_hidden_states,
            attentions=(all_self_attentions, all_idx),  # EM: Return idx of retrieved memories
        )


@add_start_docstrings(
    """
    The MPT Model transformer with a language modeling head on top (linear layer with weights tied to the input
    embeddings).
    """,
    MPT_START_DOCSTRING,
)
class ExtendedMptForCausalLM(MptPreTrainedModel):
    """Extended MPT for Causal LM."""

    _tied_weights_keys = ["lm_head.weight"]

    def __init__(self, config: ExtendedMptConfig, external_memories=None):
        super().__init__(config)
        self.transformer: ExtendedMptModel = ExtendedMptModel(config)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

        self.use_external_mind = config.use_external_mind
        self.memory_type = config.attn_config.memory_type
        self.memory_ids = None
        self.memories = None
        self.memory_device = config.attn_config.memory_device
        self.remove_special_ids = config.attn_config.remove_special_ids
        self.tokenizer_all_special_ids = config.attn_config.tokenizer_all_special_ids

        # EM: Memory token ids
        if external_memories is not None:
            self.memory_ids = external_memories
        # Initialize weights and apply final processing
        self.post_init()

    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, new_embeddings: torch.Tensor):
        self.lm_head = new_embeddings

    # EM: Clear memory cache
    def clear_memory(self):
        """Clear memory cache."""
        self.memory_ids = None
        self.memories = None

    def prepare_inputs_for_generation(
        self,
        input_ids: torch.LongTensor,
        past_key_values: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        inputs_embeds: Optional[torch.Tensor] = None,
        use_cache: Optional[bool] = None,
        **kwargs,
    ) -> dict:
        # only last token for input_ids if past is not None
        if past_key_values:
            input_ids = input_ids[:, -1].unsqueeze(-1)

        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
        if inputs_embeds is not None and past_key_values is None:
            model_inputs = {"inputs_embeds": inputs_embeds}
        else:
            model_inputs = {"input_ids": input_ids}

        model_inputs.update(
            {
                "past_key_values": past_key_values,  # NITS should it be layer_past?
                "use_cache": use_cache,
                "attention_mask": attention_mask,
                "use_external_mind": kwargs.get("use_external_mind"), # EM: Add config here
                "topk": kwargs.get("topk"),
            }
        )
        return model_inputs

    @add_start_docstrings_to_model_forward(MPT_INPUTS_DOCSTRING)
    @add_code_sample_docstrings(
        checkpoint=_CHECKPOINT_FOR_DOC,
        output_type=CausalLMOutputWithCrossAttentions,
        config_class=_CONFIG_FOR_DOC,
    )
    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,
        attention_mask: Optional[torch.Tensor] = None,
        inputs_embeds: Optional[torch.Tensor] = None,
        labels: Optional[torch.Tensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_retrieved_memory_idx: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        use_external_mind: Optional[bool] = None,
        topk: int = None,
    ) -> Union[Tuple[torch.Tensor], CausalLMOutputWithCrossAttentions]:
        r"""
        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
            Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
            `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
            are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
        """
        return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )

        # EM: Generate key value cache once on first call
        if (
            self.memory_ids is not None and self.memories is None
        ): 
            self.memories = self.generate_cache(
                self.memory_ids, cache_type=self.memory_type
            )
            # EM: Remove special tokens from memory cache
            if self.remove_special_ids:
                idx_to_remove = [
                    token_idx
                    for token_idx, token in enumerate(self.memory_ids[0])
                    if token in self.tokenizer_all_special_ids
                ]
                if self.memory_type == "manual":
                    mask = torch.ones(self.memories[0][0].size(), dtype=torch.bool)
                    mask[:, :, idx_to_remove, :] = False

                    new_size = (
                        self.memories[0][0].size(0),
                        self.memories[0][0].size(1),
                        -1,
                        self.memories[0][0].size(3),
                    )
                    self.memories = [
                        (ks[mask].view(new_size), vs[mask].view(new_size))
                        for ks, vs in self.memories
                    ]
                else:
                    kn_index, kv_index = self.memories
                    all_idx_to_remove = [
                        [
                            i
                            for i in range(0, kn_index.ntotal)
                            if (
                                i
                                % (
                                    kn_index.ntotal
                                    / (
                                        self.config.num_attention_heads
                                        * self.config.num_hidden_layers
                                    )
                                )
                            )
                            == j
                        ]
                        for j in idx_to_remove
                    ]
                    kn_index.remove_ids(
                        np.array(all_idx_to_remove).flatten().astype("int64")
                    )
                    kv_index.remove_ids(
                        np.array(all_idx_to_remove).flatten().astype("int64")
                    )

        use_external_mind = (
            use_external_mind
            if use_external_mind is not None
            else self.use_external_mind
        )
        topk = topk if topk is not None else None

        long_range_past_key_values = None
        faiss_indexes = None
        if hasattr(self, "memories") and isinstance(self.memories, list):
            long_range_past_key_values = self.memories
        elif hasattr(self, "memories"):
            faiss_indexes = self.memories

        transformer_outputs = self.transformer(
            input_ids,
            past_key_values=past_key_values,
            attention_mask=attention_mask,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_retrieved_memory_idx=output_retrieved_memory_idx,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            long_range_past_key_values=long_range_past_key_values,
            faiss_indexes=faiss_indexes,
            use_external_mind=use_external_mind,
            topk=topk,
        )
        hidden_states = transformer_outputs[0]

        lm_logits = self.lm_head(hidden_states)

        loss = None
        if labels is not None:
            # move labels to correct device to enable model parallelism
            labels = labels.to(lm_logits.device)
            # Shift so that tokens < n predict n
            shift_logits = lm_logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            batch_size, seq_length, vocab_size = shift_logits.shape
            # Flatten the tokens
            loss_fct = CrossEntropyLoss()
            loss = loss_fct(
                shift_logits.view(batch_size * seq_length, vocab_size),
                shift_labels.view(batch_size * seq_length),
            )

        if not return_dict:
            output = (lm_logits,) + transformer_outputs[1:]
            return ((loss,) + output) if loss is not None else output

        return CausalLMOutputWithCrossAttentions(
            loss=loss,
            logits=lm_logits,
            past_key_values=transformer_outputs.past_key_values,
            hidden_states=transformer_outputs.hidden_states,
            attentions=transformer_outputs.attentions,
        )

    def _reorder_cache(
        self,
        past: Tuple[Tuple[torch.Tensor, torch.Tensor], ...],
        beam_idx: torch.LongTensor,
    ) -> Tuple[Tuple[torch.Tensor, 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.

        Output shares the same memory storage as `past`.
        """
        # Get a copy of `beam_idx` on all the devices where we need those indices.
        device_to_beam_idx = {
            past_state.device: beam_idx.to(past_state.device)
            for layer_past in past
            for past_state in layer_past
        }
        reordered_past = tuple(
            (
                layer_past[0].index_select(0, device_to_beam_idx[layer_past[0].device]),
                layer_past[1].index_select(0, device_to_beam_idx[layer_past[0].device]),
            )
            for layer_past in past
        )
        return reordered_past

    # EM: Add method to generate key-value cache
    def generate_cache(
        self,
        input_ids: torch.LongTensor,
        stride: int = 512,
        max_len: int = 3072,
        cache_type: str = "manual",
    ):
        """Generate cache for long range attention."""
        if cache_type not in ["manual", "faiss"]:
            raise NotImplementedError(f"Cache type {cache_type} not implemented.")

        prev_end_loc = 0
        long_range_past_key_values = None
        faiss_indexes = None
        for b_idx in range(
            0, input_ids.size(-1), stride
        ):  # generate kv-pairs using stride
            end_loc = min(b_idx + max_len, input_ids.size(-1))
            trg_len = end_loc - prev_end_loc
            subseq = input_ids[:, b_idx:end_loc].to(self.device)
            with torch.no_grad():
                outputs = self.transformer(
                    subseq, use_cache=True, use_external_mind=False
                )
            to_cache = [
                (kv[0][:, :, -trg_len:], kv[1][:, :, -trg_len:])
                for kv in outputs.past_key_values
            ]
            long_range_past_key_values, faiss_indexes = self.cache(
                to_cache,
                cache_type,
                long_range_past_key_values=long_range_past_key_values,
                faiss_indexes=faiss_indexes,
            )

            prev_end_loc = end_loc
            if end_loc == input_ids.size(-1):
                break
        if long_range_past_key_values is not None:
            return long_range_past_key_values
        else:
            return faiss_indexes
        
    # EM: Add method to cache key value pairs
    def cache(
        self,
        to_cache: list,
        cache_type: str = "manual",
        long_range_past_key_values: list = None,
        faiss_indexes: faiss.IndexFlatIP = None,
        max_length_cache=100000,
        verbose=False,
    ):
        """Cache long range attention."""
        if (long_range_past_key_values is not None) & (faiss_indexes is not None):
            raise NotImplementedError(
                "Using faiss and passing key value pairs manually are mutually exclusive right now."
            )
        
        # To avoid spinning up a new index for each layer, we add one-hot encodings to the keys so that queries match with the appropriate layer, head
        if cache_type == "faiss":  # add one-hot encoding to match layer, head indices
            one_hot_encodings = (
                F.one_hot(torch.arange(0, self.config.n_heads * self.config.n_layers))
                * 10
            )
            # New indices, one to store normalized keys with one-hot encodings, another to retrieve kv pairs without normalization
            if faiss_indexes is None:
                faiss_indexes = (
                    faiss.IndexFlatIP(
                        to_cache[0][0].size(-1) + one_hot_encodings.size(-1)
                    ),
                    faiss.IndexFlatIP(to_cache[0][0].size(-1) * 2),
                )
            kn_index, kv_index = faiss_indexes
            for l_idx, (k, v) in enumerate(to_cache):
                k_n = (k / vector_norm(k, ord=2, dim=-1, keepdim=True)).to("cpu") #Normalize keys for cosine sim
                
                # Indices are 2 dimensional, so flatten 
                # Add normalized keys with one-hot encodings
                k_n = torch.concat(
                    [
                        rearrange(k_n, "b h s d -> b (h s) d", h=self.config.n_heads),
                        one_hot_encodings[
                            self.config.n_heads
                            * l_idx : self.config.n_heads
                            * (l_idx + 1)
                        ]
                        .unsqueeze(0)
                        .repeat_interleave(repeats=k.size(-2), dim=-2),
                    ],
                    dim=-1,
                )
                kn_index.add(k_n.squeeze().numpy())

                # Add unnormalized keys and values
                k = rearrange(k, "b h s d -> b (h s) d", h=self.config.n_heads)
                v = rearrange(v, "b h s d -> b (h s) d", h=self.config.n_heads)
                kv_index.add(
                    torch.concat([k.squeeze(), v.squeeze()], dim=1).to("cpu").numpy()
                )
        else:
            # Simply use list to store key value pairs
            if long_range_past_key_values is None:
                long_range_past_key_values = [
                    (k.to(self.memory_device), v.to(self.memory_device))
                    for k, v in to_cache
                ]
            else:
                long_range_past_key_values = [
                    (
                        torch.concat(
                            [kv[0], to_cache[ind][0].to(self.memory_device)], dim=2
                        ),
                        torch.concat(
                            [kv[1], to_cache[ind][1].to(self.memory_device)], dim=2
                        ),
                    )
                    for ind, kv in enumerate(long_range_past_key_values)
                ]
        if (
            long_range_past_key_values is not None
        ):  # set a limit on manual memory length
            if long_range_past_key_values[0][0].size(-2) > max_length_cache:
                long_range_past_key_values = [
                    (
                        kv[0][:, :, -max_length_cache:],
                        kv[1][:, :, -max_length_cache:],
                    )
                    for kv in long_range_past_key_values
                ]
        if verbose:
            if cache_type == "faiss":
                print(f"{kn_index.ntotal} keys in faiss index")
            else:
                print(f"{long_range_past_key_values[0][0].size(-2)} cached kvs")

        return (
            long_range_past_key_values,
            (kn_index, kv_index) if cache_type == "faiss" else None,
        )