reformer

mindnlp.transformers.models.reformer.configuration_reformer

Reformer model configuration

mindnlp.transformers.models.reformer.configuration_reformer.ReformerConfig

Bases: PretrainedConfig

This is the configuration class to store the configuration of a [ReformerModel]. It is used to instantiate a Reformer model according to the specified arguments, defining the model architecture. Instantiating a configuration with the defaults will yield a similar configuration to that of the ReFormer google/reformer-crime-and-punishment architecture.

Configuration objects inherit from [PretrainedConfig] and can be used to control the model outputs. Read the documentation from [PretrainedConfig] for more information.

PARAMETER	DESCRIPTION
attention_head_size	Dimensionality of the projected key, query and value vectors TYPE: `int`, *optional*, defaults to 64 DEFAULT: 64
attn_layers	List of attention layer types in ascending order. It can be chosen between a LSHSelfAttention layer ("lsh") and a LocalSelfAttention layer ("local"). For more information on LSHSelfAttention layer, see LSH Self Attention. For more information on LocalSelfAttention layer, see Local Self Attention. TYPE: `List[str]`, *optional*, defaults to `["local", "lsh", "local", "lsh", "local", "lsh"]` DEFAULT: ['local', 'lsh', 'local', 'lsh', 'local', 'lsh']
axial_pos_embds	Whether or not to use axial position embeddings. For more information on how axial position embeddings work, see Axial Position Encodings. TYPE: `bool`, *optional*, defaults to `True` DEFAULT: True
axial_norm_std	The standard deviation of the normal_initializer for initializing the weight matrices of the axial positional encodings. TYPE: `float`, *optional*, defaults to 1.0 DEFAULT: 1.0
axial_pos_shape	The position dims of the axial position encodings. During training, the product of the position dims has to be equal to the sequence length. For more information on how axial position embeddings work, see Axial Position Encodings. TYPE: `List[int]`, *optional*, defaults to `[64, 64]` DEFAULT: [64, 64]
axial_pos_embds_dim	The embedding dims of the axial position encodings. The sum of the embedding dims has to be equal to the hidden size. For more information on how axial position embeddings work, see Axial Position Encodings. TYPE: `List[int]`, *optional*, defaults to `[64, 192]` DEFAULT: [64, 192]
chunk_size_lm_head	The chunk size of the final language model feed forward head layer. A chunk size of 0 means that the feed forward layer is not chunked. A chunk size of n means that the feed forward layer processes n < sequence_length embeddings at a time. For more information on feed forward chunking, see How does Feed Forward Chunking work?. TYPE: `int`, *optional*, defaults to 0 DEFAULT: 0
eos_token_id	The token id for the end-of-sentence token. TYPE: `int`, *optional*, defaults to 2 DEFAULT: 2
feed_forward_size	Dimensionality of the feed_forward layer in the residual attention block. TYPE: `int`, *optional*, defaults to 512 DEFAULT: 512
hash_seed	Seed that can be used to make local sensitive hashing in LSHSelfAttention deterministic. This should only be set for testing purposed. For evaluation and training purposes hash_seed should be left as None to ensure fully random rotations in local sensitive hashing scheme. TYPE: `int`, *optional* DEFAULT: None
hidden_act	The non-linear activation function (function or string) in the feed forward layer in the residual attention block. If string, "gelu", "relu", "silu" and "gelu_new" are supported. TYPE: `str` or `Callable`, *optional*, defaults to `"relu"` DEFAULT: 'relu'
hidden_dropout_prob	The dropout probability for all fully connected layers in the embeddings, encoder, and pooler. TYPE: `float`, *optional*, defaults to 0.05 DEFAULT: 0.05
hidden_size	Dimensionality of the output hidden states of the residual attention blocks. TYPE: `int`, *optional*, defaults to 256 DEFAULT: 256
initializer_range	The standard deviation of the truncated_normal_initializer for initializing all weight matrices. TYPE: `float`, *optional*, defaults to 0.02 DEFAULT: 0.02
is_decoder	Whether or not to use a causal mask in addition to the attention_mask passed to [ReformerModel]. When using the Reformer for causal language modeling, this argument should be set to True. TYPE: `bool`, *optional*, defaults to `False` DEFAULT: False
layer_norm_eps	The epsilon used by the layer normalization layers. TYPE: `float`, *optional*, defaults to 1e-12 DEFAULT: 1e-12
local_chunk_length	Length of chunk which attends to itself in LocalSelfAttention. Chunking reduces memory complexity from sequence length x sequence length (self attention) to chunk length x chunk length x sequence length / chunk length (chunked self attention). TYPE: `int`, *optional*, defaults to 64
local_num_chunks_before	Number of previous neighbouring chunks to attend to in LocalSelfAttention layer to itself. TYPE: `int`, *optional*, defaults to 1 DEFAULT: 1
local_num_chunks_after	Number of following neighbouring chunks to attend to in LocalSelfAttention layer in addition to itself. TYPE: `int`, *optional*, defaults to 0 DEFAULT: 0
local_attention_probs_dropout_prob	The dropout ratio for the attention probabilities in LocalSelfAttention. TYPE: `float`, *optional*, defaults to 0.1 DEFAULT: 0.05
lsh_attn_chunk_length	Length of chunk which attends to itself in LSHSelfAttention. Chunking reduces memory complexity from sequence length x sequence length (self attention) to chunk length x chunk length x sequence length / chunk length (chunked self attention). TYPE: `int`, *optional*, defaults to 64 DEFAULT: 64
lsh_num_chunks_before	Number of previous neighbouring chunks to attend to in LSHSelfAttention layer to itself. TYPE: `int`, *optional*, defaults to 1 DEFAULT: 1
lsh_num_chunks_after	Number of following neighbouring chunks to attend to in LSHSelfAttention layer to itself. TYPE: `int`, *optional*, defaults to 0 DEFAULT: 0
lsh_attention_probs_dropout_prob	The dropout ratio for the attention probabilities in LSHSelfAttention. TYPE: `float`, *optional*, defaults to 0.1 DEFAULT: 0.0
max_position_embeddings	The maximum sequence length that this model might ever be used with. Typically set this to something large just in case (e.g., 512 or 1024 or 2048). TYPE: `int`, *optional*, defaults to 4096 DEFAULT: 4096
num_attention_heads	Number of attention heads for each attention layer in the Transformer encoder. TYPE: `int`, *optional*, defaults to 12 DEFAULT: 12
num_buckets	Number of buckets, the key query vectors can be "hashed into" using the locality sensitive hashing scheme. Each query key vector is hashed into a hash in 1, ..., num_buckets. The number of buckets can also be factorized into a list for improved memory complexity. In this case, each query key vector is hashed into a hash in 1-1, 1-2, ..., num_buckets[0]-1, ..., num_buckets[0]-num_buckets[1] if num_buckets is factorized into two factors. The number of buckets (or the product the factors) should approximately equal sequence length / lsh_chunk_length. If num_buckets not set, a good value is calculated on the fly. TYPE: `int` or `List[int]`, *optional* DEFAULT: None
num_hashes	Number of hashing rounds (e.g., number of random rotations) in Local Sensitive Hashing scheme. The higher num_hashes, the more accurate the LSHSelfAttention becomes, but also the more memory and time intensive the hashing becomes. TYPE: `int`, *optional*, defaults to 1 DEFAULT: 1
pad_token_id	The token id for the padding token. TYPE: `int`, *optional*, defaults to 0 DEFAULT: 0
vocab_size	\ Vocabulary size of the Reformer model. Defines the number of different tokens that can be represented by the inputs_ids passed when calling [ReformerModel]. TYPE: `int`, *optional*, defaults to 320 DEFAULT: 320
tie_word_embeddings	Whether to tie input and output embeddings. TYPE: `bool`, *optional*, defaults to `False` DEFAULT: False
use_cache	Whether or not the model should return the last key/values attentions (not used by all models). TYPE: `bool`, *optional*, defaults to `True` DEFAULT: True
classifier_dropout	The dropout ratio for the classification head. TYPE: `float`, *optional* DEFAULT: None

Example

>>> from transformers import ReformerConfig, ReformerModel
...
>>> # Initializing a Reformer configuration
>>> configuration = ReformerConfig()
...
>>> # Initializing a Reformer model (with random weights)
>>> model = ReformerModel(configuration)
...
>>> # Accessing the model configuration
>>> configuration = model.config

Source code in mindnlp/transformers/models/reformer/configuration_reformer.py

class ReformerConfig(PretrainedConfig):
    r"""
    This is the configuration class to store the configuration of a [`ReformerModel`]. It is used to instantiate a
    Reformer model according to the specified arguments, defining the model architecture. Instantiating a configuration
    with the defaults will yield a similar configuration to that of the ReFormer
    [google/reformer-crime-and-punishment](https://hf-mirror.com/google/reformer-crime-and-punishment) architecture.

    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
    documentation from [`PretrainedConfig`] for more information.

    Args:
        attention_head_size (`int`, *optional*, defaults to 64):
            Dimensionality of the projected key, query and value vectors
        attn_layers (`List[str]`, *optional*, defaults to `["local", "lsh", "local", "lsh", "local", "lsh"]`):
            List of attention layer types in ascending order. It can be chosen between a LSHSelfAttention layer
            (`"lsh"`) and a LocalSelfAttention layer (`"local"`).

            For more information on LSHSelfAttention layer, see [LSH Self Attention](reformer#lsh-self-attention). For
            more information on LocalSelfAttention layer, see [Local Self Attention](reformer#local-self-attention).
        axial_pos_embds (`bool`, *optional*, defaults to `True`):
            Whether or not to use axial position embeddings. For more information on how axial position embeddings
            work, see [Axial Position Encodings](reformer#axial-positional-encodings).
        axial_norm_std (`float`, *optional*, defaults to 1.0):
            The standard deviation of the normal_initializer for initializing the weight matrices of the axial
            positional encodings.
        axial_pos_shape (`List[int]`, *optional*, defaults to `[64, 64]`):
            The position dims of the axial position encodings. During training, the product of the position dims has to
            be equal to the sequence length.

            For more information on how axial position embeddings work, see [Axial Position
            Encodings](reformer#axial-positional-encodings).
        axial_pos_embds_dim (`List[int]`, *optional*, defaults to `[64, 192]`):
            The embedding dims of the axial position encodings. The sum of the embedding dims has to be equal to the
            hidden size.

            For more information on how axial position embeddings work, see [Axial Position
            Encodings](reformer#axial-positional-encodings).
        chunk_size_lm_head (`int`, *optional*, defaults to 0):
            The chunk size of the final language model feed forward head layer. A chunk size of 0 means that the feed
            forward layer is not chunked. A chunk size of n means that the feed forward layer processes n <
            sequence_length embeddings at a time.

            For more information on feed forward chunking, see [How does Feed Forward Chunking
            work?](../glossary#feed-forward-chunking).
        eos_token_id (`int`, *optional*, defaults to 2):
            The token id for the end-of-sentence token.
        feed_forward_size (`int`, *optional*, defaults to 512):
            Dimensionality of the feed_forward layer in the residual attention block.
        hash_seed (`int`, *optional*):
            Seed that can be used to make local sensitive hashing in `LSHSelfAttention` deterministic. This should only
            be set for testing purposed. For evaluation and training purposes `hash_seed` should be left as `None` to
            ensure fully random rotations in local sensitive hashing scheme.
        hidden_act (`str` or `Callable`, *optional*, defaults to `"relu"`):
            The non-linear activation function (function or string) in the feed forward layer in the residual attention
            block. If string, `"gelu"`, `"relu"`, `"silu"` and `"gelu_new"` are supported.
        hidden_dropout_prob (`float`, *optional*, defaults to 0.05):
            The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
        hidden_size (`int`, *optional*, defaults to 256):
            Dimensionality of the output hidden states of the residual attention blocks.
        initializer_range (`float`, *optional*, defaults to 0.02):
            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
        is_decoder (`bool`, *optional*, defaults to `False`):
            Whether or not to use a causal mask in addition to the `attention_mask` passed to [`ReformerModel`]. When
            using the Reformer for causal language modeling, this argument should be set to `True`.
        layer_norm_eps (`float`, *optional*, defaults to 1e-12):
            The epsilon used by the layer normalization layers.
        local_chunk_length (`int`, *optional*, defaults to 64):
            Length of chunk which attends to itself in `LocalSelfAttention`. Chunking reduces memory complexity from
            sequence length x sequence length (self attention) to chunk length x chunk length x sequence length / chunk
            length (chunked self attention).
        local_num_chunks_before (`int`, *optional*, defaults to 1):
            Number of previous neighbouring chunks to attend to in `LocalSelfAttention` layer to itself.
        local_num_chunks_after (`int`, *optional*, defaults to 0):
            Number of following neighbouring chunks to attend to in `LocalSelfAttention` layer in addition to itself.
        local_attention_probs_dropout_prob (`float`, *optional*, defaults to 0.1):
            The dropout ratio for the attention probabilities in `LocalSelfAttention`.
        lsh_attn_chunk_length (`int`, *optional*, defaults to 64):
            Length of chunk which attends to itself in `LSHSelfAttention`. Chunking reduces memory complexity from
            sequence length x sequence length (self attention) to chunk length x chunk length x sequence length / chunk
            length (chunked self attention).
        lsh_num_chunks_before (`int`, *optional*, defaults to 1):
            Number of previous neighbouring chunks to attend to in `LSHSelfAttention` layer to itself.
        lsh_num_chunks_after (`int`, *optional*, defaults to 0):
            Number of following neighbouring chunks to attend to in `LSHSelfAttention` layer to itself.
        lsh_attention_probs_dropout_prob (`float`, *optional*, defaults to 0.1):
            The dropout ratio for the attention probabilities in `LSHSelfAttention`.
        max_position_embeddings (`int`, *optional*, defaults to 4096):
            The maximum sequence length that this model might ever be used with. Typically set this to something large
            just in case (e.g., 512 or 1024 or 2048).
        num_attention_heads (`int`, *optional*, defaults to 12):
            Number of attention heads for each attention layer in the Transformer encoder.
        num_buckets (`int` or `List[int]`, *optional*):
            Number of buckets, the key query vectors can be "hashed into" using the locality sensitive hashing scheme.
            Each query key vector is hashed into a hash in `1, ..., num_buckets`. The number of buckets can also be
            factorized into a list for improved memory complexity. In this case, each query key vector is hashed into a
            hash in `1-1, 1-2, ..., num_buckets[0]-1, ..., num_buckets[0]-num_buckets[1]` if `num_buckets` is
            factorized into two factors. The number of buckets (or the product the factors) should approximately equal
            sequence length / lsh_chunk_length. If `num_buckets` not set, a good value is calculated on the fly.
        num_hashes (`int`, *optional*, defaults to 1):
            Number of hashing rounds (e.g., number of random rotations) in Local Sensitive Hashing scheme. The higher
            `num_hashes`, the more accurate the `LSHSelfAttention` becomes, but also the more memory and time intensive
            the hashing becomes.
        pad_token_id (`int`, *optional*, defaults to 0):
            The token id for the padding token.
        vocab_size (`int`, *optional*, defaults to 320):\
            Vocabulary size of the Reformer model. Defines the number of different tokens that can be represented by
            the `inputs_ids` passed when calling [`ReformerModel`].
        tie_word_embeddings (`bool`, *optional*, defaults to `False`):
            Whether to tie input and output embeddings.
        use_cache (`bool`, *optional*, defaults to `True`):
            Whether or not the model should return the last key/values attentions (not used by all models).
        classifier_dropout (`float`, *optional*):
            The dropout ratio for the classification head.

    Example:
        ```python
        >>> from transformers import ReformerConfig, ReformerModel
        ...
        >>> # Initializing a Reformer configuration
        >>> configuration = ReformerConfig()
        ...
        >>> # Initializing a Reformer model (with random weights)
        >>> model = ReformerModel(configuration)
        ...
        >>> # Accessing the model configuration
        >>> configuration = model.config
        ```
"""
    model_type = "reformer"
    keys_to_ignore_at_inference = ["past_buckets_states"]
    attribute_map = {}

    def __init__(
        self,
        attention_head_size=64,
        attn_layers=["local", "lsh", "local", "lsh", "local", "lsh"],
        axial_norm_std=1.0,
        axial_pos_embds=True,
        axial_pos_shape=[64, 64],
        axial_pos_embds_dim=[64, 192],
        chunk_size_lm_head=0,
        eos_token_id=2,
        feed_forward_size=512,
        hash_seed=None,
        hidden_act="relu",
        hidden_dropout_prob=0.05,
        hidden_size=256,
        initializer_range=0.02,
        is_decoder=False,
        layer_norm_eps=1e-12,
        local_num_chunks_before=1,
        local_num_chunks_after=0,
        local_attention_probs_dropout_prob=0.05,
        local_attn_chunk_length=64,
        lsh_attn_chunk_length=64,
        lsh_attention_probs_dropout_prob=0.0,
        lsh_num_chunks_before=1,
        lsh_num_chunks_after=0,
        max_position_embeddings=4096,
        num_attention_heads=12,
        num_buckets=None,
        num_hashes=1,
        pad_token_id=0,
        vocab_size=320,
        tie_word_embeddings=False,
        use_cache=True,
        classifier_dropout=None,
        **kwargs,
    ):
        """
        Initializes a new instance of the ReformerConfig class.

        Args:
            attention_head_size (int): The size of each attention head.
            attn_layers (list): The list of attention layer types to be used.
            axial_norm_std (float): Standard deviation for axial positional embeddings normalization.
            axial_pos_embds (bool): Whether to use axial positional embeddings.
            axial_pos_shape (list): The shape of axial positional embeddings.
            axial_pos_embds_dim (list): The dimensions of axial positional embeddings.
            chunk_size_lm_head (int): Size of chunk for the language model head.
            eos_token_id (int): The token ID for the end-of-sequence token.
            feed_forward_size (int): The size of the feed-forward network.
            hash_seed (None or int): The seed for hashing functions.
            hidden_act (str): The activation function for hidden layers.
            hidden_dropout_prob (float): The dropout probability for hidden layers.
            hidden_size (int): The size of the hidden layers.
            initializer_range (float): The range for weight initialization.
            is_decoder (bool): Whether the model is used as a decoder.
            layer_norm_eps (float): Epsilon value for layer normalization.
            local_num_chunks_before (int): Number of local attention chunks before.
            local_num_chunks_after (int): Number of local attention chunks after.
            local_attention_probs_dropout_prob (float): Dropout probability for local attention.
            local_attn_chunk_length (int): Length of chunks for local attention.
            lsh_attn_chunk_length (int): Length of chunks for LSH attention.
            lsh_attention_probs_dropout_prob (float): Dropout probability for LSH attention.
            lsh_num_chunks_before (int): Number of LSH attention chunks before.
            lsh_num_chunks_after (int): Number of LSH attention chunks after.
            max_position_embeddings (int): The maximum number of position embeddings.
            num_attention_heads (int): The number of attention heads.
            num_buckets (None or tuple): The number of buckets for hashing.
            num_hashes (int): The number of hashes for LSH attention.
            pad_token_id (int): The token ID for padding.
            vocab_size (int): The size of the vocabulary.
            tie_word_embeddings (bool): Whether to tie word embeddings.
            use_cache (bool): Whether to cache intermediate values.
            classifier_dropout (None or float): Dropout probability for classifier layers.

        Returns:
            None.

        Raises:
            None.
        """
        self.hash_seed = hash_seed
        self.vocab_size = vocab_size
        self.attention_head_size = attention_head_size
        self.hidden_size = hidden_size
        self.num_attention_heads = num_attention_heads
        self.num_hashes = num_hashes
        self.num_hidden_layers = len(attn_layers)
        self.num_buckets = tuple(num_buckets) if isinstance(num_buckets, list) else num_buckets
        self.lsh_attn_chunk_length = lsh_attn_chunk_length
        self.local_attn_chunk_length = local_attn_chunk_length
        self.lsh_num_chunks_after = lsh_num_chunks_after
        self.lsh_num_chunks_before = lsh_num_chunks_before
        self.local_num_chunks_after = local_num_chunks_after
        self.local_num_chunks_before = local_num_chunks_before
        self.hidden_act = hidden_act
        self.feed_forward_size = feed_forward_size
        self.hidden_dropout_prob = hidden_dropout_prob
        self.lsh_attention_probs_dropout_prob = lsh_attention_probs_dropout_prob
        self.local_attention_probs_dropout_prob = local_attention_probs_dropout_prob
        self.max_position_embeddings = max_position_embeddings
        self.initializer_range = initializer_range
        self.layer_norm_eps = layer_norm_eps
        self.axial_pos_embds = axial_pos_embds
        self.axial_pos_shape = tuple(axial_pos_shape)
        self.axial_pos_embds_dim = tuple(axial_pos_embds_dim)
        self.axial_norm_std = axial_norm_std
        self.chunk_size_lm_head = chunk_size_lm_head
        self.attn_layers = attn_layers
        self.use_cache = use_cache
        self.classifier_dropout = classifier_dropout
        super().__init__(
            pad_token_id=pad_token_id,
            eos_token_id=eos_token_id,
            is_decoder=is_decoder,
            tie_word_embeddings=tie_word_embeddings,
            **kwargs,
        )

mindnlp.transformers.models.reformer.configuration_reformer.ReformerConfig.__init__(attention_head_size=64, attn_layers=['local', 'lsh', 'local', 'lsh', 'local', 'lsh'], axial_norm_std=1.0, axial_pos_embds=True, axial_pos_shape=[64, 64], axial_pos_embds_dim=[64, 192], chunk_size_lm_head=0, eos_token_id=2, feed_forward_size=512, hash_seed=None, hidden_act='relu', hidden_dropout_prob=0.05, hidden_size=256, initializer_range=0.02, is_decoder=False, layer_norm_eps=1e-12, local_num_chunks_before=1, local_num_chunks_after=0, local_attention_probs_dropout_prob=0.05, local_attn_chunk_length=64, lsh_attn_chunk_length=64, lsh_attention_probs_dropout_prob=0.0, lsh_num_chunks_before=1, lsh_num_chunks_after=0, max_position_embeddings=4096, num_attention_heads=12, num_buckets=None, num_hashes=1, pad_token_id=0, vocab_size=320, tie_word_embeddings=False, use_cache=True, classifier_dropout=None, **kwargs)

Initializes a new instance of the ReformerConfig class.

PARAMETER	DESCRIPTION
attention_head_size	The size of each attention head. TYPE: int DEFAULT: 64
attn_layers	The list of attention layer types to be used. TYPE: list DEFAULT: ['local', 'lsh', 'local', 'lsh', 'local', 'lsh']
axial_norm_std	Standard deviation for axial positional embeddings normalization. TYPE: float DEFAULT: 1.0
axial_pos_embds	Whether to use axial positional embeddings. TYPE: bool DEFAULT: True
axial_pos_shape	The shape of axial positional embeddings. TYPE: list DEFAULT: [64, 64]
axial_pos_embds_dim	The dimensions of axial positional embeddings. TYPE: list DEFAULT: [64, 192]
chunk_size_lm_head	Size of chunk for the language model head. TYPE: int DEFAULT: 0
eos_token_id	The token ID for the end-of-sequence token. TYPE: int DEFAULT: 2
feed_forward_size	The size of the feed-forward network. TYPE: int DEFAULT: 512
hash_seed	The seed for hashing functions. TYPE: None or int DEFAULT: None
hidden_act	The activation function for hidden layers. TYPE: str DEFAULT: 'relu'
hidden_dropout_prob	The dropout probability for hidden layers. TYPE: float DEFAULT: 0.05
hidden_size	The size of the hidden layers. TYPE: int DEFAULT: 256
initializer_range	The range for weight initialization. TYPE: float DEFAULT: 0.02
is_decoder	Whether the model is used as a decoder. TYPE: bool DEFAULT: False
layer_norm_eps	Epsilon value for layer normalization. TYPE: float DEFAULT: 1e-12
local_num_chunks_before	Number of local attention chunks before. TYPE: int DEFAULT: 1
local_num_chunks_after	Number of local attention chunks after. TYPE: int DEFAULT: 0
local_attention_probs_dropout_prob	Dropout probability for local attention. TYPE: float DEFAULT: 0.05
local_attn_chunk_length	Length of chunks for local attention. TYPE: int DEFAULT: 64
lsh_attn_chunk_length	Length of chunks for LSH attention. TYPE: int DEFAULT: 64
lsh_attention_probs_dropout_prob	Dropout probability for LSH attention. TYPE: float DEFAULT: 0.0
lsh_num_chunks_before	Number of LSH attention chunks before. TYPE: int DEFAULT: 1
lsh_num_chunks_after	Number of LSH attention chunks after. TYPE: int DEFAULT: 0
max_position_embeddings	The maximum number of position embeddings. TYPE: int DEFAULT: 4096
num_attention_heads	The number of attention heads. TYPE: int DEFAULT: 12
num_buckets	The number of buckets for hashing. TYPE: None or tuple DEFAULT: None
num_hashes	The number of hashes for LSH attention. TYPE: int DEFAULT: 1
pad_token_id	The token ID for padding. TYPE: int DEFAULT: 0
vocab_size	The size of the vocabulary. TYPE: int DEFAULT: 320
tie_word_embeddings	Whether to tie word embeddings. TYPE: bool DEFAULT: False
use_cache	Whether to cache intermediate values. TYPE: bool DEFAULT: True
classifier_dropout	Dropout probability for classifier layers. TYPE: None or float DEFAULT: None

RETURNS	DESCRIPTION
	None.

Source code in mindnlp/transformers/models/reformer/configuration_reformer.py

def __init__(
    self,
    attention_head_size=64,
    attn_layers=["local", "lsh", "local", "lsh", "local", "lsh"],
    axial_norm_std=1.0,
    axial_pos_embds=True,
    axial_pos_shape=[64, 64],
    axial_pos_embds_dim=[64, 192],
    chunk_size_lm_head=0,
    eos_token_id=2,
    feed_forward_size=512,
    hash_seed=None,
    hidden_act="relu",
    hidden_dropout_prob=0.05,
    hidden_size=256,
    initializer_range=0.02,
    is_decoder=False,
    layer_norm_eps=1e-12,
    local_num_chunks_before=1,
    local_num_chunks_after=0,
    local_attention_probs_dropout_prob=0.05,
    local_attn_chunk_length=64,
    lsh_attn_chunk_length=64,
    lsh_attention_probs_dropout_prob=0.0,
    lsh_num_chunks_before=1,
    lsh_num_chunks_after=0,
    max_position_embeddings=4096,
    num_attention_heads=12,
    num_buckets=None,
    num_hashes=1,
    pad_token_id=0,
    vocab_size=320,
    tie_word_embeddings=False,
    use_cache=True,
    classifier_dropout=None,
    **kwargs,
):
    """
    Initializes a new instance of the ReformerConfig class.

    Args:
        attention_head_size (int): The size of each attention head.
        attn_layers (list): The list of attention layer types to be used.
        axial_norm_std (float): Standard deviation for axial positional embeddings normalization.
        axial_pos_embds (bool): Whether to use axial positional embeddings.
        axial_pos_shape (list): The shape of axial positional embeddings.
        axial_pos_embds_dim (list): The dimensions of axial positional embeddings.
        chunk_size_lm_head (int): Size of chunk for the language model head.
        eos_token_id (int): The token ID for the end-of-sequence token.
        feed_forward_size (int): The size of the feed-forward network.
        hash_seed (None or int): The seed for hashing functions.
        hidden_act (str): The activation function for hidden layers.
        hidden_dropout_prob (float): The dropout probability for hidden layers.
        hidden_size (int): The size of the hidden layers.
        initializer_range (float): The range for weight initialization.
        is_decoder (bool): Whether the model is used as a decoder.
        layer_norm_eps (float): Epsilon value for layer normalization.
        local_num_chunks_before (int): Number of local attention chunks before.
        local_num_chunks_after (int): Number of local attention chunks after.
        local_attention_probs_dropout_prob (float): Dropout probability for local attention.
        local_attn_chunk_length (int): Length of chunks for local attention.
        lsh_attn_chunk_length (int): Length of chunks for LSH attention.
        lsh_attention_probs_dropout_prob (float): Dropout probability for LSH attention.
        lsh_num_chunks_before (int): Number of LSH attention chunks before.
        lsh_num_chunks_after (int): Number of LSH attention chunks after.
        max_position_embeddings (int): The maximum number of position embeddings.
        num_attention_heads (int): The number of attention heads.
        num_buckets (None or tuple): The number of buckets for hashing.
        num_hashes (int): The number of hashes for LSH attention.
        pad_token_id (int): The token ID for padding.
        vocab_size (int): The size of the vocabulary.
        tie_word_embeddings (bool): Whether to tie word embeddings.
        use_cache (bool): Whether to cache intermediate values.
        classifier_dropout (None or float): Dropout probability for classifier layers.

    Returns:
        None.

    Raises:
        None.
    """
    self.hash_seed = hash_seed
    self.vocab_size = vocab_size
    self.attention_head_size = attention_head_size
    self.hidden_size = hidden_size
    self.num_attention_heads = num_attention_heads
    self.num_hashes = num_hashes
    self.num_hidden_layers = len(attn_layers)
    self.num_buckets = tuple(num_buckets) if isinstance(num_buckets, list) else num_buckets
    self.lsh_attn_chunk_length = lsh_attn_chunk_length
    self.local_attn_chunk_length = local_attn_chunk_length
    self.lsh_num_chunks_after = lsh_num_chunks_after
    self.lsh_num_chunks_before = lsh_num_chunks_before
    self.local_num_chunks_after = local_num_chunks_after
    self.local_num_chunks_before = local_num_chunks_before
    self.hidden_act = hidden_act
    self.feed_forward_size = feed_forward_size
    self.hidden_dropout_prob = hidden_dropout_prob
    self.lsh_attention_probs_dropout_prob = lsh_attention_probs_dropout_prob
    self.local_attention_probs_dropout_prob = local_attention_probs_dropout_prob
    self.max_position_embeddings = max_position_embeddings
    self.initializer_range = initializer_range
    self.layer_norm_eps = layer_norm_eps
    self.axial_pos_embds = axial_pos_embds
    self.axial_pos_shape = tuple(axial_pos_shape)
    self.axial_pos_embds_dim = tuple(axial_pos_embds_dim)
    self.axial_norm_std = axial_norm_std
    self.chunk_size_lm_head = chunk_size_lm_head
    self.attn_layers = attn_layers
    self.use_cache = use_cache
    self.classifier_dropout = classifier_dropout
    super().__init__(
        pad_token_id=pad_token_id,
        eos_token_id=eos_token_id,
        is_decoder=is_decoder,
        tie_word_embeddings=tie_word_embeddings,
        **kwargs,
    )

mindnlp.transformers.models.reformer.modeling_reformer

MindSpore REFORMER model.

mindnlp.transformers.models.reformer.modeling_reformer.AxialPositionEmbeddings

Bases: Module

Constructs axial position embeddings. Useful for very long input sequences to save memory and time.

Source code in mindnlp/transformers/models/reformer/modeling_reformer.py

class AxialPositionEmbeddings(nn.Module):
    """
    Constructs axial position embeddings. Useful for very long input sequences to save memory and time.
    """
    def __init__(self, config):
        """
        This method initializes an instance of the AxialPositionEmbeddings class.

        Args:
            self: The instance of the AxialPositionEmbeddings class.
            config:
                An object containing configuration parameters for the axial position embeddings.

                - axial_pos_shape: A list of integers representing the shape of the axial positions.
                - axial_pos_embds_dim: A list of integers representing the dimensions of the axial position embeddings.
                - hidden_dropout_prob: A float value representing the dropout probability.
                - hidden_size: An integer representing the hidden size of the model.

        Returns:
            None.

        Raises:
            ValueError: If the sum of axial_pos_embds_dim does not match the hidden_size specified in the configuration.
        """
        super().__init__()
        self.axial_pos_shape = config.axial_pos_shape
        self.axial_pos_embds_dim = config.axial_pos_embds_dim
        self.dropout = config.hidden_dropout_prob

        self.least_common_mult_chunk_length = _get_least_common_mult_chunk_len(config)
        self.weights = []

        if sum(self.axial_pos_embds_dim) != config.hidden_size:
            raise ValueError(
                f"Make sure that config.axial_pos_embds factors: {self.axial_pos_embds_dim} sum to "
                f"config.hidden_size: {config.hidden_size}"
            )

        # create weights
        for axis, axial_pos_embd_dim in enumerate(self.axial_pos_embds_dim):
            # create expanded shapes
            ax_shape = [1] * len(self.axial_pos_shape)
            ax_shape[axis] = self.axial_pos_shape[axis]
            ax_shape = tuple(ax_shape) + (axial_pos_embd_dim,)

            # create tensor and init
            self.weights.append(Parameter(ops.ones(ax_shape, dtype=mindspore.float32), name=f'weights.{axis}'))

        self.weights = ParameterTuple(self.weights)

    def forward(self, position_ids):
        """
        This method forwards position encodings based on the given position IDs and axial position weights.

        Args:
            self: The instance of the AxialPositionEmbeddings class.
            position_ids (torch.Tensor): A 2D tensor representing the position IDs of the input sequences.
              It has a shape of (batch_size, sequence_length).

        Returns:
            None: This method does not return any value.

        Raises:
            ValueError:
                Raised if the following conditions are met:

                1. During training, if the product of the axial_pos_shape factors does not match the sequence length.
                2. During training and dropout is enabled, if the operation cannot be performed.
                3. During inference, if the product of the axial_pos_shape factors is less than the sequence length.
                4. During inference, if the required position encodings columns exceed the available columns.
        """
        # broadcast weights to correct shape
        batch_size = position_ids.shape[0]
        sequence_length = position_ids.shape[1]

        broadcasted_weights = [
            weight.expand((batch_size,) + self.axial_pos_shape + weight.shape[-1:]) for weight in self.weights
        ]

        if self.training is True:
            if reduce(mul, self.axial_pos_shape) != sequence_length:
                raise ValueError(
                    f"If training, make sure that config.axial_pos_shape factors: {self.axial_pos_shape} multiply to "
                    f"sequence length. Got prod({self.axial_pos_shape}) != sequence_length: {sequence_length}. "
                    f"You might want to consider padding your sequence length to {reduce(mul, self.axial_pos_shape)} "
                    "or changing config.axial_pos_shape."
                )

            if self.dropout > 0:
                weights = ops.cat(broadcasted_weights, axis=-1)
                # permute weights so that 2D correctly drops dims 1 and 2
                transposed_weights = weights.swapaxes(2, 1)
                # drop entire matrix of last two dims (prev dims 1 and 2)
                dropped_transposed_weights = ops.dropout2d(
                    transposed_weights, p=self.dropout, training=self.training
                )
                dropped_weights = dropped_transposed_weights.swapaxes(2, 1)

                position_encodings = ops.reshape(dropped_weights, (batch_size, sequence_length, -1))

            else:
                position_encodings = ops.cat(
                    [ops.reshape(weight, (batch_size, sequence_length, -1)) for weight in broadcasted_weights],
                    axis=-1,
                )

        else:
            if reduce(mul, self.axial_pos_shape) < sequence_length:
                raise ValueError(
                    f"Make sure that config.axial_pos_shape factors: {self.axial_pos_shape} multiply at least to "
                    f"max(sequence_length, least_common_mult_chunk_length): max({sequence_length}, "
                    f"{self.least_common_mult_chunk_length})."
                )

            # compute how many columns are needed
            max_position_id = position_ids.max().item()
            required_pos_encodings_columns = -(-(max_position_id + 1) // self.axial_pos_shape[1])

            # cut to columns that are needed
            position_encodings = ops.cat(
                [weight[:, :required_pos_encodings_columns] for weight in broadcasted_weights], axis=-1
            )
            position_encodings = ops.reshape(position_encodings, (batch_size, -1, position_encodings.shape[-1]))

            # select correct position encodings
            position_encodings = ops.cat(
                [
                    ops.index_select(position_encodings[i], 0, position_ids[i]).unsqueeze(0)
                    for i in range(batch_size)
                ],
                axis=0,
            )

        return position_encodings

mindnlp.transformers.models.reformer.modeling_reformer.AxialPositionEmbeddings.__init__(config)

This method initializes an instance of the AxialPositionEmbeddings class.

PARAMETER	DESCRIPTION
self	The instance of the AxialPositionEmbeddings class.
config	An object containing configuration parameters for the axial position embeddings. axial_pos_shape: A list of integers representing the shape of the axial positions. axial_pos_embds_dim: A list of integers representing the dimensions of the axial position embeddings. hidden_dropout_prob: A float value representing the dropout probability. hidden_size: An integer representing the hidden size of the model.

RETURNS	DESCRIPTION
	None.

RAISES	DESCRIPTION
ValueError	If the sum of axial_pos_embds_dim does not match the hidden_size specified in the configuration.

Source code in mindnlp/transformers/models/reformer/modeling_reformer.py

def __init__(self, config):
    """
    This method initializes an instance of the AxialPositionEmbeddings class.

    Args:
        self: The instance of the AxialPositionEmbeddings class.
        config:
            An object containing configuration parameters for the axial position embeddings.

            - axial_pos_shape: A list of integers representing the shape of the axial positions.
            - axial_pos_embds_dim: A list of integers representing the dimensions of the axial position embeddings.
            - hidden_dropout_prob: A float value representing the dropout probability.
            - hidden_size: An integer representing the hidden size of the model.

    Returns:
        None.

    Raises:
        ValueError: If the sum of axial_pos_embds_dim does not match the hidden_size specified in the configuration.
    """
    super().__init__()
    self.axial_pos_shape = config.axial_pos_shape
    self.axial_pos_embds_dim = config.axial_pos_embds_dim
    self.dropout = config.hidden_dropout_prob

    self.least_common_mult_chunk_length = _get_least_common_mult_chunk_len(config)
    self.weights = []

    if sum(self.axial_pos_embds_dim) != config.hidden_size:
        raise ValueError(
            f"Make sure that config.axial_pos_embds factors: {self.axial_pos_embds_dim} sum to "
            f"config.hidden_size: {config.hidden_size}"
        )

    # create weights
    for axis, axial_pos_embd_dim in enumerate(self.axial_pos_embds_dim):
        # create expanded shapes
        ax_shape = [1] * len(self.axial_pos_shape)
        ax_shape[axis] = self.axial_pos_shape[axis]
        ax_shape = tuple(ax_shape) + (axial_pos_embd_dim,)

        # create tensor and init
        self.weights.append(Parameter(ops.ones(ax_shape, dtype=mindspore.float32), name=f'weights.{axis}'))

    self.weights = ParameterTuple(self.weights)

mindnlp.transformers.models.reformer.modeling_reformer.AxialPositionEmbeddings.forward(position_ids)

This method forwards position encodings based on the given position IDs and axial position weights.

PARAMETER	DESCRIPTION
self	The instance of the AxialPositionEmbeddings class.
position_ids	A 2D tensor representing the position IDs of the input sequences. It has a shape of (batch_size, sequence_length). TYPE: Tensor

RETURNS	DESCRIPTION
None	This method does not return any value.

RAISES

DESCRIPTION

ValueError

Raised if the following conditions are met: During training, if the product of the axial_pos_shape factors does not match the sequence length. During training and dropout is enabled, if the operation cannot be performed. During inference, if the product of the axial_pos_shape factors is less than the sequence length. During inference, if the required position encodings columns exceed the available columns.

Source code in mindnlp/transformers/models/reformer/modeling_reformer.py

def forward(self, position_ids):
    """
    This method forwards position encodings based on the given position IDs and axial position weights.

    Args:
        self: The instance of the AxialPositionEmbeddings class.
        position_ids (torch.Tensor): A 2D tensor representing the position IDs of the input sequences.
          It has a shape of (batch_size, sequence_length).

    Returns:
        None: This method does not return any value.

    Raises:
        ValueError:
            Raised if the following conditions are met:

            1. During training, if the product of the axial_pos_shape factors does not match the sequence length.
            2. During training and dropout is enabled, if the operation cannot be performed.
            3. During inference, if the product of the axial_pos_shape factors is less than the sequence length.
            4. During inference, if the required position encodings columns exceed the available columns.
    """
    # broadcast weights to correct shape
    batch_size = position_ids.shape[0]
    sequence_length = position_ids.shape[1]

    broadcasted_weights = [
        weight.expand((batch_size,) + self.axial_pos_shape + weight.shape[-1:]) for weight in self.weights
    ]

    if self.training is True:
        if reduce(mul, self.axial_pos_shape) != sequence_length:
            raise ValueError(
                f"If training, make sure that config.axial_pos_shape factors: {self.axial_pos_shape} multiply to "
                f"sequence length. Got prod({self.axial_pos_shape}) != sequence_length: {sequence_length}. "
                f"You might want to consider padding your sequence length to {reduce(mul, self.axial_pos_shape)} "
                "or changing config.axial_pos_shape."
            )

        if self.dropout > 0:
            weights = ops.cat(broadcasted_weights, axis=-1)
            # permute weights so that 2D correctly drops dims 1 and 2
            transposed_weights = weights.swapaxes(2, 1)
            # drop entire matrix of last two dims (prev dims 1 and 2)
            dropped_transposed_weights = ops.dropout2d(
                transposed_weights, p=self.dropout, training=self.training
            )
            dropped_weights = dropped_transposed_weights.swapaxes(2, 1)

            position_encodings = ops.reshape(dropped_weights, (batch_size, sequence_length, -1))

        else:
            position_encodings = ops.cat(
                [ops.reshape(weight, (batch_size, sequence_length, -1)) for weight in broadcasted_weights],
                axis=-1,
            )

    else:
        if reduce(mul, self.axial_pos_shape) < sequence_length:
            raise ValueError(
                f"Make sure that config.axial_pos_shape factors: {self.axial_pos_shape} multiply at least to "
                f"max(sequence_length, least_common_mult_chunk_length): max({sequence_length}, "
                f"{self.least_common_mult_chunk_length})."
            )

        # compute how many columns are needed
        max_position_id = position_ids.max().item()
        required_pos_encodings_columns = -(-(max_position_id + 1) // self.axial_pos_shape[1])

        # cut to columns that are needed
        position_encodings = ops.cat(
            [weight[:, :required_pos_encodings_columns] for weight in broadcasted_weights], axis=-1
        )
        position_encodings = ops.reshape(position_encodings, (batch_size, -1, position_encodings.shape[-1]))

        # select correct position encodings
        position_encodings = ops.cat(
            [
                ops.index_select(position_encodings[i], 0, position_ids[i]).unsqueeze(0)
                for i in range(batch_size)
            ],
            axis=0,
        )

    return position_encodings

mindnlp.transformers.models.reformer.modeling_reformer.ChunkReformerFeedForward

Bases: Module

This class represents a feed-forward module for chunked reformer attention output in a neural network.

The ChunkReformerFeedForward class inherits from the nn.Module class and is designed to process attention output in a chunked manner. It applies layer normalization, dense transformation, and output transformation to the input hidden states.

ATTRIBUTE	DESCRIPTION
chunk_size_feed_forward	The size of each chunk for processing the attention output. TYPE: int
seq_len_dim	The dimension representing the sequence length in the attention output. TYPE: int
layer_norm	The layer normalization module applied to the hidden states. TYPE: LayerNorm
dense	The dense transformation module applied to the layer-normalized hidden states. TYPE: ReformerFeedForwardDense
output	The final output transformation module applied to the transformed hidden states. TYPE: ReformerFeedForwardOutput

METHOD	DESCRIPTION
__init__	Initializes a new instance of the ChunkReformerFeedForward class. Args: config (object): The configuration object containing hyperparameters and settings for the module.
forward	Constructs the output of the ChunkReformerFeedForward module. Args: attention_output (Tensor): The input attention output tensor. Returns: Tensor: The forwarded output tensor after applying layer normalization, dense transformation, and output transformation.
forward_chunk	Constructs a chunk of the ChunkReformerFeedForward module. Args: hidden_states (Tensor): The input hidden states tensor. Returns: Tensor: The forwarded chunk of the feed-forward module after applying layer normalization, dense transformation, and output transformation.

Source code in mindnlp/transformers/models/reformer/modeling_reformer.py

class ChunkReformerFeedForward(nn.Module):

    """
    This class represents a feed-forward module for chunked reformer attention output in a neural network.

    The ChunkReformerFeedForward class inherits from the nn.Module class and is designed to process attention output
    in a chunked manner. It applies layer normalization, dense transformation, and output transformation to the input
    hidden states.

    Attributes:
        chunk_size_feed_forward (int): The size of each chunk for processing the attention output.
        seq_len_dim (int): The dimension representing the sequence length in the attention output.
        layer_norm (nn.LayerNorm): The layer normalization module applied to the hidden states.
        dense (ReformerFeedForwardDense): The dense transformation module applied to the layer-normalized hidden states.
        output (ReformerFeedForwardOutput): The final output transformation module applied to the
            transformed hidden states.

    Methods:
        __init__(self, config):
            Initializes a new instance of the ChunkReformerFeedForward class.

            Args:

            - config (object): The configuration object containing hyperparameters and settings for the module.

        forward(self, attention_output):
            Constructs the output of the ChunkReformerFeedForward module.

            Args:

            - attention_output (Tensor): The input attention output tensor.

            Returns:

            - Tensor: The forwarded output tensor after applying layer normalization, dense transformation,
            and output transformation.

        forward_chunk(self, hidden_states):
            Constructs a chunk of the ChunkReformerFeedForward module.

            Args:

            - hidden_states (Tensor): The input hidden states tensor.

            Returns:

            - Tensor: The forwarded chunk of the feed-forward module after applying layer normalization,
            dense transformation, and output transformation.
    """
    def __init__(self, config):

        """
        Initializes a ChunkReformerFeedForward instance.

        Args:
            self (ChunkReformerFeedForward): The ChunkReformerFeedForward instance itself.
            config:
                A configuration object containing the necessary parameters for initialization.

                - chunk_size_feed_forward (int): The chunk size for feed-forward operations.
                - hidden_size (int): The size of the hidden layers.
                - layer_norm_eps (float): The epsilon value for layer normalization.

        Returns:
            None.

        Raises:
            TypeError: If the provided config is not of the expected type.
            ValueError: If any required parameter is missing in the config.
        """
        super().__init__()
        self.chunk_size_feed_forward = config.chunk_size_feed_forward
        self.seq_len_dim = 1

        self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dense = ReformerFeedForwardDense(config)
        self.output = ReformerFeedForwardOutput(config)

    def forward(self, attention_output):

        """
        Constructs the feed-forward chunk reformer for the given attention output.

        Args:
            self (ChunkReformerFeedForward): Instance of the ChunkReformerFeedForward class.
            attention_output (Tensor): The attention output tensor to be processed by the chunk reformer.

        Returns:
            None.

        Raises:
            TypeError: If the input parameters are not of the expected types.
            ValueError: If the chunk size for feed-forward is invalid.
            RuntimeError: If there is an issue with applying chunking during the forward pass.
        """
        return apply_chunking_to_forward(
            self.forward_chunk,
            self.chunk_size_feed_forward,
            self.seq_len_dim,
            attention_output,
        )

    def forward_chunk(self, hidden_states):

        """
        Constructs a chunk of reformer feed forward layer.

        Args:
            self (ChunkReformerFeedForward): The instance of the ChunkReformerFeedForward class.
            hidden_states (tensor): The input tensor containing the hidden states.

        Returns:
            None.

        Raises:
            AttributeError: If the 'layer_norm', 'dense', or 'output' attributes are not found in the instance.
            ValueError: If the 'hidden_states' parameter is not a valid tensor.
        """
        hidden_states = self.layer_norm(hidden_states)
        hidden_states = self.dense(hidden_states)
        return self.output(hidden_states)

mindnlp.transformers.models.reformer.modeling_reformer.ChunkReformerFeedForward.__init__(config)

Initializes a ChunkReformerFeedForward instance.

PARAMETER	DESCRIPTION
self	The ChunkReformerFeedForward instance itself. TYPE: ChunkReformerFeedForward
config	A configuration object containing the necessary parameters for initialization. chunk_size_feed_forward (int): The chunk size for feed-forward operations. hidden_size (int): The size of the hidden layers. layer_norm_eps (float): The epsilon value for layer normalization.

RETURNS	DESCRIPTION
	None.

RAISES	DESCRIPTION
TypeError	If the provided config is not of the expected type.
ValueError	If any required parameter is missing in the config.

Source code in mindnlp/transformers/models/reformer/modeling_reformer.py

def __init__(self, config):

    """
    Initializes a ChunkReformerFeedForward instance.

    Args:
        self (ChunkReformerFeedForward): The ChunkReformerFeedForward instance itself.
        config:
            A configuration object containing the necessary parameters for initialization.

            - chunk_size_feed_forward (int): The chunk size for feed-forward operations.
            - hidden_size (int): The size of the hidden layers.
            - layer_norm_eps (float): The epsilon value for layer normalization.

    Returns:
        None.

    Raises:
        TypeError: If the provided config is not of the expected type.
        ValueError: If any required parameter is missing in the config.
    """
    super().__init__()
    self.chunk_size_feed_forward = config.chunk_size_feed_forward
    self.seq_len_dim = 1

    self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
    self.dense = ReformerFeedForwardDense(config)
    self.output = ReformerFeedForwardOutput(config)

mindnlp.transformers.models.reformer.modeling_reformer.ChunkReformerFeedForward.forward(attention_output)

Constructs the feed-forward chunk reformer for the given attention output.

PARAMETER	DESCRIPTION
self	Instance of the ChunkReformerFeedForward class. TYPE: ChunkReformerFeedForward
attention_output	The attention output tensor to be processed by the chunk reformer. TYPE: Tensor

RETURNS	DESCRIPTION
	None.

RAISES	DESCRIPTION
TypeError	If the input parameters are not of the expected types.
ValueError	If the chunk size for feed-forward is invalid.
RuntimeError	If there is an issue with applying chunking during the forward pass.

Source code in mindnlp/transformers/models/reformer/modeling_reformer.py

def forward(self, attention_output):

    """
    Constructs the feed-forward chunk reformer for the given attention output.

    Args:
        self (ChunkReformerFeedForward): Instance of the ChunkReformerFeedForward class.
        attention_output (Tensor): The attention output tensor to be processed by the chunk reformer.

    Returns:
        None.

    Raises:
        TypeError: If the input parameters are not of the expected types.
        ValueError: If the chunk size for feed-forward is invalid.
        RuntimeError: If there is an issue with applying chunking during the forward pass.
    """
    return apply_chunking_to_forward(
        self.forward_chunk,
        self.chunk_size_feed_forward,
        self.seq_len_dim,
        attention_output,
    )

mindnlp.transformers.models.reformer.modeling_reformer.ChunkReformerFeedForward.forward_chunk(hidden_states)

Constructs a chunk of reformer feed forward layer.

PARAMETER	DESCRIPTION
self	The instance of the ChunkReformerFeedForward class. TYPE: ChunkReformerFeedForward
hidden_states	The input tensor containing the hidden states. TYPE: tensor

RETURNS	DESCRIPTION
	None.

RAISES	DESCRIPTION
AttributeError	If the 'layer_norm', 'dense', or 'output' attributes are not found in the instance.
ValueError	If the 'hidden_states' parameter is not a valid tensor.

Source code in mindnlp/transformers/models/reformer/modeling_reformer.py

def forward_chunk(self, hidden_states):

    """
    Constructs a chunk of reformer feed forward layer.

    Args:
        self (ChunkReformerFeedForward): The instance of the ChunkReformerFeedForward class.
        hidden_states (tensor): The input tensor containing the hidden states.

    Returns:
        None.

    Raises:
        AttributeError: If the 'layer_norm', 'dense', or 'output' attributes are not found in the instance.
        ValueError: If the 'hidden_states' parameter is not a valid tensor.
    """
    hidden_states = self.layer_norm(hidden_states)
    hidden_states = self.dense(hidden_states)
    return self.output(hidden_states)

mindnlp.transformers.models.reformer.modeling_reformer.EfficientAttentionMixin

A few utilities for nn.Modules in Reformer, to be used as a mixin.

Source code in mindnlp/transformers/models/reformer/modeling_reformer.py

class EfficientAttentionMixin:
    """
    A few utilities for nn.Modules in Reformer, to be used as a mixin.
    """
    def _look_adjacent(self, vectors, num_chunks_before, num_chunks_after):
        """
        Used to implement attention between consecutive chunks.

        Args:
            vectors: array of shape [batch_size, num_attention_heads, n_chunks, chunk_len, ...]
            num_chunks_before: chunks before current chunk to include in attention
            num_chunks_after: chunks after current chunk to include in attention

        Returns:
            tensor of shape [num_chunks, N * chunk_length, ...], where N = (1 + num_chunks_before + num_chunks_after).
        """
        if num_chunks_before == 0 and num_chunks_after == 0:
            return vectors

        slices = []
        for i in range(-num_chunks_before, num_chunks_after + 1):
            if i == 0:
                slices.append(vectors)
            else:
                slices.append(ops.cat([vectors[:, :, i:, ...], vectors[:, :, :i, ...]], axis=2))
        return ops.cat(slices, axis=3)

    def _split_hidden_size_dim(self, x, num_attn_heads, attn_head_size):
        """
        splits hidden_size dim into attn_head_size and num_attn_heads
        """
        new_x_shape = x.shape[:-1] + (num_attn_heads, attn_head_size)
        x = x.view(*new_x_shape)
        return x.swapaxes(2, 1)

    def _merge_hidden_size_dims(self, x, num_attn_heads, attn_head_size):
        """
        merges attn_head_size dim and num_attn_heads dim into hidden_size
        """
        x = x.permute(0, 2, 1, 3)
        return ops.reshape(x, (x.shape[0], -1, num_attn_heads * attn_head_size))

    def _split_seq_length_dim_to(self, vectors, dim_factor_1, dim_factor_2, num_attn_heads, attn_head_size=None):
        """
        splits sequence length dim of vectors into `dim_factor_1` and `dim_factor_2` dims
        """
        batch_size = vectors.shape[0]
        split_dim_shape = (batch_size, num_attn_heads, dim_factor_1, dim_factor_2)

        if len(vectors.shape) == 4:
            return ops.reshape(vectors, split_dim_shape + (attn_head_size,))
        if len(vectors.shape) == 3:
            return ops.reshape(vectors, split_dim_shape)
        raise ValueError(f"Input vector rank should be one of [3, 4], but is: {len(vectors.shape)}")

mindnlp.transformers.models.reformer.modeling_reformer.LSHSelfAttention

Bases: Module, EfficientAttentionMixin

This class represents a self-attention mechanism using Locality Sensitive Hashing (LSH) for efficient attention computation. It inherits from nn.Module, EfficientAttentionMixin.

The class implements LSH self-attention mechanism for neural networks. It includes methods for initializing the LSH attention layer, forwarding the attention mechanism, computing attention masks, hashing vectors, and other related operations.

ATTRIBUTE	DESCRIPTION
config	Configuration parameters for the LSH self-attention layer.
chunk_length	Length of each attention chunk.
num_hashes	Number of hashes used in the LSH mechanism.
num_buckets	Number of buckets used for hashing.
num_chunks_before	Number of chunks to consider before the current chunk.
num_chunks_after	Number of chunks to consider after the current chunk.
hash_seed	Seed for random hash rotations.
is_decoder	Flag indicating if the layer is used in a decoder.
max_position_embeddings	Maximum position embeddings allowed.
dropout	Dropout probability for attention weights.
num_attention_heads	Number of attention heads.
attention_head_size	Size of each attention head.
all_head_size	Total size of all attention heads.
hidden_size	Size of the hidden layer.
query_key	Dense layer for query and key vectors.
value	Dense layer for value vectors.
self_mask_value_float16	Float16 value for masking in self-attention.
self_mask_value_float32	Float32 value for masking in self-attention.
mask_value_float16	Float16 value for general masking.
mask_value_float32	Float32 value for general masking.

METHOD	DESCRIPTION
forward	Constructs the LSH self-attention mechanism based on input hidden states and optional parameters.
_query_per_attn_head	Computes query vectors per attention head.
_value_per_attn_head	Computes value vectors per attention head.
_hash_vectors	Hashes input vectors into buckets for attention computation.
_get_sorted_bucket_idx_and_undo_sorted_bucket_idx	Computes sorted bucket indices for efficient attention calculation.
_set_num_buckets	Sets the number of buckets based on the input sequence length.
_attend	Computes attention scores and outputs based on query, key, and value vectors.
_compute_attn_mask	Computes attention mask based on query and key indices.
_get_relevant_hid_states_and_buckets	Retrieves relevant hidden states and buckets for efficient attention calculation.
_expand_to_indices_in_relevant_chunk	Expands indices for relevant chunks in hidden states.
_len_and_dim_norm	Normalizes vectors based on length and attention head size.
_len_norm	Length normalization for input vectors.
_gather_by_expansion	Expands indices and vectors for all hashes and gathers relevant elements.

Note

This class is designed for implementing efficient self-attention mechanisms using Locality Sensitive Hashing.

Source code in mindnlp/transformers/models/reformer/modeling_reformer.py

class LSHSelfAttention(nn.Module, EfficientAttentionMixin):

    """
    This class represents a self-attention mechanism using Locality Sensitive Hashing (LSH) for efficient attention
    computation.
    It inherits from nn.Module, EfficientAttentionMixin.

    The class implements LSH self-attention mechanism for neural networks. It includes methods for initializing the
    LSH attention layer, forwarding the attention mechanism, computing attention masks, hashing vectors, and other
    related operations.

    Attributes:
        config: Configuration parameters for the LSH self-attention layer.
        chunk_length: Length of each attention chunk.
        num_hashes: Number of hashes used in the LSH mechanism.
        num_buckets: Number of buckets used for hashing.
        num_chunks_before: Number of chunks to consider before the current chunk.
        num_chunks_after: Number of chunks to consider after the current chunk.
        hash_seed: Seed for random hash rotations.
        is_decoder: Flag indicating if the layer is used in a decoder.
        max_position_embeddings: Maximum position embeddings allowed.
        dropout: Dropout probability for attention weights.
        num_attention_heads: Number of attention heads.
        attention_head_size: Size of each attention head.
        all_head_size: Total size of all attention heads.
        hidden_size: Size of the hidden layer.
        query_key: Dense layer for query and key vectors.
        value: Dense layer for value vectors.
        self_mask_value_float16: Float16 value for masking in self-attention.
        self_mask_value_float32: Float32 value for masking in self-attention.
        mask_value_float16: Float16 value for general masking.
        mask_value_float32: Float32 value for general masking.

    Methods:
        forward: Constructs the LSH self-attention mechanism based on input hidden states and optional parameters.
        _query_per_attn_head: Computes query vectors per attention head.
        _value_per_attn_head: Computes value vectors per attention head.
        _hash_vectors: Hashes input vectors into buckets for attention computation.
        _get_sorted_bucket_idx_and_undo_sorted_bucket_idx: Computes sorted bucket indices for efficient attention
            calculation.
        _set_num_buckets: Sets the number of buckets based on the input sequence length.
        _attend: Computes attention scores and outputs based on query, key, and value vectors.
        _compute_attn_mask: Computes attention mask based on query and key indices.
        _get_relevant_hid_states_and_buckets: Retrieves relevant hidden states and buckets for efficient attention
            calculation.
        _expand_to_indices_in_relevant_chunk: Expands indices for relevant chunks in hidden states.
        _len_and_dim_norm: Normalizes vectors based on length and attention head size.
        _len_norm: Length normalization for input vectors.
        _gather_by_expansion: Expands indices and vectors for all hashes and gathers relevant elements.

    Note:
        This class is designed for implementing efficient self-attention mechanisms using Locality Sensitive Hashing.
    """
    def __init__(self, config):
        """
        Initializes the LSHSelfAttention class.

        Args:
            self (LSHSelfAttention): An instance of the LSHSelfAttention class.
            config (object): An object containing the configuration parameters.

        Returns:
            None

        Raises:
            None
        """
        super().__init__()
        self.config = config

        self.chunk_length = config.lsh_attn_chunk_length
        self.num_hashes = config.num_hashes
        self.num_buckets = config.num_buckets
        self.num_chunks_before = config.lsh_num_chunks_before
        self.num_chunks_after = config.lsh_num_chunks_after
        self.hash_seed = config.hash_seed
        self.is_decoder = config.is_decoder
        self.max_position_embeddings = config.max_position_embeddings

        self.dropout = config.lsh_attention_probs_dropout_prob

        self.num_attention_heads = config.num_attention_heads
        self.attention_head_size = config.attention_head_size
        self.all_head_size = self.num_attention_heads * self.attention_head_size
        self.hidden_size = config.hidden_size

        # projection matrices
        self.query_key = nn.Linear(self.hidden_size, self.all_head_size, bias=False)
        self.value = nn.Linear(self.hidden_size, self.all_head_size, bias=False)

        # save mask value here. Need fp32 and fp16 mask values
        self.self_mask_value_float16 = mindspore.tensor(-1e3)
        self.self_mask_value_float32 = mindspore.tensor(-1e5)
        self.mask_value_float16 = mindspore.tensor(-1e4)
        self.mask_value_float32 = mindspore.tensor(-1e9)

    def forward(
        self,
        hidden_states,
        attention_mask=None,
        head_mask=None,
        num_hashes=None,
        buckets=None,
        past_buckets_states=None,
        use_cache=False,
        output_attentions=False,
        **kwargs,
    ):
        """
        Constructs the LSH self-attention mechanism.

        Args:
            self: The object instance.
            hidden_states (torch.Tensor): The input hidden states with shape (batch_size, sequence_length, hidden_size).
            attention_mask (torch.Tensor, optional): Masking tensor for attention scores with shape
                (batch_size, sequence_length). Defaults to None.
            head_mask (torch.Tensor, optional): Masking tensor for individual attention heads with shape
                (num_attention_heads). Defaults to None.
            num_hashes (int, optional): Number of hash functions for LSH. Defaults to None.
            buckets (torch.Tensor, optional): Hash buckets for attention mechanism with shape
                (batch_size, num_attention_heads, num_hashes, sequence_length). Defaults to None.
            past_buckets_states (tuple, optional): Tuple containing past hash buckets and states. Defaults to None.
            use_cache (bool, optional): Flag to enable caching for attention mechanism. Defaults to False.
            output_attentions (bool, optional): Flag to output attention probabilities. Defaults to False.

        Returns:
            None

        Raises:
            AssertionError: If the input sequence length is not equal to 1 when `past_buckets_states` is passed.
            AssertionError: If the last dimension of query_key_vectors or value_vectors does not match the attention
                head size.
            AssertionError: If the last dimension of the buckets tensor does not match the expected value based on the
                number of hashes and sequence length.
            AssertionError: If the shape of out_vectors does not match the expected shape.
            AssertionError: If the configuration parameters are not set correctly when chunk_length is None.
        """
        sequence_length = hidden_states.shape[1]
        batch_size = hidden_states.shape[0]

        # num hashes can optionally be overwritten by user
        num_hashes = num_hashes if num_hashes is not None else self.num_hashes

        do_cached_attention = use_cache and past_buckets_states[1] is not None

        # check if cache shall be used and that hidden states are already cached
        if do_cached_attention:
            assert sequence_length == 1, (
                "At the moment, auto-regressive language generation is only possible one word at a time. Make sure"
                f" that input sequence length {sequence_length} equals 1, when `past_buckets_states` is passed."
            )
            past_buckets = past_buckets_states[0]
            past_states = past_buckets_states[1]

            # get query vector
            query_vectors = self.query_key(hidden_states)
            query_vectors = self._split_hidden_size_dim(
                query_vectors, self.num_attention_heads, self.attention_head_size
            )

            if past_buckets is not None:
                key_value_hidden_states, sorted_bucket_idx, buckets = self._get_relevant_hid_states_and_buckets(
                    query_vectors=query_vectors,
                    attention_mask=attention_mask,
                    num_hashes=num_hashes,
                    hidden_states=hidden_states,
                    past_states=past_states,
                    past_buckets=past_buckets,
                )

                query_key_vectors = self._query_per_attn_head(key_value_hidden_states)
                value_vectors = self._value_per_attn_head(key_value_hidden_states)

                # split key & value vectors by num hashes to apply
                # self attention on each separately
                query_key_vectors = self._split_seq_length_dim_to(
                    query_key_vectors,
                    num_hashes,
                    -1,
                    self.num_attention_heads,
                    self.attention_head_size,
                )
                value_vectors = self._split_seq_length_dim_to(
                    value_vectors,
                    num_hashes,
                    -1,
                    self.num_attention_heads,
                    self.attention_head_size,
                )
                # repeat query vectors across hash dimension
                query_vectors = query_vectors.unsqueeze(2).repeat(1, 1, num_hashes, 1, 1)
            else:
                key_value_hidden_states = ops.cat([past_states, hidden_states], axis=1)

                query_key_vectors = self.query_key(key_value_hidden_states)
                value_vectors = self.value(key_value_hidden_states)

        else:
            # project hidden_states to query_key and value
            query_vectors = None
            query_key_vectors = self.query_key(hidden_states)
            value_vectors = self.value(hidden_states)

        # if query key is not already split
        if not do_cached_attention or past_buckets is None:
            query_key_vectors = self._split_hidden_size_dim(
                query_key_vectors, self.num_attention_heads, self.attention_head_size
            )
            value_vectors = self._split_hidden_size_dim(
                value_vectors, self.num_attention_heads, self.attention_head_size
            )

        # cache buckets for next incremental decoding
        if do_cached_attention and past_buckets is None and key_value_hidden_states.shape[1] >= self.chunk_length:
            buckets = self._hash_vectors(query_key_vectors, num_hashes, attention_mask)

        # free memory
        del hidden_states

        assert (
            query_key_vectors.shape[-1] == self.attention_head_size
        ), f"last dim of query_key_vectors is {query_key_vectors.shape[-1]} but should be {self.attention_head_size}."
        assert (
            value_vectors.shape[-1] == self.attention_head_size
        ), f"last dim of value_vectors is {value_vectors.shape[-1]} but should be {self.attention_head_size}."

        do_standard_self_attention = (sequence_length <= self.chunk_length) or (
            use_cache and past_buckets_states[1] is not None
        )
        # LSH attention only makes sense if chunked attention should be performed
        if not do_standard_self_attention:
            # set `num_buckets` on the fly, recommended way to do it
            if self.num_buckets is None:
                self._set_num_buckets(sequence_length)

            # use cached buckets for backprop only
            if buckets is None:
                # hash query key vectors into buckets
                buckets = self._hash_vectors(query_key_vectors, num_hashes, attention_mask)
            else:
                # make sure buckets has correct shape for LSH attention
                buckets = buckets.view(batch_size, self.num_attention_heads, num_hashes * sequence_length)

            assert (
                int(buckets.shape[-1]) == num_hashes * sequence_length
            ), f"last dim of buckets is {buckets.shape[-1]}, but should be {num_hashes * sequence_length}"

            sorted_bucket_idx, undo_sorted_bucket_idx = self._get_sorted_bucket_idx_and_undo_sorted_bucket_idx(
                sequence_length, buckets, num_hashes
            )

            # make sure bucket idx is not longer then sequence length
            sorted_bucket_idx_per_hash = sorted_bucket_idx % sequence_length

            # cluster query key value vectors according to hashed buckets
            query_key_vectors = self._gather_by_expansion(query_key_vectors, sorted_bucket_idx_per_hash, num_hashes)
            value_vectors = self._gather_by_expansion(value_vectors, sorted_bucket_idx_per_hash, num_hashes)
            query_key_vectors = self._split_seq_length_dim_to(
                query_key_vectors,
                -1,
                self.chunk_length,
                self.num_attention_heads,
                self.attention_head_size,
            )
            value_vectors = self._split_seq_length_dim_to(
                value_vectors,
                -1,
                self.chunk_length,
                self.num_attention_heads,
                self.attention_head_size,
            )

            if self.chunk_length is None:
                assert self.num_chunks_before == 0 and self.num_chunks_after == 0, (
                    "If `config.chunk_length` is `None`, make sure `config.num_chunks_after` and"
                    " `config.num_chunks_before` are set to 0."
                )
        elif do_cached_attention and past_buckets is not None:
            # use max sequence length
            sorted_bucket_idx_per_hash = sorted_bucket_idx
        else:
            # get sequence length indices
            sorted_bucket_idx_per_hash = ops.arange(sequence_length).repeat(
                batch_size, self.num_attention_heads, 1
            )

        # scale key vectors
        sqrt_num = np.sqrt(self.attention_head_size)
        key_vectors = self._len_and_dim_norm(query_key_vectors, sqrt_num)

        # set query_vectors to query key vectors if LSH self attention
        query_vectors = query_vectors if query_vectors is not None else query_key_vectors

        # free memory
        del query_key_vectors

        # get attention probs
        out_vectors, logits, attention_probs = self._attend(
            query_vectors=query_vectors,
            key_vectors=key_vectors,
            value_vectors=value_vectors,
            sorted_bucket_idx_per_hash=sorted_bucket_idx_per_hash,
            attention_mask=attention_mask,
            head_mask=head_mask,
            do_standard_self_attention=do_standard_self_attention,
            do_cached_attention=do_cached_attention,
        )

        # free memory
        del key_vectors, value_vectors

        # re-order out_vectors and logits
        if not do_standard_self_attention:
            # sort clusters back to correct ordering
            out_vectors, logits = ReverseSort()(out_vectors, logits, sorted_bucket_idx, undo_sorted_bucket_idx)

        if not do_standard_self_attention or (do_cached_attention and past_buckets is not None):
            # sum up all hash rounds
            if num_hashes > 1:
                out_vectors = self._split_seq_length_dim_to(
                    out_vectors,
                    num_hashes,
                    sequence_length,
                    self.num_attention_heads,
                    self.attention_head_size,
                )
                logits = self._split_seq_length_dim_to(
                    logits,
                    num_hashes,
                    sequence_length,
                    self.num_attention_heads,
                    self.attention_head_size,
                ).unsqueeze(-1)

                probs_vectors = ops.exp(logits - ops.logsumexp(logits, axis=2, keep_dims=True))
                out_vectors = ops.sum(out_vectors * probs_vectors, dim=2)
                # free memory
                del probs_vectors

            # free memory
            del logits

        assert out_vectors.shape == (
            batch_size,
            self.num_attention_heads,
            sequence_length,
            self.attention_head_size,
        ), (
            "out_vectors have be of shape `[batch_size, config.num_attention_heads, sequence_length,"
            " config.attention_head_size]`."
        )

        out_vectors = self._merge_hidden_size_dims(out_vectors, self.num_attention_heads, self.attention_head_size)

        if output_attentions is False:
            attention_probs = ()

        if buckets is not None:
            buckets = buckets.view(batch_size, self.num_attention_heads, num_hashes, -1)

        return LSHSelfAttentionOutput(hidden_states=out_vectors, attention_probs=attention_probs, buckets=buckets)

    def _query_per_attn_head(self, hidden_states):
        """
        This method computes the query key vectors for each attention head.

        Args:
            self (LSHSelfAttention): The instance of LSHSelfAttention class.
            hidden_states (Tensor): The input hidden states with shape (batch_size, sequence_length, hidden_size).

        Returns:
            None: This method does not return any value but updates the query_key_vectors attribute of the
                LSHSelfAttention instance.

        Raises:
            ValueError: If the dimensions of the input hidden_states or query_key weight are incompatible
                for the matrix multiplication.
            RuntimeError: If the einsum operation fails due to any reason.
        """
        per_head_query_key = self.query_key.weight.reshape(
            self.num_attention_heads, self.attention_head_size, self.hidden_size
        ).swapaxes(-2, -1)
        # only relevant for inference and no bias => we can use einsum here
        query_key_vectors = ops.einsum("balh,ahr->balr", hidden_states, per_head_query_key)
        return query_key_vectors

    def _value_per_attn_head(self, hidden_states):
        """
        This method calculates the value vectors per attention head in the LSHSelfAttention class.

        Args:
            self (LSHSelfAttention): The instance of the LSHSelfAttention class.
            hidden_states (Tensor): The input hidden states tensor with shape [batch_size, sequence_length, hidden_size].

        Returns:
            None: This method does not return a value but updates the value vectors based on the input hidden states.

        Raises:
            ValueError: If the dimensions of the input hidden_states tensor are incorrect for the matrix multiplication.
            RuntimeError: If any runtime error occurs during the calculation process.
        """
        per_head_value = self.value.weight.reshape(
            self.num_attention_heads, self.attention_head_size, self.hidden_size
        ).swapaxes(-2, -1)
        # only relevant for inference and no bias => we can use einsum here
        value_vectors = ops.einsum("balh,ahr->balr", hidden_states, per_head_value)
        return value_vectors

    def _hash_vectors(self, vectors, num_hashes, attention_mask, increase_num_buckets=False):
        """
        This method '_hash_vectors' is defined in the class 'LSHSelfAttention', and it performs hashing operations
        on input vectors.

        Args:
            self: The instance of the LSHSelfAttention class.
            vectors (Tensor): Input vectors to be hashed. Shape should be (batch_size, sequence_length, hidden_size).
            num_hashes (int): Number of hashes to be generated.
            attention_mask (Tensor or None): Masking tensor to apply attention masking. Should be of shape
                (batch_size, sequence_length).
            increase_num_buckets (bool): Flag indicating whether to increase the number of buckets.

        Returns:
            None: This method does not return any value but updates the 'offset_buckets' attribute.

        Raises:
            AssertionError:
                - If the number of buckets is not an even integer.
                - If the input bucket factors are not even integers.
                - If the attention mask sum is less than the expected value based on batch size and sequence length.
            ValueError:
                - If the number of buckets is not an integer or a list of integers.
            IndexError: If the tensors' dimensions are incompatible for operations like indexing and concatenation.
            TypeError: If the data type of tensors is not supported for the operations performed in the method.
        """
        batch_size = vectors.shape[0]

        # See https://arxiv.org/pdf/1509.02897.pdf
        # We sample a different random rotation for each round of hashing to
        # decrease the probability of hash misses.
        if isinstance(self.num_buckets, int):
            assert (
                self.num_buckets % 2 == 0
            ), f"There should be an even number of buckets, but `self.num_buckets`: {self.num_buckets}"
            rotation_size = self.num_buckets
            num_buckets = self.num_buckets
        else:
            # Factorize the hash if self.num_buckets is a list or tuple
            rotation_size, num_buckets = 0, 1
            for bucket_factor in self.num_buckets:
                assert (
                    bucket_factor % 2 == 0
                ), f"The number of buckets should be even, but `num_bucket`: {bucket_factor}"
                rotation_size = rotation_size + bucket_factor
                num_buckets = num_buckets * bucket_factor

        if self.hash_seed is not None:
            # for determinism
            mindspore.set_seed(self.hash_seed)

        rotations_shape = (self.num_attention_heads, vectors.shape[-1], num_hashes, rotation_size // 2)
        # create a random self.attention_head_size x num_hashes x num_buckets/2
        random_rotations = ops.randn(rotations_shape, dtype=vectors.dtype)
        # Output dim: Batch_Size x Num_Attn_Heads x Num_Hashes x Seq_Len x Num_Buckets/2
        rotated_vectors = ops.einsum("bmtd,mdhr->bmhtr", vectors, random_rotations)

        if isinstance(self.num_buckets, int) or len(self.num_buckets) == 1:
            rotated_vectors = ops.cat([rotated_vectors, -rotated_vectors], axis=-1)
            buckets = ops.argmax(rotated_vectors, dim=-1)
        else:
            # Get the buckets for them and combine.
            buckets, cur_sum, cur_product = None, 0, 1
            for bucket_factor in self.num_buckets:
                rotated_vectors_factor = rotated_vectors[..., cur_sum : cur_sum + (bucket_factor // 2)]
                cur_sum = cur_sum + bucket_factor // 2
                rotated_vectors_factor = ops.cat([rotated_vectors_factor, -rotated_vectors_factor], axis=-1)
                if buckets is None:
                    buckets = ops.argmax(rotated_vectors_factor, dim=-1)
                else:
                    buckets = buckets + (cur_product * ops.argmax(rotated_vectors_factor, dim=-1))

                cur_product = cur_product * bucket_factor

        if attention_mask is not None and (attention_mask.sum().item() < batch_size * attention_mask.shape[-1]):
            # add an extra bucket for padding tokens only
            num_buckets = num_buckets + 1
            # assign padding tokens extra bucket
            buckets_mask = attention_mask.to(mindspore.bool_)[:, None, None, :].expand(buckets.shape)
            buckets = ops.where(
                buckets_mask, buckets, mindspore.tensor(num_buckets - 1, dtype=mindspore.int64)
            )
        elif increase_num_buckets:
            num_buckets = num_buckets + 1

        # buckets is now (Batch_size x Num_Attn_Heads x Num_Hashes x Seq_Len).
        # Next we add offsets so that bucket numbers from different hashing rounds don't overlap.
        offsets = ops.arange(num_hashes)
        offsets = (offsets * num_buckets).view((1, 1, -1, 1))

        # expand to batch size and num attention heads
        offsets = offsets.expand((batch_size, self.num_attention_heads) + offsets.shape[-2:])
        offset_buckets = (buckets + offsets).flatten(start_dim=2, end_dim=3)

        return offset_buckets

    def _get_sorted_bucket_idx_and_undo_sorted_bucket_idx(self, sequence_length, buckets, num_hashes):
        """
        Method to get the sorted bucket indices and create an undo mapping for sorting purposes in LSHSelfAttention class.

        Args:
            self: The instance of the LSHSelfAttention class.
            sequence_length (int): The length of the input sequence.
            buckets (Tensor): A tensor containing bucket values for each element in the input sequence.
            num_hashes (int): The number of hash functions used for bucketing.

        Returns:
            two tensors: sorted_bucket_idx and undo_sorted_bucket_idx.

        Raises:
            ValueError: If the input sequence length is not a positive integer.
            TypeError: If the buckets tensor is not a valid tensor object.
            ValueError: If the number of hashes is not a positive integer.
        """
        # hash-based sort
        sorted_bucket_idx = _stable_argsort(buckets, dim=-1)

        # create simple indices to scatter to, to have undo sort
        indices = (
            ops.arange(sorted_bucket_idx.shape[-1])
            .view(1, 1, -1)
            .expand(sorted_bucket_idx.shape)
        )

        # get undo sort
        undo_sorted_bucket_idx = sorted_bucket_idx.new_zeros(sorted_bucket_idx.shape).astype(indices.dtype)
        undo_sorted_bucket_idx = undo_sorted_bucket_idx.scatter(-1, sorted_bucket_idx, indices)

        return sorted_bucket_idx, undo_sorted_bucket_idx

    def _set_num_buckets(self, sequence_length):
        """
        This method _set_num_buckets calculates the number of buckets to be used for locality-sensitive hashing (LSH)
        self-attention in the LSHSelfAttention class.

        Args:
            self (LSHSelfAttention): The instance of the LSHSelfAttention class.
            sequence_length (int): The total length of the input sequence.

        Returns:
            None: This method does not return any value. It sets the calculated number of buckets in the configuration
                and instance variable.

        Raises:
            None.
        """
        # `num_buckets` should be set to 2 * sequence_length // chunk_length as recommended in paper
        num_buckets_pow_2 = (2 * (sequence_length // self.chunk_length)).bit_length() - 1
        # make sure buckets are power of 2
        num_buckets = 2**num_buckets_pow_2

        # factorize `num_buckets` if `num_buckets` becomes too large
        num_buckets_limit = 2 * max(
            int((self.max_position_embeddings // self.chunk_length) ** (0.5)),
            self.chunk_length,
        )
        if num_buckets > num_buckets_limit:
            num_buckets = [2 ** (num_buckets_pow_2 // 2), 2 ** (num_buckets_pow_2 - num_buckets_pow_2 // 2)]

        logger.warning(f"config.num_buckets is not set. Setting config.num_buckets to {num_buckets}...")

        # set num buckets in config to be properly saved
        self.config.num_buckets = num_buckets
        self.num_buckets = num_buckets

    def _attend(
        self,
        query_vectors,
        key_vectors,
        value_vectors,
        sorted_bucket_idx_per_hash,
        attention_mask,
        head_mask,
        do_standard_self_attention,
        do_cached_attention,
    ):
        '''
        This method performs LSH (Locality Sensitive Hashing) self-attention calculation for the LSHSelfAttention class.

        Args:
            self (LSHSelfAttention): The LSHSelfAttention object.
            query_vectors (Tensor): The input query vectors for the attention calculation.
            key_vectors (Tensor): The input key vectors for the attention calculation.
            value_vectors (Tensor): The input value vectors for the attention calculation.
            sorted_bucket_idx_per_hash (Tensor): The sorted bucket indices per hash for the attention calculation.
            attention_mask (Tensor): The attention mask to be applied in the calculation.
            head_mask (Tensor, optional): The optional mask to be applied to the attention scores.
            do_standard_self_attention (bool): A flag indicating whether to perform standard self-attention calculation.
            do_cached_attention (bool): A flag indicating whether to use cached attention for the calculation.

        Returns:
            None.

        Raises:
            ValueError: If the input tensors do not have the expected shapes or types.
            RuntimeError: If there is an issue with the computation during the method execution.
        '''
        # look at previous and following chunks if chunked attention
        if not do_standard_self_attention:
            key_vectors = self._look_adjacent(key_vectors, self.num_chunks_before, self.num_chunks_after)
            value_vectors = self._look_adjacent(value_vectors, self.num_chunks_before, self.num_chunks_after)

        # get logits and dots
        # (BS, NumAttn, NumHash x NumChunk, Chunk_L x Hidden),(BS, NumAttn, NumHash x NumChunk, Chunk_L * (Num_bef + Num_aft + 1) x Hidden) -> (BS, NumAttn, NumHash x NumChunk, Chunk_L, Chunk_L * (1 + Num_bef+ Num_aft))
        query_key_dots = ops.matmul(query_vectors, key_vectors.swapaxes(-1, -2))

        # free memory
        del query_vectors, key_vectors

        # if chunked attention split bucket idxs to query and key
        if not do_standard_self_attention:
            query_bucket_idx = self._split_seq_length_dim_to(
                sorted_bucket_idx_per_hash, -1, self.chunk_length, self.num_attention_heads
            )
            key_value_bucket_idx = self._look_adjacent(query_bucket_idx, self.num_chunks_before, self.num_chunks_after)
        elif do_cached_attention and query_key_dots.ndim > 4:
            key_value_bucket_idx = sorted_bucket_idx_per_hash
            query_bucket_idx = (
                key_value_bucket_idx.new_ones(key_value_bucket_idx.shape[:-1] + (1,)) * key_value_bucket_idx.max()
            )
        elif do_cached_attention and query_key_dots.ndim <= 4:
            query_bucket_idx = (query_key_dots.shape[-1] - 1) * ops.ones_like(query_key_dots)[:, :, :, -1]
            key_value_bucket_idx = ops.arange(
                query_key_dots.shape[-1], dtype=mindspore.int64
            )[None, None, :].expand(query_bucket_idx.shape[:2] + (-1,))
        else:
            query_bucket_idx = key_value_bucket_idx = sorted_bucket_idx_per_hash

        # get correct mask values depending on precision
        if query_key_dots.dtype == mindspore.float16:
            self_mask_value = self.self_mask_value_float16.half()
            mask_value = self.mask_value_float16.half()
        else:
            self_mask_value = self.self_mask_value_float32
            mask_value = self.mask_value_float32

        if not do_cached_attention:
            mask = self._compute_attn_mask(
                query_bucket_idx,
                key_value_bucket_idx,
                attention_mask,
                query_key_dots.shape,
                do_standard_self_attention,
            )

            if mask is not None:
                query_key_dots = ops.where(mask, query_key_dots, mask_value)

            # free memory
            del mask

        # Self mask is ALWAYS applied.
        # From the reformer paper (https://arxiv.org/pdf/2001.04451.pdf):
        # " While attention to the future is not allowed, typical implementations of the
        # Transformer do allow a position to attend to itself.
        # Such behavior is undesirable in a shared-QK formulation because the dot-product
        # of a query vector with itself will almost always be greater than the dot product of a
        # query vector with a vector at another position. We therefore modify the masking
        # to forbid a token from attending to itself, except in situations
        # where a token has no other valid attention targets (e.g. the first token in a sequence) "

        self_mask = ops.ne(query_bucket_idx.unsqueeze(-1), key_value_bucket_idx.unsqueeze(-2))

        # apply self_mask
        query_key_dots = ops.where(self_mask, query_key_dots, self_mask_value)

        # free memory
        del self_mask

        logits = ops.logsumexp(query_key_dots, axis=-1, keep_dims=True)
        # dots shape is `[batch_size, num_attn_heads, num_hashes * seq_len // chunk_length, chunk_length, chunk_length * (1 + num_chunks_before + num_chunks_after)]`
        attention_probs = ops.exp(query_key_dots - logits)

        # free memory
        del query_key_dots

        # dropout
        attention_probs = ops.dropout(attention_probs, p=self.dropout, training=self.training)

        # Mask heads if we want to
        if head_mask is not None:
            attention_probs = attention_probs * head_mask

        # attend values
        out_vectors = ops.matmul(attention_probs, value_vectors)

        # free memory
        del value_vectors

        # merge chunk length
        if out_vectors.ndim > 4:
            logits = logits.flatten(start_dim=2, end_dim=3).squeeze(-1)
            out_vectors = out_vectors.flatten(start_dim=2, end_dim=3)

        return out_vectors, logits, attention_probs

    def _compute_attn_mask(
        self, query_indices, key_indices, attention_mask, query_key_dot_shape, do_standard_self_attention
    ):

        """
        Compute attention mask for LSH self-attention.

        This method computes the attention mask for LSH self-attention based on the given query and key indices,
        attention mask, query-key dot shape, and the flag indicating whether to use standard self-attention.

        Args:
            self (LSHSelfAttention): An instance of the LSHSelfAttention class.
            query_indices (Tensor): The indices of the query.
            key_indices (Tensor): The indices of the key.
            attention_mask (Tensor): The attention mask to be applied (optional).
            query_key_dot_shape (tuple): The shape of the dot product between query and key.
            do_standard_self_attention (bool): A flag indicating whether to use standard self-attention.

        Returns:
            Tensor: The computed attention mask.

        Raises:
            TypeError: If the attention_mask is not of type Tensor.
        """
        # attention mask for LSH
        if attention_mask is not None:
            # if chunked attention, the attention mask has to correspond to LSH order
            attention_mask = attention_mask.to(mindspore.bool_)[:, None, :]
            if not do_standard_self_attention:
                # expand attn_mask to fit with key_value_bucket_idx shape
                attention_mask = attention_mask[:, None, :]
                attention_mask = attention_mask.expand(query_indices.shape[:-1] + (-1,))
                # extract attention mask from LSH sorted key_indices
                attention_mask = ops.gather_elements(attention_mask, -1, key_indices)

            attention_mask = attention_mask.unsqueeze(-2).expand(query_key_dot_shape)

        # Causal mask
        if self.is_decoder is True:
            causal_mask = ops.ge(query_indices.unsqueeze(-1), key_indices.unsqueeze(-2))

            # add attention mask if not None
            if attention_mask is not None:
                attention_mask = causal_mask & attention_mask
            else:
                attention_mask = causal_mask

        return attention_mask

    def _get_relevant_hid_states_and_buckets(
        self, query_vectors, attention_mask, num_hashes, hidden_states, past_states, past_buckets
    ):

        """
        This method '_get_relevant_hid_states_and_buckets' is defined in the class 'LSHSelfAttention'.

        Args:
            self (LSHSelfAttention): The instance of the LSHSelfAttention class.
            query_vectors (Tensor): The input query vectors for which relevant hidden states and buckets need
                to be retrieved.
            attention_mask (Tensor): The attention mask to be applied during the computation.
            num_hashes (int): The number of hashes to be used in the computation.
            hidden_states (Tensor): The current hidden states that include both past and new states.
            past_states (Tensor): The past hidden states that need to be included in the computation.
            past_buckets (Tensor): The past bucket values that need to be considered.

        Returns:
            None: This method does not return any value but updates the relevant hidden states and buckets.

        Raises:
            AssertionError:

                - If the shape of 'bucket_idx' does not match the expected shape.
                - If the shape of 'relevant_hidden_states' or 'relevant_bucket_idx_chunk' does not match the
                expected shapes.
                - If the number of 'hidden_states' or 'bucket_idx' does not match the expected values.
            IndexError: If there is an index error during the computation.
        """
        # concat hidden states
        hidden_states = ops.cat([past_states, hidden_states], axis=1)

        # batch_size hidden
        batch_size = hidden_states.shape[0]
        sequence_length = hidden_states.shape[1]

        # check if cached buckets include pad bucket
        max_bucket = self.num_buckets if isinstance(self.num_buckets, int) else reduce(mul, self.num_buckets)

        # if pad bucket was cached => need to increase num buckets for caching
        increase_num_buckets = past_buckets.max() > num_hashes * max_bucket - 1

        # retrieve query buckets
        query_buckets = self._hash_vectors(
            query_vectors, num_hashes, attention_mask, increase_num_buckets=increase_num_buckets
        )

        # concat buckets
        concat_buckets = ops.cat([past_buckets, query_buckets.unsqueeze(-1)], axis=-1)

        # hash-based sort
        bucket_idx = _stable_argsort(concat_buckets, dim=-1)

        # bucket_idx has shape: BatchSize x NumAttnHeads x NumHashes x SequenceLength
        assert bucket_idx.shape == (
            batch_size,
            self.num_attention_heads,
            num_hashes,
            sequence_length,
        ), (
            f"bucket_idx should have shape {(batch_size, self.num_attention_heads, num_hashes, sequence_length)}, but"
            f" has shape {bucket_idx.shape}."
        )

        # find indices of new bucket indices
        relevant_bucket_idx = (bucket_idx == (bucket_idx.shape[-1] - 1)).nonzero()

        # expand relevant bucket indices to its chunks
        relevant_bucket_idx_chunk = self._expand_to_indices_in_relevant_chunk(relevant_bucket_idx, sequence_length)
        relevant_bucket_idx_chunk = bucket_idx[tuple(relevant_bucket_idx_chunk.swapaxes(0, 1))]

        # adapt bucket_idx for batch and hidden states for index select
        offset = ops.arange(relevant_bucket_idx_chunk.shape[-1], dtype=mindspore.int64)
        bucket_idx_batch_offset = sequence_length * (
            batch_size * ops.div(offset, relevant_bucket_idx_chunk.shape[-1], rounding_mode="floor")
        )

        # add batch offset
        relevant_bucket_idx_chunk_all_batch = relevant_bucket_idx_chunk + bucket_idx_batch_offset
        hidden_states = hidden_states.reshape((-1, self.hidden_size))

        # select all relevant hidden states
        relevant_hidden_states = hidden_states.index_select(0, relevant_bucket_idx_chunk_all_batch)

        # reshape hidden states and bucket_idx to correct output
        relevant_hidden_states = relevant_hidden_states.reshape(
            batch_size, self.num_attention_heads, -1, self.hidden_size
        )
        relevant_bucket_idx_chunk = relevant_bucket_idx_chunk.reshape(
            batch_size, self.num_attention_heads, num_hashes, -1
        )

        assert (
            relevant_hidden_states.shape[2]
            == (self.num_chunks_before + self.num_chunks_after + 1) * self.chunk_length * num_hashes
        ), (
            "There should be"
            f" {(self.num_chunks_before + self.num_chunks_after + 1) * self.chunk_length * num_hashes} `hidden_states`,"
            f" there are {relevant_hidden_states.shape[2]} `hidden_states`."
        )

        assert (
            relevant_bucket_idx_chunk.shape[-1]
            == (self.num_chunks_before + self.num_chunks_after + 1) * self.chunk_length
        ), (
            "There should be"
            f" {(self.num_chunks_before + self.num_chunks_after + 1) * self.chunk_length} `hidden_states`, there are"
            f" {relevant_bucket_idx_chunk.shape[-1]} `bucket_idx`."
        )

        return relevant_hidden_states, relevant_bucket_idx_chunk, query_buckets

    def _expand_to_indices_in_relevant_chunk(self, indices, sequence_length):

        """
        This method '_expand_to_indices_in_relevant_chunk' is defined in the class 'LSHSelfAttention'.

        Args:
            self:
                An instance of the LSHSelfAttention class.

                - Purpose: Represents the current instance of the LSHSelfAttention class.
                - Restrictions: None.

            indices:
                A tensor containing indices.

                - Type: Tensor
                - Purpose: Specifies the indices to be expanded within the relevant chunk.
                - Restrictions: Should be a 2D tensor.

            sequence_length:
                An integer specifying the total length of the sequence.

                - Type: int
                - Purpose: Indicates the total length of the sequence for boundary calculations.
                - Restrictions: Must be a positive integer.

        Returns:
            None:

                - Type: None
                - Purpose: The method does not return any value, it modifies the 'indices' tensor in place.

        Raises:
            None.
        """
        # get relevant indices of where chunk starts and its size
        start_indices_chunk = ((indices[:, -1] // self.chunk_length) - self.num_chunks_before) * self.chunk_length
        total_chunk_size = self.chunk_length * (1 + self.num_chunks_before + self.num_chunks_after)

        # expand start indices and add correct chunk offset via arange
        expanded_start_indices = start_indices_chunk.unsqueeze(-1).expand(indices.shape[0], total_chunk_size)
        chunk_sequence_indices = expanded_start_indices + ops.arange(
            total_chunk_size, dtype=mindspore.int64
        ).unsqueeze(0).expand(indices.shape[0], total_chunk_size)

        # make sure that circular logic holds via % seq len
        chunk_sequence_indices = chunk_sequence_indices.flatten() % sequence_length

        # expand indices and set indices correctly
        indices = indices.unsqueeze(1).expand((indices.shape[0], total_chunk_size, -1)).flatten(start_dim=0, end_dim=1).copy()
        indices[:, -1] = chunk_sequence_indices

        return indices

    def _len_and_dim_norm(self, vectors, sqrt_num):
        """
        length and attention head size dim normalization
        """
        vectors = self._len_norm(vectors)
        vectors = vectors / sqrt_num
        return vectors

    def _len_norm(self, x, eps=1e-6):
        """
        length normalization
        """
        variance = ops.mean(x**2, -1, keep_dims=True)
        norm_x = x * ops.rsqrt(variance + epsilon)
        return norm_x

    def _gather_by_expansion(self, vectors, idxs, num_hashes):
        """
        expand dims of idxs and vectors for all hashes and gather
        """
        expanded_idxs = idxs.unsqueeze(-1).expand(-1, -1, -1, self.attention_head_size)
        vectors = vectors.repeat(1, 1, num_hashes, 1)
        return ops.gather_elements(vectors, 2, expanded_idxs)

mindnlp.transformers.models.reformer.modeling_reformer.LSHSelfAttention.__init__(config)

Initializes the LSHSelfAttention class.

PARAMETER	DESCRIPTION
self	An instance of the LSHSelfAttention class. TYPE: LSHSelfAttention
config	An object containing the configuration parameters. TYPE: object

RETURNS	DESCRIPTION
	None

Source code in mindnlp/transformers/models/reformer/modeling_reformer.py

def __init__(self, config):
    """
    Initializes the LSHSelfAttention class.

    Args:
        self (LSHSelfAttention): An instance of the LSHSelfAttention class.
        config (object): An object containing the configuration parameters.

    Returns:
        None

    Raises:
        None
    """
    super().__init__()
    self.config = config

    self.chunk_length = config.lsh_attn_chunk_length
    self.num_hashes = config.num_hashes
    self.num_buckets = config.num_buckets
    self.num_chunks_before = config.lsh_num_chunks_before
    self.num_chunks_after = config.lsh_num_chunks_after
    self.hash_seed = config.hash_seed
    self.is_decoder = config.is_decoder
    self.max_position_embeddings = config.max_position_embeddings

    self.dropout = config.lsh_attention_probs_dropout_prob

    self.num_attention_heads = config.num_attention_heads
    self.attention_head_size = config.attention_head_size
    self.all_head_size = self.num_attention_heads * self.attention_head_size
    self.hidden_size = config.hidden_size

    # projection matrices
    self.query_key = nn.Linear(self.hidden_size, self.all_head_size, bias=False)
    self.value = nn.Linear(self.hidden_size, self.all_head_size, bias=False)

    # save mask value here. Need fp32 and fp16 mask values
    self.self_mask_value_float16 = mindspore.tensor(-1e3)
    self.self_mask_value_float32 = mindspore.tensor(-1e5)
    self.mask_value_float16 = mindspore.tensor(-1e4)
    self.mask_value_float32 = mindspore.tensor(-1e9)

mindnlp.transformers.models.reformer.modeling_reformer.LSHSelfAttention.forward(hidden_states, attention_mask=None, head_mask=None, num_hashes=None, buckets=None, past_buckets_states=None, use_cache=False, output_attentions=False, **kwargs)

Constructs the LSH self-attention mechanism.

PARAMETER	DESCRIPTION
self	The object instance.
hidden_states	The input hidden states with shape (batch_size, sequence_length, hidden_size). TYPE: Tensor
attention_mask	Masking tensor for attention scores with shape (batch_size, sequence_length). Defaults to None. TYPE: Tensor DEFAULT: None
head_mask	Masking tensor for individual attention heads with shape (num_attention_heads). Defaults to None. TYPE: Tensor DEFAULT: None
num_hashes	Number of hash functions for LSH. Defaults to None. TYPE: int DEFAULT: None
buckets	Hash buckets for attention mechanism with shape (batch_size, num_attention_heads, num_hashes, sequence_length). Defaults to None. TYPE: Tensor DEFAULT: None
past_buckets_states	Tuple containing past hash buckets and states. Defaults to None. TYPE: tuple DEFAULT: None
use_cache	Flag to enable caching for attention mechanism. Defaults to False. TYPE: bool DEFAULT: False
output_attentions	Flag to output attention probabilities. Defaults to False. TYPE: bool DEFAULT: False

RETURNS	DESCRIPTION
	None

RAISES	DESCRIPTION
AssertionError	If the input sequence length is not equal to 1 when past_buckets_states is passed.
AssertionError	If the last dimension of query_key_vectors or value_vectors does not match the attention head size.
AssertionError	If the last dimension of the buckets tensor does not match the expected value based on the number of hashes and sequence length.
AssertionError	If the shape of out_vectors does not match the expected shape.
AssertionError	If the configuration parameters are not set correctly when chunk_length is None.

Source code in mindnlp/transformers/models/reformer/modeling_reformer.py

def forward(
    self,
    hidden_states,
    attention_mask=None,
    head_mask=None,
    num_hashes=None,
    buckets=None,
    past_buckets_states=None,
    use_cache=False,
    output_attentions=False,
    **kwargs,
):
    """
    Constructs the LSH self-attention mechanism.

    Args:
        self: The object instance.
        hidden_states (torch.Tensor): The input hidden states with shape (batch_size, sequence_length, hidden_size).
        attention_mask (torch.Tensor, optional): Masking tensor for attention scores with shape
            (batch_size, sequence_length). Defaults to None.
        head_mask (torch.Tensor, optional): Masking tensor for individual attention heads with shape
            (num_attention_heads). Defaults to None.
        num_hashes (int, optional): Number of hash functions for LSH. Defaults to None.
        buckets (torch.Tensor, optional): Hash buckets for attention mechanism with shape
            (batch_size, num_attention_heads, num_hashes, sequence_length). Defaults to None.
        past_buckets_states (tuple, optional): Tuple containing past hash buckets and states. Defaults to None.
        use_cache (bool, optional): Flag to enable caching for attention mechanism. Defaults to False.
        output_attentions (bool, optional): Flag to output attention probabilities. Defaults to False.

    Returns:
        None

    Raises:
        AssertionError: If the input sequence length is not equal to 1 when `past_buckets_states` is passed.
        AssertionError: If the last dimension of query_key_vectors or value_vectors does not match the attention
            head size.
        AssertionError: If the last dimension of the buckets tensor does not match the expected value based on the
            number of hashes and sequence length.
        AssertionError: If the shape of out_vectors does not match the expected shape.
        AssertionError: If the configuration parameters are not set correctly when chunk_length is None.
    """
    sequence_length = hidden_states.shape[1]
    batch_size = hidden_states.shape[0]

    # num hashes can optionally be overwritten by user
    num_hashes = num_hashes if num_hashes is not None else self.num_hashes

    do_cached_attention = use_cache and past_buckets_states[1] is not None

    # check if cache shall be used and that hidden states are already cached
    if do_cached_attention:
        assert sequence_length == 1, (
            "At the moment, auto-regressive language generation is only possible one word at a time. Make sure"
            f" that input sequence length {sequence_length} equals 1, when `past_buckets_states` is passed."
        )
        past_buckets = past_buckets_states[0]
        past_states = past_buckets_states[1]

        # get query vector
        query_vectors = self.query_key(hidden_states)
        query_vectors = self._split_hidden_size_dim(
            query_vectors, self.num_attention_heads, self.attention_head_size
        )

        if past_buckets is not None:
            key_value_hidden_states, sorted_bucket_idx, buckets = self._get_relevant_hid_states_and_buckets(
                query_vectors=query_vectors,
                attention_mask=attention_mask,
                num_hashes=num_hashes,
                hidden_states=hidden_states,
                past_states=past_states,
                past_buckets=past_buckets,
            )

            query_key_vectors = self._query_per_attn_head(key_value_hidden_states)
            value_vectors = self._value_per_attn_head(key_value_hidden_states)

            # split key & value vectors by num hashes to apply
            # self attention on each separately
            query_key_vectors = self._split_seq_length_dim_to(
                query_key_vectors,
                num_hashes,
                -1,
                self.num_attention_heads,
                self.attention_head_size,
            )
            value_vectors = self._split_seq_length_dim_to(
                value_vectors,
                num_hashes,
                -1,
                self.num_attention_heads,
                self.attention_head_size,
            )
            # repeat query vectors across hash dimension
            query_vectors = query_vectors.unsqueeze(2).repeat(1, 1, num_hashes, 1, 1)
        else:
            key_value_hidden_states = ops.cat([past_states, hidden_states], axis=1)

            query_key_vectors = self.query_key(key_value_hidden_states)
            value_vectors = self.value(key_value_hidden_states)

    else:
        # project hidden_states to query_key and value
        query_vectors = None
        query_key_vectors = self.query_key(hidden_states)
        value_vectors = self.value(hidden_states)

    # if query key is not already split
    if not do_cached_attention or past_buckets is None:
        query_key_vectors = self._split_hidden_size_dim(
            query_key_vectors, self.num_attention_heads, self.attention_head_size
        )
        value_vectors = self._split_hidden_size_dim(
            value_vectors, self.num_attention_heads, self.attention_head_size
        )

    # cache buckets for next incremental decoding
    if do_cached_attention and past_buckets is None and key_value_hidden_states.shape[1] >= self.chunk_length:
        buckets = self._hash_vectors(query_key_vectors, num_hashes, attention_mask)

    # free memory
    del hidden_states

    assert (
        query_key_vectors.shape[-1] == self.attention_head_size
    ), f"last dim of query_key_vectors is {query_key_vectors.shape[-1]} but should be {self.attention_head_size}."
    assert (
        value_vectors.shape[-1] == self.attention_head_size
    ), f"last dim of value_vectors is {value_vectors.shape[-1]} but should be {self.attention_head_size}."

    do_standard_self_attention = (sequence_length <= self.chunk_length) or (
        use_cache and past_buckets_states[1] is not None
    )
    # LSH attention only makes sense if chunked attention should be performed
    if not do_standard_self_attention:
        # set `num_buckets` on the fly, recommended way to do it
        if self.num_buckets is None:
            self._set_num_buckets(sequence_length)

        # use cached buckets for backprop only
        if buckets is None:
            # hash query key vectors into buckets
            buckets = self._hash_vectors(query_key_vectors, num_hashes, attention_mask)
        else:
            # make sure buckets has correct shape for LSH attention
            buckets = buckets.view(batch_size, self.num_attention_heads, num_hashes * sequence_length)

        assert (
            int(buckets.shape[-1]) == num_hashes * sequence_length
        ), f"last dim of buckets is {buckets.shape[-1]}, but should be {num_hashes * sequence_length}"

        sorted_bucket_idx, undo_sorted_bucket_idx = self._get_sorted_bucket_idx_and_undo_sorted_bucket_idx(
            sequence_length, buckets, num_hashes
        )

        # make sure bucket idx is not longer then sequence length
        sorted_bucket_idx_per_hash = sorted_bucket_idx % sequence_length

        # cluster query key value vectors according to hashed buckets
        query_key_vectors = self._gather_by_expansion(query_key_vectors, sorted_bucket_idx_per_hash, num_hashes)
        value_vectors = self._gather_by_expansion(value_vectors, sorted_bucket_idx_per_hash, num_hashes)
        query_key_vectors = self._split_seq_length_dim_to(
            query_key_vectors,
            -1,
            self.chunk_length,
            self.num_attention_heads,
            self.attention_head_size,
        )
        value_vectors = self._split_seq_length_dim_to(
            value_vectors,
            -1,
            self.chunk_length,
            self.num_attention_heads,
            self.attention_head_size,
        )

        if self.chunk_length is None:
            assert self.num_chunks_before == 0 and self.num_chunks_after == 0, (
                "If `config.chunk_length` is `None`, make sure `config.num_chunks_after` and"
                " `config.num_chunks_before` are set to 0."
            )
    elif do_cached_attention and past_buckets is not None:
        # use max sequence length
        sorted_bucket_idx_per_hash = sorted_bucket_idx
    else:
        # get sequence length indices
        sorted_bucket_idx_per_hash = ops.arange(sequence_length).repeat(
            batch_size, self.num_attention_heads, 1
        )

    # scale key vectors
    sqrt_num = np.sqrt(self.attention_head_size)
    key_vectors = self._len_and_dim_norm(query_key_vectors, sqrt_num)

    # set query_vectors to query key vectors if LSH self attention
    query_vectors = query_vectors if query_vectors is not None else query_key_vectors

    # free memory
    del query_key_vectors

    # get attention probs
    out_vectors, logits, attention_probs = self._attend(
        query_vectors=query_vectors,
        key_vectors=key_vectors,
        value_vectors=value_vectors,
        sorted_bucket_idx_per_hash=sorted_bucket_idx_per_hash,
        attention_mask=attention_mask,
        head_mask=head_mask,
        do_standard_self_attention=do_standard_self_attention,
        do_cached_attention=do_cached_attention,
    )

    # free memory
    del key_vectors, value_vectors

    # re-order out_vectors and logits
    if not do_standard_self_attention:
        # sort clusters back to correct ordering
        out_vectors, logits = ReverseSort()(out_vectors, logits, sorted_bucket_idx, undo_sorted_bucket_idx)

    if not do_standard_self_attention or (do_cached_attention and past_buckets is not None):
        # sum up all hash rounds
        if num_hashes > 1:
            out_vectors = self._split_seq_length_dim_to(
                out_vectors,
                num_hashes,
                sequence_length,
                self.num_attention_heads,
                self.attention_head_size,
            )
            logits = self._split_seq_length_dim_to(
                logits,
                num_hashes,
                sequence_length,
                self.num_attention_heads,
                self.attention_head_size,
            ).unsqueeze(-1)

            probs_vectors = ops.exp(logits - ops.logsumexp(logits, axis=2, keep_dims=True))
            out_vectors = ops.sum(out_vectors * probs_vectors, dim=2)
            # free memory
            del probs_vectors

        # free memory
        del logits

    assert out_vectors.shape == (
        batch_size,
        self.num_attention_heads,
        sequence_length,
        self.attention_head_size,
    ), (
        "out_vectors have be of shape `[batch_size, config.num_attention_heads, sequence_length,"
        " config.attention_head_size]`."
    )

    out_vectors = self._merge_hidden_size_dims(out_vectors, self.num_attention_heads, self.attention_head_size)

    if output_attentions is False:
        attention_probs = ()

    if buckets is not None:
        buckets = buckets.view(batch_size, self.num_attention_heads, num_hashes, -1)

    return LSHSelfAttentionOutput(hidden_states=out_vectors, attention_probs=attention_probs, buckets=buckets)

mindnlp.transformers.models.reformer.modeling_reformer.LocalSelfAttention

Bases: Module, EfficientAttentionMixin

The LocalSelfAttention class is a subclass of nn.Module and EfficientAttentionMixin that represents a local self-attention mechanism. This mechanism is commonly used in transformer-based models for processing sequential data.

ATTRIBUTE	DESCRIPTION
`num_attention_heads`	The number of attention heads. TYPE: int
`chunk_length`	The length of each attention chunk. TYPE: int
`num_chunks_before`	The number of chunks before the current position. TYPE: int
`num_chunks_after`	The number of chunks after the current position. TYPE: int
`is_decoder`	Indicates whether the attention is used in a decoder architecture. TYPE: bool
`pad_token_id`	The token ID used for padding. TYPE: int
`attention_head_size`	The size of each attention head. TYPE: int
`all_head_size`	The total size of all attention heads. TYPE: int
`hidden_size`	The hidden size of the input. TYPE: int
`query`	The dense layer used for computing query vectors. TYPE: Linear
`key`	The dense layer used for computing key vectors. TYPE: Linear
`value`	The dense layer used for computing value vectors. TYPE: Linear
`dropout`	The dropout rate for attention probabilities. TYPE: float
`mask_value_float16`	The mask value for float16 data type. TYPE: tensor
`mask_value_float32`	The mask value for float32 data type. TYPE: tensor

METHOD	DESCRIPTION
`forward`	Computes the local self-attention mechanism for the given hidden states.
`_compute_attn_mask`	Computes the attention mask based on query and key indices.
`_retrieve_relevant_hidden_states`	Retrieves the relevant hidden states from previous states.

Source code in mindnlp/transformers/models/reformer/modeling_reformer.py

class LocalSelfAttention(nn.Module, EfficientAttentionMixin):

    """
    The `LocalSelfAttention` class is a subclass of `nn.Module` and `EfficientAttentionMixin` that represents
    a local self-attention mechanism. This mechanism is commonly used in transformer-based models for
    processing sequential data.

    Attributes:
        `num_attention_heads` (int): The number of attention heads.
        `chunk_length` (int): The length of each attention chunk.
        `num_chunks_before` (int): The number of chunks before the current position.
        `num_chunks_after` (int): The number of chunks after the current position.
        `is_decoder` (bool): Indicates whether the attention is used in a decoder architecture.
        `pad_token_id` (int): The token ID used for padding.
        `attention_head_size` (int): The size of each attention head.
        `all_head_size` (int): The total size of all attention heads.
        `hidden_size` (int): The hidden size of the input.
        `query` (nn.Linear): The dense layer used for computing query vectors.
        `key` (nn.Linear): The dense layer used for computing key vectors.
        `value` (nn.Linear): The dense layer used for computing value vectors.
        `dropout` (float): The dropout rate for attention probabilities.
        `mask_value_float16` (mindspore.tensor): The mask value for float16 data type.
        `mask_value_float32` (mindspore.tensor): The mask value for float32 data type.

    Methods:
        `forward`: Computes the local self-attention mechanism for the given hidden states.
        `_compute_attn_mask`: Computes the attention mask based on query and key indices.
        `_retrieve_relevant_hidden_states`: Retrieves the relevant hidden states from previous states.

    """
    def __init__(self, config):

        """
        Initializes the LocalSelfAttention class.

        Args:
            self: The instance of the LocalSelfAttention class.
            config: An object containing configuration parameters for the attention mechanism.
                This parameter is expected to have the following attributes:

                - num_attention_heads (int): The number of attention heads.
                - local_attn_chunk_length (int): The length of attention chunks.
                - local_num_chunks_before (int): The number of chunks before the current position.
                - local_num_chunks_after (int): The number of chunks after the current position.
                - is_decoder (bool): Indicates if the attention mechanism is used in a decoder.
                - pad_token_id (int): The token ID used for padding.
                - attention_head_size (int): The size of each attention head.
                - hidden_size (int): The size of the hidden layer.
                - local_attention_probs_dropout_prob (float): The dropout probability for attention weights.

        Returns:
            None.

        Raises:
            ValueError: If the configuration parameters provided are invalid or missing.
            TypeError: If any of the configuration attribute types are incorrect.
        """
        super().__init__()

        self.num_attention_heads = config.num_attention_heads
        self.chunk_length = config.local_attn_chunk_length
        self.num_chunks_before = config.local_num_chunks_before
        self.num_chunks_after = config.local_num_chunks_after
        self.is_decoder = config.is_decoder
        self.pad_token_id = config.pad_token_id

        self.attention_head_size = config.attention_head_size
        self.all_head_size = self.num_attention_heads * self.attention_head_size
        self.hidden_size = config.hidden_size

        # projection matrices
        self.query = nn.Linear(self.hidden_size, self.all_head_size, bias=False)
        self.key = nn.Linear(self.hidden_size, self.all_head_size, bias=False)
        self.value = nn.Linear(self.hidden_size, self.all_head_size, bias=False)

        self.dropout = float(config.local_attention_probs_dropout_prob)

        # save mask value here
        self.mask_value_float16 = mindspore.tensor(-1e4)
        self.mask_value_float32 = mindspore.tensor(-1e9)

    def forward(
        self,
        hidden_states,
        attention_mask=None,
        head_mask=None,
        past_buckets_states=None,
        use_cache=False,
        output_attentions=False,
        **kwargs,
    ):

        """
        Constructs the local self-attention mechanism.

        Args:
            self (object): The instance of the LocalSelfAttention class.
            hidden_states (Tensor): The input hidden states with shape (batch_size, sequence_length, hidden_size).
            attention_mask (Tensor, optional): Mask to prevent attention to certain positions. Shape should be
                (batch_size, 1, sequence_length, sequence_length). Defaults to None.
            head_mask (Tensor, optional): Mask to prevent attention to certain heads.
                Shape should be (num_attention_heads,) or (num_layers, num_attention_heads). Defaults to None.
            past_buckets_states (Tuple[Tensor, Tensor], optional): Tuple containing the previous key and
                value hidden states. Defaults to None.
            use_cache (bool, optional): Flag indicating whether to use cached key and value states. Defaults to False.
            output_attentions (bool, optional): Flag indicating whether to output attention probabilities.
                Defaults to False.

        Returns:
            LocalSelfAttentionOutput: An object containing the output hidden states and attention probabilities.

        Raises:
            AssertionError: Raised if the conditions specified in the method are not met.
        """
        sequence_length = hidden_states.shape[1]
        batch_size = hidden_states.shape[0]

        # check if cache shall be used and that hidden states are already cached
        if use_cache and past_buckets_states[1] is not None:
            assert past_buckets_states[0] is None, (
                "LocalSelfAttention should not make use of `buckets`. There seems to be an error when caching"
                " hidden_states_and_buckets."
            )
            key_value_hidden_states = self._retrieve_relevant_hidden_states(
                past_buckets_states[1], self.chunk_length, self.num_chunks_before
            )
            key_value_hidden_states = ops.cat([key_value_hidden_states, hidden_states], axis=1)

            # only query vector for last token
            query_vectors = self.query(hidden_states)
            # compute key and value for relevant chunk
            key_vectors = self.key(key_value_hidden_states)
            value_vectors = self.value(key_value_hidden_states)

            # free memory
            del key_value_hidden_states
        else:
            # project hidden_states to query, key and value
            query_vectors = self.query(hidden_states)
            key_vectors = self.key(hidden_states)
            value_vectors = self.value(hidden_states)

        # split last dim into `config.num_attention_heads` and `config.attention_head_size`
        query_vectors = self._split_hidden_size_dim(query_vectors, self.num_attention_heads, self.attention_head_size)
        key_vectors = self._split_hidden_size_dim(key_vectors, self.num_attention_heads, self.attention_head_size)
        value_vectors = self._split_hidden_size_dim(value_vectors, self.num_attention_heads, self.attention_head_size)

        assert (
            query_vectors.shape[-1] == self.attention_head_size
        ), f"last dim of query_key_vectors is {query_vectors.shape[-1]} but should be {self.attention_head_size}."
        assert (
            key_vectors.shape[-1] == self.attention_head_size
        ), f"last dim of query_key_vectors is {key_vectors.shape[-1]} but should be {self.attention_head_size}."
        assert (
            value_vectors.shape[-1] == self.attention_head_size
        ), f"last dim of query_key_vectors is {value_vectors.shape[-1]} but should be {self.attention_head_size}."

        if self.chunk_length is None:
            assert self.num_chunks_before == 0 and self.num_chunks_after == 0, (
                "If `config.chunk_length` is `None`, make sure `config.num_chunks_after` and"
                " `config.num_chunks_before` are set to 0."
            )

        # normalize key vectors
        key_vectors = key_vectors / np.sqrt(self.attention_head_size)

        # get sequence length indices
        indices = ops.arange(sequence_length).repeat(
            batch_size, self.num_attention_heads, 1
        )

        # if one should do normal n^2 self-attention
        do_standard_self_attention = sequence_length <= self.chunk_length

        # if input should be chunked
        if not do_standard_self_attention:
            # chunk vectors
            # B x Num_Attn_Head x Seq_Len // chunk_len x chunk_len  x  attn_head_size
            query_vectors = self._split_seq_length_dim_to(
                query_vectors,
                -1,
                self.chunk_length,
                self.num_attention_heads,
                self.attention_head_size,
            )
            key_vectors = self._split_seq_length_dim_to(
                key_vectors,
                -1,
                self.chunk_length,
                self.num_attention_heads,
                self.attention_head_size,
            )
            value_vectors = self._split_seq_length_dim_to(
                value_vectors,
                -1,
                self.chunk_length,
                self.num_attention_heads,
                self.attention_head_size,
            )

            # chunk indices
            query_indices = self._split_seq_length_dim_to(indices, -1, self.chunk_length, self.num_attention_heads)
            key_indices = self._split_seq_length_dim_to(indices, -1, self.chunk_length, self.num_attention_heads)

            # append chunks before and after
            key_vectors = self._look_adjacent(key_vectors, self.num_chunks_before, self.num_chunks_after)
            value_vectors = self._look_adjacent(value_vectors, self.num_chunks_before, self.num_chunks_after)
            key_indices = self._look_adjacent(key_indices, self.num_chunks_before, self.num_chunks_after)
        else:
            query_indices = key_indices = indices

        # query-key matmul: QK^T
        query_key_dots = ops.matmul(query_vectors, key_vectors.swapaxes(-1, -2))

        # free memory
        del query_vectors, key_vectors

        mask = self._compute_attn_mask(
            query_indices, key_indices, attention_mask, query_key_dots.shape, do_standard_self_attention
        )

        if mask is not None:
            # get mask tensor depending on half precision or not
            if query_key_dots.dtype == mindspore.float16:
                mask_value = self.mask_value_float16.half()
            else:
                mask_value = self.mask_value_float32

            query_key_dots = ops.where(mask, query_key_dots, mask_value)

        # free memory
        del mask

        # softmax
        logits = ops.logsumexp(query_key_dots, axis=-1, keep_dims=True)
        attention_probs = ops.exp(query_key_dots - logits)

        # free memory
        del logits

        # dropout
        attention_probs = ops.dropout(attention_probs, p=self.dropout, training=self.training)

        # Mask heads if we want to
        if head_mask is not None:
            attention_probs = attention_probs * head_mask

        # attend values
        out_vectors = ops.matmul(attention_probs, value_vectors)

        # free memory
        del value_vectors

        # merge chunk length
        if not do_standard_self_attention:
            out_vectors = out_vectors.flatten(start_dim=2, end_dim=3)

        assert out_vectors.shape == (
            batch_size,
            self.num_attention_heads,
            sequence_length,
            self.attention_head_size,
        )

        out_vectors = self._merge_hidden_size_dims(out_vectors, self.num_attention_heads, self.attention_head_size)

        if output_attentions is False:
            attention_probs = ()

        return LocalSelfAttentionOutput(hidden_states=out_vectors, attention_probs=attention_probs)

    def _compute_attn_mask(
        self, query_indices, key_indices, attention_mask, query_key_dots_shape, do_standard_self_attention
    ):

        """
        Computes the attention mask for the LocalSelfAttention module.

        Args:
            self (LocalSelfAttention): The instance of the LocalSelfAttention class.
            query_indices (Tensor): The indices of the query sequence. Shape: (batch_size, query_seq_length).
            key_indices (Tensor): The indices of the key sequence. Shape: (batch_size, key_seq_length).
            attention_mask (Tensor): The attention mask tensor. Shape: (batch_size, query_seq_length, key_seq_length).
            query_key_dots_shape (Tensor): The shape of the query-key dot products.
                Shape: (batch_size, num_attention_heads, query_seq_length, key_seq_length).
            do_standard_self_attention (bool): A flag indicating whether to use standard self-attention.

        Returns:
            Tensor or None: The computed attention mask tensor.
                Shape: (batch_size, num_attention_heads, query_seq_length, key_seq_length).

        Raises:
            TypeError: If the input arguments are not of the expected types.
            ValueError: If the shapes of the input arguments are not valid.
        """
        # chunk attention mask and look before and after
        if attention_mask is not None:
            attention_mask = attention_mask.to(mindspore.bool_)[:, None, :]

            if not do_standard_self_attention:
                attention_mask = self._split_seq_length_dim_to(attention_mask, -1, self.chunk_length, 1)
                attention_mask = self._look_adjacent(attention_mask, self.num_chunks_before, self.num_chunks_after)
            # create attn_mask
            attention_mask = attention_mask.unsqueeze(-2).expand(query_key_dots_shape)

        # Causal mask
        if self.is_decoder is True:
            causal_mask = ops.ge(query_indices.unsqueeze(-1), key_indices.unsqueeze(-2))

            # add attention mask if not None
            if attention_mask is not None:
                attention_mask = causal_mask & attention_mask
            else:
                attention_mask = causal_mask

        return attention_mask

    @staticmethod
    def _retrieve_relevant_hidden_states(previous_hidden_states, chunk_length, num_chunks_before):

        """
        Retrieves relevant hidden states from previous hidden states based on the provided parameters.

        Args:
            previous_hidden_states (ndarray): An array of shape (batch_size, sequence_length, hidden_size)
                representing the previous hidden states.
            chunk_length (int): The length of each chunk of hidden states.
            num_chunks_before (int): The number of chunks before the relevant hidden states to retrieve.

        Returns:
            None

        Raises:
            None

        This static method retrieves relevant hidden states from previous hidden states based on the given parameters.
        It calculates the start position of the relevant hidden states based on the chunk length and the number of
        chunks before. The retrieved hidden states are then returned as an array with shape
        (batch_size, sequence_length, hidden_size), starting from the calculated start position.
        """
        start_position = ((previous_hidden_states.shape[1] // chunk_length) - num_chunks_before) * chunk_length
        return previous_hidden_states[:, start_position:]

mindnlp.transformers.models.reformer.modeling_reformer.LocalSelfAttention.__init__(config)

Initializes the LocalSelfAttention class.

PARAMETER	DESCRIPTION
self	The instance of the LocalSelfAttention class.
config	An object containing configuration parameters for the attention mechanism. This parameter is expected to have the following attributes: num_attention_heads (int): The number of attention heads. local_attn_chunk_length (int): The length of attention chunks. local_num_chunks_before (int): The number of chunks before the current position. local_num_chunks_after (int): The number of chunks after the current position. is_decoder (bool): Indicates if the attention mechanism is used in a decoder. pad_token_id (int): The token ID used for padding. attention_head_size (int): The size of each attention head. hidden_size (int): The size of the hidden layer. local_attention_probs_dropout_prob (float): The dropout probability for attention weights.

RETURNS	DESCRIPTION
	None.

RAISES	DESCRIPTION
ValueError	If the configuration parameters provided are invalid or missing.
TypeError	If any of the configuration attribute types are incorrect.

Source code in mindnlp/transformers/models/reformer/modeling_reformer.py

def __init__(self, config):

    """
    Initializes the LocalSelfAttention class.

    Args:
        self: The instance of the LocalSelfAttention class.
        config: An object containing configuration parameters for the attention mechanism.
            This parameter is expected to have the following attributes:

            - num_attention_heads (int): The number of attention heads.
            - local_attn_chunk_length (int): The length of attention chunks.
            - local_num_chunks_before (int): The number of chunks before the current position.
            - local_num_chunks_after (int): The number of chunks after the current position.
            - is_decoder (bool): Indicates if the attention mechanism is used in a decoder.
            - pad_token_id (int): The token ID used for padding.
            - attention_head_size (int): The size of each attention head.
            - hidden_size (int): The size of the hidden layer.
            - local_attention_probs_dropout_prob (float): The dropout probability for attention weights.

    Returns:
        None.

    Raises:
        ValueError: If the configuration parameters provided are invalid or missing.
        TypeError: If any of the configuration attribute types are incorrect.
    """
    super().__init__()

    self.num_attention_heads = config.num_attention_heads
    self.chunk_length = config.local_attn_chunk_length
    self.num_chunks_before = config.local_num_chunks_before
    self.num_chunks_after = config.local_num_chunks_after
    self.is_decoder = config.is_decoder
    self.pad_token_id = config.pad_token_id

    self.attention_head_size = config.attention_head_size
    self.all_head_size = self.num_attention_heads * self.attention_head_size
    self.hidden_size = config.hidden_size

    # projection matrices
    self.query = nn.Linear(self.hidden_size, self.all_head_size, bias=False)
    self.key = nn.Linear(self.hidden_size, self.all_head_size, bias=False)
    self.value = nn.Linear(self.hidden_size, self.all_head_size, bias=False)

    self.dropout = float(config.local_attention_probs_dropout_prob)

    # save mask value here
    self.mask_value_float16 = mindspore.tensor(-1e4)
    self.mask_value_float32 = mindspore.tensor(-1e9)

mindnlp.transformers.models.reformer.modeling_reformer.LocalSelfAttention.forward(hidden_states, attention_mask=None, head_mask=None, past_buckets_states=None, use_cache=False, output_attentions=False, **kwargs)

Constructs the local self-attention mechanism.

PARAMETER	DESCRIPTION
self	The instance of the LocalSelfAttention class. TYPE: object
hidden_states	The input hidden states with shape (batch_size, sequence_length, hidden_size). TYPE: Tensor
attention_mask	Mask to prevent attention to certain positions. Shape should be (batch_size, 1, sequence_length, sequence_length). Defaults to None. TYPE: Tensor DEFAULT: None
head_mask	Mask to prevent attention to certain heads. Shape should be (num_attention_heads,) or (num_layers, num_attention_heads). Defaults to None. TYPE: Tensor DEFAULT: None
past_buckets_states	Tuple containing the previous key and value hidden states. Defaults to None. TYPE: Tuple[Tensor, Tensor] DEFAULT: None
use_cache	Flag indicating whether to use cached key and value states. Defaults to False. TYPE: bool DEFAULT: False
output_attentions	Flag indicating whether to output attention probabilities. Defaults to False. TYPE: bool DEFAULT: False

RETURNS	DESCRIPTION
LocalSelfAttentionOutput	An object containing the output hidden states and attention probabilities.

RAISES	DESCRIPTION
AssertionError	Raised if the conditions specified in the method are not met.

Source code in mindnlp/transformers/models/reformer/modeling_reformer.py

def forward(
    self,
    hidden_states,
    attention_mask=None,
    head_mask=None,
    past_buckets_states=None,
    use_cache=False,
    output_attentions=False,
    **kwargs,
):

    """
    Constructs the local self-attention mechanism.

    Args:
        self (object): The instance of the LocalSelfAttention class.
        hidden_states (Tensor): The input hidden states with shape (batch_size, sequence_length, hidden_size).
        attention_mask (Tensor, optional): Mask to prevent attention to certain positions. Shape should be
            (batch_size, 1, sequence_length, sequence_length). Defaults to None.
        head_mask (Tensor, optional): Mask to prevent attention to certain heads.
            Shape should be (num_attention_heads,) or (num_layers, num_attention_heads). Defaults to None.
        past_buckets_states (Tuple[Tensor, Tensor], optional): Tuple containing the previous key and
            value hidden states. Defaults to None.
        use_cache (bool, optional): Flag indicating whether to use cached key and value states. Defaults to False.
        output_attentions (bool, optional): Flag indicating whether to output attention probabilities.
            Defaults to False.

    Returns:
        LocalSelfAttentionOutput: An object containing the output hidden states and attention probabilities.

    Raises:
        AssertionError: Raised if the conditions specified in the method are not met.
    """
    sequence_length = hidden_states.shape[1]
    batch_size = hidden_states.shape[0]

    # check if cache shall be used and that hidden states are already cached
    if use_cache and past_buckets_states[1] is not None:
        assert past_buckets_states[0] is None, (
            "LocalSelfAttention should not make use of `buckets`. There seems to be an error when caching"
            " hidden_states_and_buckets."
        )
        key_value_hidden_states = self._retrieve_relevant_hidden_states(
            past_buckets_states[1], self.chunk_length, self.num_chunks_before
        )
        key_value_hidden_states = ops.cat([key_value_hidden_states, hidden_states], axis=1)

        # only query vector for last token
        query_vectors = self.query(hidden_states)
        # compute key and value for relevant chunk
        key_vectors = self.key(key_value_hidden_states)
        value_vectors = self.value(key_value_hidden_states)

        # free memory
        del key_value_hidden_states
    else:
        # project hidden_states to query, key and value
        query_vectors = self.query(hidden_states)
        key_vectors = self.key(hidden_states)
        value_vectors = self.value(hidden_states)

    # split last dim into `config.num_attention_heads` and `config.attention_head_size`
    query_vectors = self._split_hidden_size_dim(query_vectors, self.num_attention_heads, self.attention_head_size)
    key_vectors = self._split_hidden_size_dim(key_vectors, self.num_attention_heads, self.attention_head_size)
    value_vectors = self._split_hidden_size_dim(value_vectors, self.num_attention_heads, self.attention_head_size)

    assert (
        query_vectors.shape[-1] == self.attention_head_size
    ), f"last dim of query_key_vectors is {query_vectors.shape[-1]} but should be {self.attention_head_size}."
    assert (
        key_vectors.shape[-1] == self.attention_head_size
    ), f"last dim of query_key_vectors is {key_vectors.shape[-1]} but should be {self.attention_head_size}."
    assert (
        value_vectors.shape[-1] == self.attention_head_size
    ), f"last dim of query_key_vectors is {value_vectors.shape[-1]} but should be {self.attention_head_size}."

    if self.chunk_length is None:
        assert self.num_chunks_before == 0 and self.num_chunks_after == 0, (
            "If `config.chunk_length` is `None`, make sure `config.num_chunks_after` and"
            " `config.num_chunks_before` are set to 0."
        )

    # normalize key vectors
    key_vectors = key_vectors / np.sqrt(self.attention_head_size)

    # get sequence length indices
    indices = ops.arange(sequence_length).repeat(
        batch_size, self.num_attention_heads, 1
    )

    # if one should do normal n^2 self-attention
    do_standard_self_attention = sequence_length <= self.chunk_length

    # if input should be chunked
    if not do_standard_self_attention:
        # chunk vectors
        # B x Num_Attn_Head x Seq_Len // chunk_len x chunk_len  x  attn_head_size
        query_vectors = self._split_seq_length_dim_to(
            query_vectors,
            -1,
            self.chunk_length,
            self.num_attention_heads,
            self.attention_head_size,
        )
        key_vectors = self._split_seq_length_dim_to(
            key_vectors,
            -1,
            self.chunk_length,
            self.num_attention_heads,
            self.attention_head_size,
        )
        value_vectors = self._split_seq_length_dim_to(
            value_vectors,
            -1,
            self.chunk_length,
            self.num_attention_heads,
            self.attention_head_size,
        )

        # chunk indices
        query_indices = self._split_seq_length_dim_to(indices, -1, self.chunk_length, self.num_attention_heads)
        key_indices = self._split_seq_length_dim_to(indices, -1, self.chunk_length, self.num_attention_heads)

        # append chunks before and after
        key_vectors = self._look_adjacent(key_vectors, self.num_chunks_before, self.num_chunks_after)
        value_vectors = self._look_adjacent(value_vectors, self.num_chunks_before, self.num_chunks_after)
        key_indices = self._look_adjacent(key_indices, self.num_chunks_before, self.num_chunks_after)
    else:
        query_indices = key_indices = indices

    # query-key matmul: QK^T
    query_key_dots = ops.matmul(query_vectors, key_vectors.swapaxes(-1, -2))

    # free memory
    del query_vectors, key_vectors

    mask = self._compute_attn_mask(
        query_indices, key_indices, attention_mask, query_key_dots.shape, do_standard_self_attention
    )

    if mask is not None:
        # get mask tensor depending on half precision or not
        if query_key_dots.dtype == mindspore.float16:
            mask_value = self.mask_value_float16.half()
        else:
            mask_value = self.mask_value_float32

        query_key_dots = ops.where(mask, query_key_dots, mask_value)

    # free memory
    del mask

    # softmax
    logits = ops.logsumexp(query_key_dots, axis=-1, keep_dims=True)
    attention_probs = ops.exp(query_key_dots - logits)

    # free memory
    del logits

    # dropout
    attention_probs = ops.dropout(attention_probs, p=self.dropout, training=self.training)

    # Mask heads if we want to
    if head_mask is not None:
        attention_probs = attention_probs * head_mask

    # attend values
    out_vectors = ops.matmul(attention_probs, value_vectors)

    # free memory
    del value_vectors

    # merge chunk length
    if not do_standard_self_attention:
        out_vectors = out_vectors.flatten(start_dim=2, end_dim=3)

    assert out_vectors.shape == (
        batch_size,
        self.num_attention_heads,
        sequence_length,
        self.attention_head_size,
    )

    out_vectors = self._merge_hidden_size_dims(out_vectors, self.num_attention_heads, self.attention_head_size)

    if output_attentions is False:
        attention_probs = ()

    return LocalSelfAttentionOutput(hidden_states=out_vectors, attention_probs=attention_probs)

mindnlp.transformers.models.reformer.modeling_reformer.PositionEmbeddings

Bases: Module

Constructs conventional position embeddings of shape [max_pos_embeddings, hidden_size].

Source code in mindnlp/transformers/models/reformer/modeling_reformer.py

class PositionEmbeddings(nn.Module):
    """Constructs conventional position embeddings of shape `[max_pos_embeddings, hidden_size]`."""
    def __init__(self, config):
        """
        Initializes an instance of the PositionEmbeddings class.

        Args:
            self: The instance of the PositionEmbeddings class.
            config:
                An instance of the configuration class containing the following attributes:

                - hidden_dropout_prob (float): The dropout probability used for the hidden layers.
                - max_position_embeddings (int): The maximum number of position embeddings.
                - hidden_size (int): The size of the hidden layers.

        Returns:
            None

        Raises:
            None
        """
        super().__init__()
        self.dropout = config.hidden_dropout_prob
        self.embedding = nn.Embedding(config.max_position_embeddings, config.hidden_size)

    def forward(self, position_ids):
        """
        Constructs position embeddings based on given position IDs.

        Args:
            self (PositionEmbeddings): The instance of the PositionEmbeddings class.
            position_ids (Tensor): A tensor containing the position IDs for which embeddings need to be forwarded.
                It should be a 1D tensor of integers representing the positions.

        Returns:
            position_embeddings: The method returns the forwarded position embeddings.

        Raises:
            ValueError: If the position_ids tensor is not provided or is not a valid 1D tensor.
            TypeError: If the dropout rate is not a float or the training flag is not a boolean.
        """
        position_embeddings = self.embedding(position_ids)
        position_embeddings = ops.dropout(position_embeddings, p=self.dropout, training=self.training)
        return position_embeddings

mindnlp.transformers.models.reformer.modeling_reformer.PositionEmbeddings.__init__(config)

Initializes an instance of the PositionEmbeddings class.

PARAMETER	DESCRIPTION
self	The instance of the PositionEmbeddings class.
config	An instance of the configuration class containing the following attributes: hidden_dropout_prob (float): The dropout probability used for the hidden layers. max_position_embeddings (int): The maximum number of position embeddings. hidden_size (int): The size of the hidden layers.

RETURNS	DESCRIPTION
	None

Source code in mindnlp/transformers/models/reformer/modeling_reformer.py

def __init__(self, config):
    """
    Initializes an instance of the PositionEmbeddings class.

    Args:
        self: The instance of the PositionEmbeddings class.
        config:
            An instance of the configuration class containing the following attributes:

            - hidden_dropout_prob (float): The dropout probability used for the hidden layers.
            - max_position_embeddings (int): The maximum number of position embeddings.
            - hidden_size (int): The size of the hidden layers.

    Returns:
        None

    Raises:
        None
    """
    super().__init__()
    self.dropout = config.hidden_dropout_prob
    self.embedding = nn.Embedding(config.max_position_embeddings, config.hidden_size)

mindnlp.transformers.models.reformer.modeling_reformer.PositionEmbeddings.forward(position_ids)

Constructs position embeddings based on given position IDs.

PARAMETER	DESCRIPTION
self	The instance of the PositionEmbeddings class. TYPE: PositionEmbeddings
position_ids	A tensor containing the position IDs for which embeddings need to be forwarded. It should be a 1D tensor of integers representing the positions. TYPE: Tensor

RETURNS	DESCRIPTION
position_embeddings	The method returns the forwarded position embeddings.

RAISES	DESCRIPTION
ValueError	If the position_ids tensor is not provided or is not a valid 1D tensor.
TypeError	If the dropout rate is not a float or the training flag is not a boolean.

Source code in mindnlp/transformers/models/reformer/modeling_reformer.py

def forward(self, position_ids):
    """
    Constructs position embeddings based on given position IDs.

    Args:
        self (PositionEmbeddings): The instance of the PositionEmbeddings class.
        position_ids (Tensor): A tensor containing the position IDs for which embeddings need to be forwarded.
            It should be a 1D tensor of integers representing the positions.

    Returns:
        position_embeddings: The method returns the forwarded position embeddings.

    Raises:
        ValueError: If the position_ids tensor is not provided or is not a valid 1D tensor.
        TypeError: If the dropout rate is not a float or the training flag is not a boolean.
    """
    position_embeddings = self.embedding(position_ids)
    position_embeddings = ops.dropout(position_embeddings, p=self.dropout, training=self.training)
    return position_embeddings

mindnlp.transformers.models.reformer.modeling_reformer.ReformerAttention

Bases: Module

This class represents a ReformerAttention module, which is used in the Reformer model for attention mechanisms. It inherits from the nn.Module class.

ATTRIBUTE	DESCRIPTION
layer_id	The ID of the attention layer. TYPE: int
attn_layers	The types of attention layers used in the Reformer model. TYPE: list
layer_norm	Layer normalization module applied to the input hidden states. TYPE: LayerNorm
self_attention	The self-attention module used based on the attention layer type. TYPE: LSHSelfAttention or LocalSelfAttention
output	The module responsible for the final output of the self-attention mechanism. TYPE: ReformerSelfOutput

METHOD	DESCRIPTION
forward	Applies the ReformerAttention module to the input hidden_states and returns the attention output.

Note

• Only 'lsh' and 'local' attention layer types are supported. The attention layer types can be selected from ['lsh', 'local'] only.

RAISES	DESCRIPTION
NotImplementedError	If the input attention layer types are not 'lsh' or 'local'.

Source code in mindnlp/transformers/models/reformer/modeling_reformer.py

class ReformerAttention(nn.Module):

    """
    This class represents a ReformerAttention module, which is used in the Reformer model for attention mechanisms.
    It inherits from the nn.Module class.

    Attributes:
        layer_id (int): The ID of the attention layer.
        attn_layers (list): The types of attention layers used in the Reformer model.
        layer_norm (nn.LayerNorm): Layer normalization module applied to the input hidden states.
        self_attention (LSHSelfAttention or LocalSelfAttention): The self-attention module used based on
            the attention layer type.
        output (ReformerSelfOutput): The module responsible for the final output of the self-attention mechanism.

    Methods:
        forward: Applies the ReformerAttention module to the input hidden_states and returns the attention output.

    Note:
        - Only 'lsh' and 'local' attention layer types are supported. The attention layer types can be selected from
        ['lsh', 'local'] only.

    Raises:
        NotImplementedError: If the input attention layer types are not 'lsh' or 'local'.

    """
    def __init__(self, config, layer_id=0):

        """
        Initialize the ReformerAttention class.

        Args:
            self (object): The instance of the ReformerAttention class.
            config (object): An object containing configuration settings for the attention layer.
            layer_id (int, optional): The ID of the layer within the attention module. Defaults to 0.

        Returns:
            None.

        Raises:
            NotImplementedError: If the specified attention layer types are not 'lsh' or 'local'.
        """
        super().__init__()
        self.layer_id = layer_id
        self.attn_layers = config.attn_layers

        self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)

        if len(set(self.attn_layers)) == 1 and self.attn_layers[0] == "lsh":
            self.self_attention = LSHSelfAttention(config)
        elif len(set(self.attn_layers)) == 1 and self.attn_layers[0] == "local":
            self.self_attention = LocalSelfAttention(config)
        elif len(set(self.attn_layers)) == 2 and set(self.attn_layers) == {"lsh", "local"}:
            # get correct attn layers
            if self.attn_layers[self.layer_id] == "lsh":
                self.self_attention = LSHSelfAttention(config)
            else:
                self.self_attention = LocalSelfAttention(config)
        else:
            raise NotImplementedError(
                f"Only attn layer types 'lsh' and 'local' exist, but got `config.attn_layers`: {self.attn_layers}. "
                "Select attn layer types from ['lsh', 'local'] only."
            )
        self.output = ReformerSelfOutput(config)

    def forward(
        self,
        hidden_states,
        attention_mask=None,
        head_mask=None,
        num_hashes=None,
        past_buckets_states=None,
        use_cache=False,
        orig_sequence_length=None,
        output_attentions=False,
        buckets=None,
    ):

        """
        This method forwards the attention output in the ReformerAttention class.

        Args:
            self: Reference to the class instance.
            hidden_states (torch.Tensor): Input tensor containing the hidden states.
            attention_mask (torch.Tensor, optional): Mask tensor for controlling attention computation.
            head_mask (torch.Tensor, optional): Mask tensor for controlling attention within each head.
            num_hashes (int, optional): Number of hashes to use in hashing attention.
            past_buckets_states (torch.Tensor, optional): Tensor containing past bucket states.
            use_cache (bool, optional): Flag indicating whether to use cache for attention computation.
            orig_sequence_length (int, optional): Original length of the sequence.
            output_attentions (bool, optional): Flag indicating whether to output attention weights.
            buckets (torch.Tensor, optional): Tensor containing bucket information.

        Returns:
            None.

        Raises:
            ValueError: If the orig_sequence_length is less than 1.
            AttributeError: If the self_attention_outputs object does not have the 'buckets' attribute.
            IndexError: If the layer_id in past_buckets_states is out of bounds.
            RuntimeError: If there is an issue with concatenating tensors using ops.cat.
            TypeError: If the input types are not as expected.
            Exception: For any other unforeseen errors during the method execution.
        """
        hidden_states = self.layer_norm(hidden_states)

        # make sure cached hidden states is set to None for backward pass
        if past_buckets_states is not None:
            past_buckets_states_layer = past_buckets_states[self.layer_id]
        else:
            past_buckets_states_layer = None

        # use cached buckets for backprob if buckets not None for LSHSelfAttention
        self_attention_outputs = self.self_attention(
            hidden_states=hidden_states,
            head_mask=head_mask,
            attention_mask=attention_mask,
            num_hashes=num_hashes,
            past_buckets_states=past_buckets_states_layer,
            use_cache=use_cache,
            output_attentions=output_attentions,
            buckets=buckets,
        )

        # add buckets if necessary
        if hasattr(self_attention_outputs, "buckets"):
            buckets = self_attention_outputs.buckets
        else:
            buckets = None

        # cache hidden states for future use
        if use_cache:
            if past_buckets_states[self.layer_id][0] is None:
                # padded input should not be cached
                past_buckets = (
                    buckets[:, :, :, :orig_sequence_length]
                    if (buckets is not None and orig_sequence_length > 1)
                    else buckets
                )
            else:
                past_buckets = ops.cat([past_buckets_states[self.layer_id][0], buckets], axis=-1)

            if past_buckets_states[self.layer_id][1] is None:
                # padded input should not be cached
                past_states = hidden_states[:, :orig_sequence_length]
            else:
                past_states = ops.cat([past_buckets_states[self.layer_id][1], hidden_states], axis=1)

            past_buckets_states[self.layer_id] = (past_buckets, past_states)
        # compute attention feed forward output
        attention_output = self.output(self_attention_outputs.hidden_states)

        return AttentionOutput(
            hidden_states=attention_output,
            attention_probs=self_attention_outputs.attention_probs,
            buckets=buckets,
        )

mindnlp.transformers.models.reformer.modeling_reformer.ReformerAttention.__init__(config, layer_id=0)

Initialize the ReformerAttention class.

PARAMETER	DESCRIPTION
self	The instance of the ReformerAttention class. TYPE: object
config	An object containing configuration settings for the attention layer. TYPE: object
layer_id	The ID of the layer within the attention module. Defaults to 0. TYPE: int DEFAULT: 0

RETURNS	DESCRIPTION
	None.

RAISES	DESCRIPTION
NotImplementedError	If the specified attention layer types are not 'lsh' or 'local'.

Source code in mindnlp/transformers/models/reformer/modeling_reformer.py

def __init__(self, config, layer_id=0):

    """
    Initialize the ReformerAttention class.

    Args:
        self (object): The instance of the ReformerAttention class.
        config (object): An object containing configuration settings for the attention layer.
        layer_id (int, optional): The ID of the layer within the attention module. Defaults to 0.

    Returns:
        None.

    Raises:
        NotImplementedError: If the specified attention layer types are not 'lsh' or 'local'.
    """
    super().__init__()
    self.layer_id = layer_id
    self.attn_layers = config.attn_layers

    self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)

    if len(set(self.attn_layers)) == 1 and self.attn_layers[0] == "lsh":
        self.self_attention = LSHSelfAttention(config)
    elif len(set(self.attn_layers)) == 1 and self.attn_layers[0] == "local":
        self.self_attention = LocalSelfAttention(config)
    elif len(set(self.attn_layers)) == 2 and set(self.attn_layers) == {"lsh", "local"}:
        # get correct attn layers
        if self.attn_layers[self.layer_id] == "lsh":
            self.self_attention = LSHSelfAttention(config)
        else:
            self.self_attention = LocalSelfAttention(config)
    else:
        raise NotImplementedError(
            f"Only attn layer types 'lsh' and 'local' exist, but got `config.attn_layers`: {self.attn_layers}. "
            "Select attn layer types from ['lsh', 'local'] only."
        )
    self.output = ReformerSelfOutput(config)

mindnlp.transformers.models.reformer.modeling_reformer.ReformerAttention.forward(hidden_states, attention_mask=None, head_mask=None, num_hashes=None, past_buckets_states=None, use_cache=False, orig_sequence_length=None, output_attentions=False, buckets=None)

This method forwards the attention output in the ReformerAttention class.

PARAMETER	DESCRIPTION
self	Reference to the class instance.
hidden_states	Input tensor containing the hidden states. TYPE: Tensor
attention_mask	Mask tensor for controlling attention computation. TYPE: Tensor DEFAULT: None
head_mask	Mask tensor for controlling attention within each head. TYPE: Tensor DEFAULT: None
num_hashes	Number of hashes to use in hashing attention. TYPE: int DEFAULT: None
past_buckets_states	Tensor containing past bucket states. TYPE: Tensor DEFAULT: None
use_cache	Flag indicating whether to use cache for attention computation. TYPE: bool DEFAULT: False
orig_sequence_length	Original length of the sequence. TYPE: int DEFAULT: None
output_attentions	Flag indicating whether to output attention weights. TYPE: bool DEFAULT: False
buckets	Tensor containing bucket information. TYPE: Tensor DEFAULT: None

RETURNS	DESCRIPTION
	None.

RAISES	DESCRIPTION
ValueError	If the orig_sequence_length is less than 1.
AttributeError	If the self_attention_outputs object does not have the 'buckets' attribute.
IndexError	If the layer_id in past_buckets_states is out of bounds.
RuntimeError	If there is an issue with concatenating tensors using ops.cat.
TypeError	If the input types are not as expected.
Exception	For any other unforeseen errors during the method execution.

Source code in mindnlp/transformers/models/reformer/modeling_reformer.py

def forward(
    self,
    hidden_states,
    attention_mask=None,
    head_mask=None,
    num_hashes=None,
    past_buckets_states=None,
    use_cache=False,
    orig_sequence_length=None,
    output_attentions=False,
    buckets=None,
):

    """
    This method forwards the attention output in the ReformerAttention class.

    Args:
        self: Reference to the class instance.
        hidden_states (torch.Tensor): Input tensor containing the hidden states.
        attention_mask (torch.Tensor, optional): Mask tensor for controlling attention computation.
        head_mask (torch.Tensor, optional): Mask tensor for controlling attention within each head.
        num_hashes (int, optional): Number of hashes to use in hashing attention.
        past_buckets_states (torch.Tensor, optional): Tensor containing past bucket states.
        use_cache (bool, optional): Flag indicating whether to use cache for attention computation.
        orig_sequence_length (int, optional): Original length of the sequence.
        output_attentions (bool, optional): Flag indicating whether to output attention weights.
        buckets (torch.Tensor, optional): Tensor containing bucket information.

    Returns:
        None.

    Raises:
        ValueError: If the orig_sequence_length is less than 1.
        AttributeError: If the self_attention_outputs object does not have the 'buckets' attribute.
        IndexError: If the layer_id in past_buckets_states is out of bounds.
        RuntimeError: If there is an issue with concatenating tensors using ops.cat.
        TypeError: If the input types are not as expected.
        Exception: For any other unforeseen errors during the method execution.
    """
    hidden_states = self.layer_norm(hidden_states)

    # make sure cached hidden states is set to None for backward pass
    if past_buckets_states is not None:
        past_buckets_states_layer = past_buckets_states[self.layer_id]
    else:
        past_buckets_states_layer = None

    # use cached buckets for backprob if buckets not None for LSHSelfAttention
    self_attention_outputs = self.self_attention(
        hidden_states=hidden_states,
        head_mask=head_mask,
        attention_mask=attention_mask,
        num_hashes=num_hashes,
        past_buckets_states=past_buckets_states_layer,
        use_cache=use_cache,
        output_attentions=output_attentions,
        buckets=buckets,
    )

    # add buckets if necessary
    if hasattr(self_attention_outputs, "buckets"):
        buckets = self_attention_outputs.buckets
    else:
        buckets = None

    # cache hidden states for future use
    if use_cache:
        if past_buckets_states[self.layer_id][0] is None:
            # padded input should not be cached
            past_buckets = (
                buckets[:, :, :, :orig_sequence_length]
                if (buckets is not None and orig_sequence_length > 1)
                else buckets
            )
        else:
            past_buckets = ops.cat([past_buckets_states[self.layer_id][0], buckets], axis=-1)

        if past_buckets_states[self.layer_id][1] is None:
            # padded input should not be cached
            past_states = hidden_states[:, :orig_sequence_length]
        else:
            past_states = ops.cat([past_buckets_states[self.layer_id][1], hidden_states], axis=1)

        past_buckets_states[self.layer_id] = (past_buckets, past_states)
    # compute attention feed forward output
    attention_output = self.output(self_attention_outputs.hidden_states)

    return AttentionOutput(
        hidden_states=attention_output,
        attention_probs=self_attention_outputs.attention_probs,
        buckets=buckets,
    )

mindnlp.transformers.models.reformer.modeling_reformer.ReformerClassificationHead

Bases: Module

Head for sentence-level classification tasks.

Source code in mindnlp/transformers/models/reformer/modeling_reformer.py

class ReformerClassificationHead(nn.Module):
    """Head for sentence-level classification tasks."""
    def __init__(self, config):

        """
        Initializes the ReformerClassificationHead class.

        Args:
            self (ReformerClassificationHead): The instance of the ReformerClassificationHead class.
            config (object): The configuration object containing the settings for the Reformer model.
                It should have attributes such as hidden_size (int), classifier_dropout (float, optional),
                hidden_dropout_prob (float), and num_labels (int). The config object is required and should not be None.

        Returns:
            None.

        Raises:
            ValueError: If the config parameter is None or if any of the required attributes in the
                config object are missing.
            TypeError: If the config parameter is not of the expected type or if any attribute in the
                config object has an unexpected type.
            RuntimeError: If there is an issue with the initialization of the Dense and Dropout layers.
        """
        super().__init__()
        self.dense = nn.Linear(2 * config.hidden_size, config.hidden_size)
        classifier_dropout = (
            config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
        )
        self.dropout = nn.Dropout(p=classifier_dropout)
        self.out_proj = nn.Linear(config.hidden_size, config.num_labels)

    def forward(self, hidden_states, **kwargs):

        """
        Constructs the classification head for the Reformer model.

        Args:
            self (ReformerClassificationHead): An instance of the ReformerClassificationHead class.
            hidden_states (torch.Tensor): The hidden states of the input sequence.
                It should have shape (batch_size, sequence_length, hidden_size).

        Returns:
            None: This method does not return any value.

        Raises:
            None: No exceptions are raised by this method.
        """
        hidden_states = hidden_states[:, 0, :]  # take <s> token (equiv. to [CLS])
        hidden_states = self.dropout(hidden_states)
        hidden_states = self.dense(hidden_states)
        hidden_states = ops.tanh(hidden_states)
        hidden_states = self.dropout(hidden_states)
        hidden_states = self.out_proj(hidden_states)
        return hidden_states

mindnlp.transformers.models.reformer.modeling_reformer.ReformerClassificationHead.__init__(config)

Initializes the ReformerClassificationHead class.

PARAMETER	DESCRIPTION
self	The instance of the ReformerClassificationHead class. TYPE: ReformerClassificationHead
config	The configuration object containing the settings for the Reformer model. It should have attributes such as hidden_size (int), classifier_dropout (float, optional), hidden_dropout_prob (float), and num_labels (int). The config object is required and should not be None. TYPE: object

RETURNS	DESCRIPTION
	None.

RAISES	DESCRIPTION
ValueError	If the config parameter is None or if any of the required attributes in the config object are missing.
TypeError	If the config parameter is not of the expected type or if any attribute in the config object has an unexpected type.
RuntimeError	If there is an issue with the initialization of the Dense and Dropout layers.

Source code in mindnlp/transformers/models/reformer/modeling_reformer.py

def __init__(self, config):

    """
    Initializes the ReformerClassificationHead class.

    Args:
        self (ReformerClassificationHead): The instance of the ReformerClassificationHead class.
        config (object): The configuration object containing the settings for the Reformer model.
            It should have attributes such as hidden_size (int), classifier_dropout (float, optional),
            hidden_dropout_prob (float), and num_labels (int). The config object is required and should not be None.

    Returns:
        None.

    Raises:
        ValueError: If the config parameter is None or if any of the required attributes in the
            config object are missing.
        TypeError: If the config parameter is not of the expected type or if any attribute in the
            config object has an unexpected type.
        RuntimeError: If there is an issue with the initialization of the Dense and Dropout layers.
    """
    super().__init__()
    self.dense = nn.Linear(2 * config.hidden_size, config.hidden_size)
    classifier_dropout = (
        config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
    )
    self.dropout = nn.Dropout(p=classifier_dropout)
    self.out_proj = nn.Linear(config.hidden_size, config.num_labels)

mindnlp.transformers.models.reformer.modeling_reformer.ReformerClassificationHead.forward(hidden_states, **kwargs)

Constructs the classification head for the Reformer model.

PARAMETER	DESCRIPTION
self	An instance of the ReformerClassificationHead class. TYPE: ReformerClassificationHead
hidden_states	The hidden states of the input sequence. It should have shape (batch_size, sequence_length, hidden_size). TYPE: Tensor

RETURNS	DESCRIPTION
None	This method does not return any value.

RAISES	DESCRIPTION
None	No exceptions are raised by this method.

Source code in mindnlp/transformers/models/reformer/modeling_reformer.py

def forward(self, hidden_states, **kwargs):

    """
    Constructs the classification head for the Reformer model.

    Args:
        self (ReformerClassificationHead): An instance of the ReformerClassificationHead class.
        hidden_states (torch.Tensor): The hidden states of the input sequence.
            It should have shape (batch_size, sequence_length, hidden_size).

    Returns:
        None: This method does not return any value.

    Raises:
        None: No exceptions are raised by this method.
    """
    hidden_states = hidden_states[:, 0, :]  # take <s> token (equiv. to [CLS])
    hidden_states = self.dropout(hidden_states)
    hidden_states = self.dense(hidden_states)
    hidden_states = ops.tanh(hidden_states)
    hidden_states = self.dropout(hidden_states)
    hidden_states = self.out_proj(hidden_states)
    return hidden_states

mindnlp.transformers.models.reformer.modeling_reformer.ReformerEmbeddings

Bases: Module

Construct the embeddings from word, position and token_type embeddings.

Source code in mindnlp/transformers/models/reformer/modeling_reformer.py

class ReformerEmbeddings(nn.Module):
    """Construct the embeddings from word, position and token_type embeddings."""
    def __init__(self, config):
        """
        Initializes the ReformerEmbeddings class.

        Args:
            self (object): The instance of the ReformerEmbeddings class.
            config (object):
                An object containing configuration parameters for the embeddings.

                - max_position_embeddings (int): The maximum number of positions for position embeddings.
                - hidden_dropout_prob (float): The dropout probability for hidden layers.
                - vocab_size (int): The size of the vocabulary.
                - hidden_size (int): The size of the hidden layers.
                - axial_pos_embds (bool): A flag indicating whether to use axial position embeddings.

        Returns:
            None.

        Raises:
            None.
        """
        super().__init__()
        self.max_position_embeddings = config.max_position_embeddings
        self.dropout = float(config.hidden_dropout_prob)

        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size)
        self.position_embeddings = (
            AxialPositionEmbeddings(config) if config.axial_pos_embds else PositionEmbeddings(config)
        )

    def forward(self, input_ids=None, position_ids=None, inputs_embeds=None, start_idx_pos_encodings=0):
        """
        Constructs the embeddings for the Reformer model.

        Args:
            self (ReformerEmbeddings): An instance of the ReformerEmbeddings class.
            input_ids (Optional[Tensor]): The input tensor representing the tokenized input sequence.
                It has shape [batch_size, sequence_length] and each element corresponds to a token ID.
            position_ids (Optional[Tensor]): The tensor representing the position encodings for each token
                in the input sequence. It has shape [batch_size, sequence_length] and each element corresponds
                to a position ID.
            inputs_embeds (Optional[Tensor]): The tensor representing the embedded input sequence.
                It has shape [batch_size, sequence_length, embedding_size] and each element corresponds
                to an embedding vector.
            start_idx_pos_encodings (int): The starting index for the positional encodings.

        Returns:
            Tensor: The embeddings tensor representing the input sequence with positional information.
                It has shape [batch_size, sequence_length, embedding_size].

        Raises:
            ValueError: If the sequence length of the position_ids tensor is greater than the maximum allowed
                sequence length.
        """
        if input_ids is not None:
            input_shape = input_ids.shape
        else:
            input_shape = inputs_embeds.shape[:-1]

        seq_length = input_shape[1]
        if position_ids is None:
            position_ids = ops.arange(
                start_idx_pos_encodings, start_idx_pos_encodings + seq_length, dtype=mindspore.int64
            )
            position_ids = position_ids.unsqueeze(0).expand(input_shape)

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

        if position_ids.shape[-1] > self.max_position_embeddings:
            raise ValueError(
                f"Sequence Length: {position_ids.shape[-1]} has to be less or equal than "
                f"config.max_position_embeddings {self.max_position_embeddings}."
            )

        # dropout
        embeddings = ops.dropout(inputs_embeds, p=self.dropout, training=self.training)

        # add positional embeddings
        position_embeddings = self.position_embeddings(position_ids)
        embeddings = embeddings + position_embeddings
        return embeddings

mindnlp.transformers.models.reformer.modeling_reformer.ReformerEmbeddings.__init__(config)

Initializes the ReformerEmbeddings class.

PARAMETER	DESCRIPTION
self	The instance of the ReformerEmbeddings class. TYPE: object
config	An object containing configuration parameters for the embeddings. max_position_embeddings (int): The maximum number of positions for position embeddings. hidden_dropout_prob (float): The dropout probability for hidden layers. vocab_size (int): The size of the vocabulary. hidden_size (int): The size of the hidden layers. axial_pos_embds (bool): A flag indicating whether to use axial position embeddings. TYPE: object

RETURNS	DESCRIPTION
	None.

Source code in mindnlp/transformers/models/reformer/modeling_reformer.py

def __init__(self, config):
    """
    Initializes the ReformerEmbeddings class.

    Args:
        self (object): The instance of the ReformerEmbeddings class.
        config (object):
            An object containing configuration parameters for the embeddings.

            - max_position_embeddings (int): The maximum number of positions for position embeddings.
            - hidden_dropout_prob (float): The dropout probability for hidden layers.
            - vocab_size (int): The size of the vocabulary.
            - hidden_size (int): The size of the hidden layers.
            - axial_pos_embds (bool): A flag indicating whether to use axial position embeddings.

    Returns:
        None.

    Raises:
        None.
    """
    super().__init__()
    self.max_position_embeddings = config.max_position_embeddings
    self.dropout = float(config.hidden_dropout_prob)

    self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size)
    self.position_embeddings = (
        AxialPositionEmbeddings(config) if config.axial_pos_embds else PositionEmbeddings(config)
    )

mindnlp.transformers.models.reformer.modeling_reformer.ReformerEmbeddings.forward(input_ids=None, position_ids=None, inputs_embeds=None, start_idx_pos_encodings=0)

Constructs the embeddings for the Reformer model.

PARAMETER	DESCRIPTION
self	An instance of the ReformerEmbeddings class. TYPE: ReformerEmbeddings
input_ids	The input tensor representing the tokenized input sequence. It has shape [batch_size, sequence_length] and each element corresponds to a token ID. TYPE: Optional[Tensor] DEFAULT: None
position_ids	The tensor representing the position encodings for each token in the input sequence. It has shape [batch_size, sequence_length] and each element corresponds to a position ID. TYPE: Optional[Tensor] DEFAULT: None
inputs_embeds	The tensor representing the embedded input sequence. It has shape [batch_size, sequence_length, embedding_size] and each element corresponds to an embedding vector. TYPE: Optional[Tensor] DEFAULT: None
start_idx_pos_encodings	The starting index for the positional encodings. TYPE: int DEFAULT: 0

RETURNS	DESCRIPTION
Tensor	The embeddings tensor representing the input sequence with positional information. It has shape [batch_size, sequence_length, embedding_size].

RAISES	DESCRIPTION
ValueError	If the sequence length of the position_ids tensor is greater than the maximum allowed sequence length.

Source code in mindnlp/transformers/models/reformer/modeling_reformer.py

def forward(self, input_ids=None, position_ids=None, inputs_embeds=None, start_idx_pos_encodings=0):
    """
    Constructs the embeddings for the Reformer model.

    Args:
        self (ReformerEmbeddings): An instance of the ReformerEmbeddings class.
        input_ids (Optional[Tensor]): The input tensor representing the tokenized input sequence.
            It has shape [batch_size, sequence_length] and each element corresponds to a token ID.
        position_ids (Optional[Tensor]): The tensor representing the position encodings for each token
            in the input sequence. It has shape [batch_size, sequence_length] and each element corresponds
            to a position ID.
        inputs_embeds (Optional[Tensor]): The tensor representing the embedded input sequence.
            It has shape [batch_size, sequence_length, embedding_size] and each element corresponds
            to an embedding vector.
        start_idx_pos_encodings (int): The starting index for the positional encodings.

    Returns:
        Tensor: The embeddings tensor representing the input sequence with positional information.
            It has shape [batch_size, sequence_length, embedding_size].

    Raises:
        ValueError: If the sequence length of the position_ids tensor is greater than the maximum allowed
            sequence length.
    """
    if input_ids is not None:
        input_shape = input_ids.shape
    else:
        input_shape = inputs_embeds.shape[:-1]

    seq_length = input_shape[1]
    if position_ids is None:
        position_ids = ops.arange(
            start_idx_pos_encodings, start_idx_pos_encodings + seq_length, dtype=mindspore.int64
        )
        position_ids = position_ids.unsqueeze(0).expand(input_shape)

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

    if position_ids.shape[-1] > self.max_position_embeddings:
        raise ValueError(
            f"Sequence Length: {position_ids.shape[-1]} has to be less or equal than "
            f"config.max_position_embeddings {self.max_position_embeddings}."
        )

    # dropout
    embeddings = ops.dropout(inputs_embeds, p=self.dropout, training=self.training)

    # add positional embeddings
    position_embeddings = self.position_embeddings(position_ids)
    embeddings = embeddings + position_embeddings
    return embeddings

mindnlp.transformers.models.reformer.modeling_reformer.ReformerEncoder

Bases: Module

The 'ReformerEncoder' class is a Python class that represents the encoder component of the Reformer model. It inherits from the 'nn.Module' class.

ATTRIBUTE	DESCRIPTION
dropout	The dropout probability for the hidden states. TYPE: float
layers	A list of 'ReformerLayer' instances representing the layers of the encoder. TYPE: ModuleList
layer_norm	A layer normalization module. TYPE: LayerNorm

METHOD	DESCRIPTION
__init__	Initializes a new instance of the 'ReformerEncoder' class.
forward	Constructs the encoder by applying the Reformer layers to the input hidden states.

Source code in mindnlp/transformers/models/reformer/modeling_reformer.py

class ReformerEncoder(nn.Module):

    """
    The 'ReformerEncoder' class is a Python class that represents the encoder component of the Reformer model.
    It inherits from the 'nn.Module' class.

    Attributes:
        dropout (float): The dropout probability for the hidden states.
        layers (nn.ModuleList): A list of 'ReformerLayer' instances representing the layers of the encoder.
        layer_norm (nn.LayerNorm): A layer normalization module.

    Methods:
        __init__: Initializes a new instance of the 'ReformerEncoder' class.
        forward: Constructs the encoder by applying the Reformer layers to the input hidden states.

    """
    def __init__(self, config):

        """
        Initializes a ReformerEncoder instance.

        Args:
            self (ReformerEncoder): The ReformerEncoder instance to be initialized.
            config (Config):
                A configuration object containing settings for the ReformerEncoder.

                - config.hidden_dropout_prob (float): The dropout probability for hidden layers.
                - config.num_hidden_layers (int): The number of hidden layers in the ReformerEncoder.
                - config.hidden_size (int): The size of hidden layers.
                - config.layer_norm_eps (float): The epsilon value for layer normalization.

        Returns:
            None.

        Raises:
            TypeError: If config is not of type Config.
            ValueError: If config is missing any required attributes.
            ValueError: If config.hidden_dropout_prob is not a float.
            ValueError: If config.num_hidden_layers is not an integer.
            ValueError: If config.hidden_size is not an integer.
            ValueError: If config.layer_norm_eps is not a float.
        """
        super().__init__()
        self.dropout = float(config.hidden_dropout_prob)

        self.layers = nn.ModuleList([ReformerLayer(config, i) for i in range(config.num_hidden_layers)])
        # Reformer is using Rev Nets, thus last layer outputs are concatenated and
        # Layer Norm is done over 2 * hidden_size
        self.layer_norm = nn.LayerNorm(2 * config.hidden_size, eps=config.layer_norm_eps)

    def forward(
        self,
        hidden_states,
        attention_mask=None,
        head_mask=None,
        num_hashes=None,
        past_buckets_states=None,
        use_cache=False,
        orig_sequence_length=None,
        output_hidden_states=False,
        output_attentions=False,
    ):

        """
        Constructs the ReformerEncoder output given the input parameters.

        Args:
            self: The class instance.
            hidden_states (Tensor): The input hidden states. Shape [batch_size, sequence_length, hidden_size].
            attention_mask (Tensor, optional): The attention mask. Shape [batch_size, sequence_length, sequence_length].
                Masks the attention scores for padding tokens. Defaults to None.
            head_mask (Tensor, optional): The head mask. Shape [num_attention_heads, sequence_length, sequence_length].
                Masks the attention scores for specific attention heads. Defaults to None.
            num_hashes (int, optional): The number of hashes to use for LSH attention. Defaults to None.
            past_buckets_states (list, optional): The list of past bucket states. Each element is a tuple of two tensors.
                Shape [(buckets, num_hashes, sequence_length // bucket_size, embedding_dim),
                (buckets, num_hashes, sequence_length // bucket_size)].
                Defaults to None.
            use_cache (bool, optional): Whether to use cache for fast decoding. Defaults to False.
            orig_sequence_length (int, optional): The original sequence length before padding. Defaults to None.
            output_hidden_states (bool, optional): Whether to output all hidden states. Defaults to False.
            output_attentions (bool, optional): Whether to output all attention matrices. Defaults to False.

        Returns:
            ReformerEncoderOutput:
                An instance of the ReformerEncoderOutput class containing the following attributes:

                - hidden_states (Tensor): The output hidden states. Shape [batch_size, sequence_length, hidden_size].
                - all_hidden_states (list): List of hidden states at each layer. Each element has shape
                [batch_size, sequence_length, hidden_size].
                - all_attentions (list): List of attention matrices at each layer. Each element has shape
                [batch_size, num_attention_heads, sequence_length, sequence_length].
                - past_buckets_states (list): List of past bucket states for fast decoding.
                Each element is a tuple of two tensors.
                Shape [(buckets, num_hashes, sequence_length // bucket_size, embedding_dim),
                (buckets, num_hashes, sequence_length // bucket_size)].

        Raises:
            None.
        """
        # hidden_states and attention lists to be filled if wished
        all_hidden_states = []
        all_attentions = []

        # init cached hidden states if necessary
        if past_buckets_states is None:
            past_buckets_states = [((None), (None)) for i in range(len(self.layers))]

        # concat same tensor for reversible ResNet
        hidden_states = ops.cat([hidden_states, hidden_states], axis=-1)
        hidden_states = _ReversibleFunction()(
            hidden_states,
            self.layers,
            attention_mask,
            head_mask,
            num_hashes,
            all_hidden_states,
            all_attentions,
            past_buckets_states,
            use_cache,
            orig_sequence_length,
            output_hidden_states,
            output_attentions,
        )
        # Apply layer norm to concatenated hidden states
        hidden_states = self.layer_norm(hidden_states)

        # Apply dropout
        hidden_states = ops.dropout(hidden_states, p=self.dropout, training=self.training)

        return ReformerEncoderOutput(
            hidden_states=hidden_states,
            all_hidden_states=all_hidden_states,
            all_attentions=all_attentions,
            past_buckets_states=past_buckets_states,
        )

mindnlp.transformers.models.reformer.modeling_reformer.ReformerEncoder.__init__(config)

Initializes a ReformerEncoder instance.

PARAMETER	DESCRIPTION
self	The ReformerEncoder instance to be initialized. TYPE: ReformerEncoder
config	A configuration object containing settings for the ReformerEncoder. config.hidden_dropout_prob (float): The dropout probability for hidden layers. config.num_hidden_layers (int): The number of hidden layers in the ReformerEncoder. config.hidden_size (int): The size of hidden layers. config.layer_norm_eps (float): The epsilon value for layer normalization. TYPE: Config

RETURNS	DESCRIPTION
	None.

RAISES	DESCRIPTION
TypeError	If config is not of type Config.
ValueError	If config is missing any required attributes.
ValueError	If config.hidden_dropout_prob is not a float.
ValueError	If config.num_hidden_layers is not an integer.
ValueError	If config.hidden_size is not an integer.
ValueError	If config.layer_norm_eps is not a float.

Source code in mindnlp/transformers/models/reformer/modeling_reformer.py

def __init__(self, config):

    """
    Initializes a ReformerEncoder instance.

    Args:
        self (ReformerEncoder): The ReformerEncoder instance to be initialized.
        config (Config):
            A configuration object containing settings for the ReformerEncoder.

            - config.hidden_dropout_prob (float): The dropout probability for hidden layers.
            - config.num_hidden_layers (int): The number of hidden layers in the ReformerEncoder.
            - config.hidden_size (int): The size of hidden layers.
            - config.layer_norm_eps (float): The epsilon value for layer normalization.

    Returns:
        None.

    Raises:
        TypeError: If config is not of type Config.
        ValueError: If config is missing any required attributes.
        ValueError: If config.hidden_dropout_prob is not a float.
        ValueError: If config.num_hidden_layers is not an integer.
        ValueError: If config.hidden_size is not an integer.
        ValueError: If config.layer_norm_eps is not a float.
    """
    super().__init__()
    self.dropout = float(config.hidden_dropout_prob)

    self.layers = nn.ModuleList([ReformerLayer(config, i) for i in range(config.num_hidden_layers)])
    # Reformer is using Rev Nets, thus last layer outputs are concatenated and
    # Layer Norm is done over 2 * hidden_size
    self.layer_norm = nn.LayerNorm(2 * config.hidden_size, eps=config.layer_norm_eps)

mindnlp.transformers.models.reformer.modeling_reformer.ReformerEncoder.forward(hidden_states, attention_mask=None, head_mask=None, num_hashes=None, past_buckets_states=None, use_cache=False, orig_sequence_length=None, output_hidden_states=False, output_attentions=False)

Constructs the ReformerEncoder output given the input parameters.

PARAMETER	DESCRIPTION
self	The class instance.
hidden_states	The input hidden states. Shape [batch_size, sequence_length, hidden_size]. TYPE: Tensor
attention_mask	The attention mask. Shape [batch_size, sequence_length, sequence_length]. Masks the attention scores for padding tokens. Defaults to None. TYPE: Tensor DEFAULT: None
head_mask	The head mask. Shape [num_attention_heads, sequence_length, sequence_length]. Masks the attention scores for specific attention heads. Defaults to None. TYPE: Tensor DEFAULT: None
num_hashes	The number of hashes to use for LSH attention. Defaults to None. TYPE: int DEFAULT: None
past_buckets_states	The list of past bucket states. Each element is a tuple of two tensors. Shape [(buckets, num_hashes, sequence_length // bucket_size, embedding_dim), (buckets, num_hashes, sequence_length // bucket_size)]. Defaults to None. TYPE: list DEFAULT: None
use_cache	Whether to use cache for fast decoding. Defaults to False. TYPE: bool DEFAULT: False
orig_sequence_length	The original sequence length before padding. Defaults to None. TYPE: int DEFAULT: None
output_hidden_states	Whether to output all hidden states. Defaults to False. TYPE: bool DEFAULT: False
output_attentions	Whether to output all attention matrices. Defaults to False. TYPE: bool DEFAULT: False

RETURNS

DESCRIPTION

ReformerEncoderOutput

An instance of the ReformerEncoderOutput class containing the following attributes: hidden_states (Tensor): The output hidden states. Shape [batch_size, sequence_length, hidden_size]. all_hidden_states (list): List of hidden states at each layer. Each element has shape [batch_size, sequence_length, hidden_size]. all_attentions (list): List of attention matrices at each layer. Each element has shape [batch_size, num_attention_heads, sequence_length, sequence_length]. past_buckets_states (list): List of past bucket states for fast decoding. Each element is a tuple of two tensors. Shape [(buckets, num_hashes, sequence_length // bucket_size, embedding_dim), (buckets, num_hashes, sequence_length // bucket_size)].

Source code in mindnlp/transformers/models/reformer/modeling_reformer.py

def forward(
    self,
    hidden_states,
    attention_mask=None,
    head_mask=None,
    num_hashes=None,
    past_buckets_states=None,
    use_cache=False,
    orig_sequence_length=None,
    output_hidden_states=False,
    output_attentions=False,
):

    """
    Constructs the ReformerEncoder output given the input parameters.

    Args:
        self: The class instance.
        hidden_states (Tensor): The input hidden states. Shape [batch_size, sequence_length, hidden_size].
        attention_mask (Tensor, optional): The attention mask. Shape [batch_size, sequence_length, sequence_length].
            Masks the attention scores for padding tokens. Defaults to None.
        head_mask (Tensor, optional): The head mask. Shape [num_attention_heads, sequence_length, sequence_length].
            Masks the attention scores for specific attention heads. Defaults to None.
        num_hashes (int, optional): The number of hashes to use for LSH attention. Defaults to None.
        past_buckets_states (list, optional): The list of past bucket states. Each element is a tuple of two tensors.
            Shape [(buckets, num_hashes, sequence_length // bucket_size, embedding_dim),
            (buckets, num_hashes, sequence_length // bucket_size)].
            Defaults to None.
        use_cache (bool, optional): Whether to use cache for fast decoding. Defaults to False.
        orig_sequence_length (int, optional): The original sequence length before padding. Defaults to None.
        output_hidden_states (bool, optional): Whether to output all hidden states. Defaults to False.
        output_attentions (bool, optional): Whether to output all attention matrices. Defaults to False.

    Returns:
        ReformerEncoderOutput:
            An instance of the ReformerEncoderOutput class containing the following attributes:

            - hidden_states (Tensor): The output hidden states. Shape [batch_size, sequence_length, hidden_size].
            - all_hidden_states (list): List of hidden states at each layer. Each element has shape
            [batch_size, sequence_length, hidden_size].
            - all_attentions (list): List of attention matrices at each layer. Each element has shape
            [batch_size, num_attention_heads, sequence_length, sequence_length].
            - past_buckets_states (list): List of past bucket states for fast decoding.
            Each element is a tuple of two tensors.
            Shape [(buckets, num_hashes, sequence_length // bucket_size, embedding_dim),
            (buckets, num_hashes, sequence_length // bucket_size)].

    Raises:
        None.
    """
    # hidden_states and attention lists to be filled if wished
    all_hidden_states = []
    all_attentions = []

    # init cached hidden states if necessary
    if past_buckets_states is None:
        past_buckets_states = [((None), (None)) for i in range(len(self.layers))]

    # concat same tensor for reversible ResNet
    hidden_states = ops.cat([hidden_states, hidden_states], axis=-1)
    hidden_states = _ReversibleFunction()(
        hidden_states,
        self.layers,
        attention_mask,
        head_mask,
        num_hashes,
        all_hidden_states,
        all_attentions,
        past_buckets_states,
        use_cache,
        orig_sequence_length,
        output_hidden_states,
        output_attentions,
    )
    # Apply layer norm to concatenated hidden states
    hidden_states = self.layer_norm(hidden_states)

    # Apply dropout
    hidden_states = ops.dropout(hidden_states, p=self.dropout, training=self.training)

    return ReformerEncoderOutput(
        hidden_states=hidden_states,
        all_hidden_states=all_hidden_states,
        all_attentions=all_attentions,
        past_buckets_states=past_buckets_states,
    )

mindnlp.transformers.models.reformer.modeling_reformer.ReformerFeedForwardDense

Bases: Module

ReformerFeedForwardDense represents a feedforward dense layer used in a Reformer model for neural network operations.

ATTRIBUTE	DESCRIPTION
dropout	The dropout rate for the hidden states. TYPE: float
act_fn	The activation function used for the hidden states. TYPE: function
dense	The dense layer for transforming hidden states. TYPE: Linear

METHOD	DESCRIPTION
__init__	Initializes the ReformerFeedForwardDense instance with the provided configuration.
forward	Constructs the feedforward dense layer by applying dense transformation, dropout, and activation function to the hidden states.

This class inherits from nn.Module and includes methods to initialize and forward the feedforward dense layer in a Reformer model.

Source code in mindnlp/transformers/models/reformer/modeling_reformer.py

class ReformerFeedForwardDense(nn.Module):

    """
    ReformerFeedForwardDense represents a feedforward dense layer used in a Reformer model for neural network operations.

    Attributes:
        dropout (float): The dropout rate for the hidden states.
        act_fn (function): The activation function used for the hidden states.
        dense (nn.Linear): The dense layer for transforming hidden states.

    Methods:
        __init__: Initializes the ReformerFeedForwardDense instance with the provided configuration.
        forward: Constructs the feedforward dense layer by applying dense transformation, dropout,
            and activation function to the hidden states.

    This class inherits from nn.Module and includes methods to initialize and forward the feedforward dense layer in
    a Reformer model.
    """
    def __init__(self, config):

        """
        Initializes a ReformerFeedForwardDense object with the specified configuration.

        Args:
            self (ReformerFeedForwardDense): The instance of the ReformerFeedForwardDense class.
            config (object): The configuration object containing various settings for the dense layer.
                Expected to have the following attributes:

                - hidden_dropout_prob (float): The dropout probability for the hidden layer.
                - hidden_act (str or function): The activation function for the hidden layer.
                - hidden_size (int): The size of the hidden layer.
                - feed_forward_size (int): The size of the feed-forward layer.

        Returns:
            None.

        Raises:
            TypeError: If the config.hidden_dropout_prob is not a float.
            KeyError: If the config.hidden_act is not a valid activation function name.
            AttributeError: If the config object is missing any of the required attributes.
        """
        super().__init__()
        self.dropout = float(config.hidden_dropout_prob)

        if isinstance(config.hidden_act, str):
            self.act_fn = ACT2FN[config.hidden_act]
        else:
            self.act_fn = config.hidden_act

        self.dense = nn.Linear(config.hidden_size, config.feed_forward_size)

    def forward(self, hidden_states):

        """
        Constructs the feedforward dense layer for the Reformer model.

        Args:
            self (ReformerFeedForwardDense): An instance of the ReformerFeedForwardDense class.
            hidden_states (tensor): The input hidden states to be processed by the feedforward dense layer.

        Returns:
            tensor: The processed hidden states after passing through the feedforward dense layer.

        Raises:
            ValueError: If the hidden_states tensor is not provided.
            TypeError: If the input hidden_states tensor is not of type tensor.
            RuntimeError: If an error occurs during the dropout operation.
        """
        hidden_states = self.dense(hidden_states)
        hidden_states = ops.dropout(hidden_states, p=self.dropout, training=self.training)
        hidden_states = self.act_fn(hidden_states)
        return hidden_states

mindnlp.transformers.models.reformer.modeling_reformer.ReformerFeedForwardDense.__init__(config)

Initializes a ReformerFeedForwardDense object with the specified configuration.

PARAMETER	DESCRIPTION
self	The instance of the ReformerFeedForwardDense class. TYPE: ReformerFeedForwardDense
config	The configuration object containing various settings for the dense layer. Expected to have the following attributes: hidden_dropout_prob (float): The dropout probability for the hidden layer. hidden_act (str or function): The activation function for the hidden layer. hidden_size (int): The size of the hidden layer. feed_forward_size (int): The size of the feed-forward layer. TYPE: object

RETURNS	DESCRIPTION
	None.

RAISES	DESCRIPTION
TypeError	If the config.hidden_dropout_prob is not a float.
KeyError	If the config.hidden_act is not a valid activation function name.
AttributeError	If the config object is missing any of the required attributes.

Source code in mindnlp/transformers/models/reformer/modeling_reformer.py

def __init__(self, config):

    """
    Initializes a ReformerFeedForwardDense object with the specified configuration.

    Args:
        self (ReformerFeedForwardDense): The instance of the ReformerFeedForwardDense class.
        config (object): The configuration object containing various settings for the dense layer.
            Expected to have the following attributes:

            - hidden_dropout_prob (float): The dropout probability for the hidden layer.
            - hidden_act (str or function): The activation function for the hidden layer.
            - hidden_size (int): The size of the hidden layer.
            - feed_forward_size (int): The size of the feed-forward layer.

    Returns:
        None.

    Raises:
        TypeError: If the config.hidden_dropout_prob is not a float.
        KeyError: If the config.hidden_act is not a valid activation function name.
        AttributeError: If the config object is missing any of the required attributes.
    """
    super().__init__()
    self.dropout = float(config.hidden_dropout_prob)

    if isinstance(config.hidden_act, str):
        self.act_fn = ACT2FN[config.hidden_act]
    else:
        self.act_fn = config.hidden_act

    self.dense = nn.Linear(config.hidden_size, config.feed_forward_size)

mindnlp.transformers.models.reformer.modeling_reformer.ReformerFeedForwardDense.forward(hidden_states)

Constructs the feedforward dense layer for the Reformer model.

PARAMETER	DESCRIPTION
self	An instance of the ReformerFeedForwardDense class. TYPE: ReformerFeedForwardDense
hidden_states	The input hidden states to be processed by the feedforward dense layer. TYPE: tensor

RETURNS	DESCRIPTION
tensor	The processed hidden states after passing through the feedforward dense layer.

RAISES	DESCRIPTION
ValueError	If the hidden_states tensor is not provided.
TypeError	If the input hidden_states tensor is not of type tensor.
RuntimeError	If an error occurs during the dropout operation.

Source code in mindnlp/transformers/models/reformer/modeling_reformer.py

def forward(self, hidden_states):

    """
    Constructs the feedforward dense layer for the Reformer model.

    Args:
        self (ReformerFeedForwardDense): An instance of the ReformerFeedForwardDense class.
        hidden_states (tensor): The input hidden states to be processed by the feedforward dense layer.

    Returns:
        tensor: The processed hidden states after passing through the feedforward dense layer.

    Raises:
        ValueError: If the hidden_states tensor is not provided.
        TypeError: If the input hidden_states tensor is not of type tensor.
        RuntimeError: If an error occurs during the dropout operation.
    """
    hidden_states = self.dense(hidden_states)
    hidden_states = ops.dropout(hidden_states, p=self.dropout, training=self.training)
    hidden_states = self.act_fn(hidden_states)
    return hidden_states

mindnlp.transformers.models.reformer.modeling_reformer.ReformerFeedForwardOutput

Bases: Module

Represents the output of the feed forward layer in a Reformer neural network.

This class inherits from nn.Module and contains methods for initializing and forwarding the feed forward layer output.

ATTRIBUTE	DESCRIPTION
dropout	The dropout rate for the hidden units. TYPE: float
dense	The fully connected layer for transforming input hidden states. TYPE: Linear

METHOD	DESCRIPTION
__init__	Initializes the ReformerFeedForwardOutput with the given configuration.
forward	Constructs the output of the feed forward layer using the provided hidden states.

Source code in mindnlp/transformers/models/reformer/modeling_reformer.py

class ReformerFeedForwardOutput(nn.Module):

    """
    Represents the output of the feed forward layer in a Reformer neural network.

    This class inherits from nn.Module and contains methods for initializing and forwarding the feed forward layer output.

    Attributes:
        dropout (float): The dropout rate for the hidden units.
        dense (nn.Linear): The fully connected layer for transforming input hidden states.

    Methods:
        __init__: Initializes the ReformerFeedForwardOutput with the given configuration.
        forward: Constructs the output of the feed forward layer using the provided hidden states.
    """
    def __init__(self, config):

        """
        Initializes an instance of the ReformerFeedForwardOutput class.

        Args:
            self: The object instance.
            config:
                An object containing the configuration parameters.

                - Type: Any valid object.
                - Purpose: Specifies the configuration settings for the ReformerFeedForwardOutput instance.
                - Restrictions: None.

        Returns:
            None

        Raises:
            None
        """
        super().__init__()
        self.dropout = float(config.hidden_dropout_prob)

        self.dense = nn.Linear(config.feed_forward_size, config.hidden_size)

    def forward(self, hidden_states):

        """
        Constructs the output of the feed-forward layer in the Reformer model.

        Args:
            self (ReformerFeedForwardOutput): The instance of the ReformerFeedForwardOutput class.
            hidden_states (tensor): The input hidden states to be processed by the feed-forward layer.

        Returns:
            tensor: The processed hidden states after passing through the feed-forward layer.

        Raises:
            ValueError: If the hidden_states tensor is not valid or has incorrect dimensions.
            RuntimeError: If an error occurs during the computation of the output tensor.
        """
        hidden_states = self.dense(hidden_states)
        hidden_states = ops.dropout(hidden_states, p=self.dropout, training=self.training)
        return hidden_states

mindnlp.transformers.models.reformer.modeling_reformer.ReformerFeedForwardOutput.__init__(config)

Initializes an instance of the ReformerFeedForwardOutput class.

PARAMETER	DESCRIPTION
self	The object instance.
config	An object containing the configuration parameters. Type: Any valid object. Purpose: Specifies the configuration settings for the ReformerFeedForwardOutput instance. Restrictions: None.

RETURNS	DESCRIPTION
	None

Source code in mindnlp/transformers/models/reformer/modeling_reformer.py

def __init__(self, config):

    """
    Initializes an instance of the ReformerFeedForwardOutput class.

    Args:
        self: The object instance.
        config:
            An object containing the configuration parameters.

            - Type: Any valid object.
            - Purpose: Specifies the configuration settings for the ReformerFeedForwardOutput instance.
            - Restrictions: None.

    Returns:
        None

    Raises:
        None
    """
    super().__init__()
    self.dropout = float(config.hidden_dropout_prob)

    self.dense = nn.Linear(config.feed_forward_size, config.hidden_size)

mindnlp.transformers.models.reformer.modeling_reformer.ReformerFeedForwardOutput.forward(hidden_states)

Constructs the output of the feed-forward layer in the Reformer model.

PARAMETER	DESCRIPTION
self	The instance of the ReformerFeedForwardOutput class. TYPE: ReformerFeedForwardOutput
hidden_states	The input hidden states to be processed by the feed-forward layer. TYPE: tensor

RETURNS	DESCRIPTION
tensor	The processed hidden states after passing through the feed-forward layer.

RAISES	DESCRIPTION
ValueError	If the hidden_states tensor is not valid or has incorrect dimensions.
RuntimeError	If an error occurs during the computation of the output tensor.

Source code in mindnlp/transformers/models/reformer/modeling_reformer.py

def forward(self, hidden_states):

    """
    Constructs the output of the feed-forward layer in the Reformer model.

    Args:
        self (ReformerFeedForwardOutput): The instance of the ReformerFeedForwardOutput class.
        hidden_states (tensor): The input hidden states to be processed by the feed-forward layer.

    Returns:
        tensor: The processed hidden states after passing through the feed-forward layer.

    Raises:
        ValueError: If the hidden_states tensor is not valid or has incorrect dimensions.
        RuntimeError: If an error occurs during the computation of the output tensor.
    """
    hidden_states = self.dense(hidden_states)
    hidden_states = ops.dropout(hidden_states, p=self.dropout, training=self.training)
    return hidden_states

mindnlp.transformers.models.reformer.modeling_reformer.ReformerForMaskedLM

Bases: ReformerPreTrainedModel

A Reformer model with a language modeling head for masked language modeling tasks.

This class inherits from ReformerPreTrainedModel and utilizes the Reformer architecture and a language modeling head for masking language modeling tasks. The class is capable of generating output embeddings and setting new embeddings with the provided methods, get_output_embeddings() and set_output_embeddings(), respectively, and the forward() method forwards the model and computes the masked language modeling loss if the labels argument is provided.

The ReformerForMaskedLM class takes a config argument, which is an instance of ReformerConfig. The class implements the __init__() method that initializes the parent class with the provided config. The method also checks that config.is_decoder=False for bi-directional self-attention.

The class has the following methods:

• get_output_embeddings(): Returns the decoder for the language modeling head.
• set_output_embeddings(new_embeddings): Sets the decoder for the language modeling head to new_embeddings.

• forward(input_ids=None, position_ids=None, attention_mask=None, head_mask=None, inputs_embeds=None, num_hashes=None, labels=None, output_hidden_states=None, output_attentions=None, return_dict=None): Constructs the model and computes the masked language modeling loss if the labels argument is provided. This method takes several optional input arguments and returns a tuple with the following elements:

  • If labels is not None, returns a tuple (masked_lm_loss, logits, hidden_states, attentions).

  • If labels is None, returns a tuple (logits, hidden_states, attentions).

    • masked_lm_loss is the masked language modeling loss, computed as cross-entropy loss between the logits and the labels.
    • logits is a tensor of shape (batch_size, sequence_length, vocab_size) containing the unnormalized log probabilities for each vocabulary token.
    • hidden_states is a tuple of length config.num_hidden_layers + 1 with each tensor of shape (batch_size, sequence_length, hidden_size) representing the hidden states of the model at each layer.
    • attentions is a tuple of length config.num_hidden_layers with each tensor of shape (batch_size, num_heads, sequence_length, sequence_length) representing the attention weights for each layer.

Note

This class utilizes a false checkpoint since there is no available pre-trained model for the masked language modeling task with the Reformer architecture.

Example

>>> from transformers import ReformerForMaskedLM, ReformerConfig
...
>>> # Initializing a Reformer configuration
>>> config = ReformerConfig()
...
>>> # Initializing a ReformerForMaskedLM model with the configuration
>>> model = ReformerForMaskedLM(config)
...
>>> # Getting the decoder for the language modeling head
>>> decoder = model.get_output_embeddings()
...
>>> # Setting new embeddings for the language modeling head
>>> model.set_output_embeddings(new_embeddings)
...
>>> # Constructing the model and computing the masked language modeling loss
>>> masked_lm_loss, logits, hidden_states, attentions = model.forward(input_ids, position_ids, attention_mask,
...     head_mask, inputs_embeds, num_hashes, labels, output_hidden_states, output_attentions, return_dict)

Source code in mindnlp/transformers/models/reformer/modeling_reformer.py

class ReformerForMaskedLM(ReformerPreTrainedModel):
    r"""
    A Reformer model with a language modeling head for masked language modeling tasks.

    This class inherits from `ReformerPreTrainedModel` and utilizes the Reformer architecture and a language modeling head for
    masking language modeling tasks. The class is capable of generating output embeddings and setting new embeddings with
    the provided methods, `get_output_embeddings()` and `set_output_embeddings()`, respectively, and the `forward()`
    method forwards the model and computes the masked language modeling loss if the `labels` argument is provided.

    The `ReformerForMaskedLM` class takes a `config` argument, which is an instance of `ReformerConfig`. The class
    implements the `__init__()` method that initializes the parent class with the provided `config`. The method also
    checks that `config.is_decoder=False` for bi-directional self-attention.

    The class has the following methods:

    - `get_output_embeddings()`: Returns the decoder for the language modeling head.
    - `set_output_embeddings(new_embeddings)`: Sets the decoder for the language modeling head to `new_embeddings`.
    - `forward(input_ids=None, position_ids=None, attention_mask=None, head_mask=None, inputs_embeds=None,
    num_hashes=None, labels=None, output_hidden_states=None, output_attentions=None, return_dict=None)`:
    Constructs the model and computes the masked language modeling loss if the `labels` argument is provided.
    This method takes several optional input arguments and returns a tuple with the following elements:

        - If `labels` is not `None`, returns a tuple `(masked_lm_loss, logits, hidden_states, attentions)`.
        - If `labels` is `None`, returns a tuple `(logits, hidden_states, attentions)`.

            - `masked_lm_loss` is the masked language modeling loss, computed as cross-entropy loss between the logits
              and the labels.
            - `logits` is a tensor of shape `(batch_size, sequence_length, vocab_size)` containing the unnormalized
              log probabilities for each vocabulary token.
            - `hidden_states` is a tuple of length `config.num_hidden_layers + 1` with each tensor of shape
              `(batch_size, sequence_length, hidden_size)` representing the hidden states of the model at each layer.
            - `attentions` is a tuple of length `config.num_hidden_layers` with each tensor of shape
              `(batch_size, num_heads, sequence_length, sequence_length)` representing the attention weights for each
              layer.

    Note:
        This class utilizes a false checkpoint since there is no available pre-trained model for the masked language
        modeling task with the Reformer architecture.

    Example:
        ```python
        >>> from transformers import ReformerForMaskedLM, ReformerConfig
        ...
        >>> # Initializing a Reformer configuration
        >>> config = ReformerConfig()
        ...
        >>> # Initializing a ReformerForMaskedLM model with the configuration
        >>> model = ReformerForMaskedLM(config)
        ...
        >>> # Getting the decoder for the language modeling head
        >>> decoder = model.get_output_embeddings()
        ...
        >>> # Setting new embeddings for the language modeling head
        >>> model.set_output_embeddings(new_embeddings)
        ...
        >>> # Constructing the model and computing the masked language modeling loss
        >>> masked_lm_loss, logits, hidden_states, attentions = model.forward(input_ids, position_ids, attention_mask,
        ...     head_mask, inputs_embeds, num_hashes, labels, output_hidden_states, output_attentions, return_dict)
        ```

    """
    _tied_weights_keys = ["lm_head.decoder.weight", "lm_head.decoder.bias"]

    def __init__(self, config):

        """
        Initializes an instance of the ReformerForMaskedLM class.

        Args:
            self: The instance of the class.
            config: An instance of the configuration class that specifies the model configuration.

        Returns:
            None

        Raises:
            AssertionError: If `config.is_decoder` is set to True. To use the ReformerForMaskedLM, `config.is_decoder`
                should be set to False for bi-directional self-attention.

        """
        super().__init__(config)
        assert not config.is_decoder, (
            "If you want to use `ReformerForMaskedLM` make sure `config.is_decoder=False` for bi-directional"
            " self-attention."
        )
        self.reformer = ReformerModel(config)
        self.lm_head = ReformerOnlyLMHead(config)

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

    def get_output_embeddings(self):

        """
        This method retrieves the output embeddings from the ReformerForMaskedLM model.

        Args:
            self: ReformerForMaskedLM - The instance of the ReformerForMaskedLM class.

        Returns:
            None: This method returns None as it only retrieves the output embeddings without any additional processing.

        Raises:
            None.
        """
        return self.lm_head.decoder

    def set_output_embeddings(self, new_embeddings):

        """
        Sets the output embeddings for the Reformer model.

        Args:
            self (ReformerForMaskedLM): The instance of the ReformerForMaskedLM class.
            new_embeddings: The new embeddings to be set as the output embeddings. It can be of any type.

        Returns:
            None.

        Raises:
            None.

        This method sets the output embeddings of the ReformerForMaskedLM model to the provided new_embeddings.
        The new_embeddings can be any type and will be assigned to the decoder of the lm_head.
        """
        self.lm_head.decoder = new_embeddings

    def forward(
        self,
        input_ids: Optional[mindspore.Tensor] = None,
        position_ids: Optional[mindspore.Tensor] = None,
        attention_mask: Optional[mindspore.Tensor] = None,
        head_mask: Optional[mindspore.Tensor] = None,
        inputs_embeds: Optional[mindspore.Tensor] = None,
        num_hashes: Optional[int] = None,
        labels: Optional[mindspore.Tensor] = None,
        output_hidden_states: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, MaskedLMOutput]:
        r"""
        Args:
            labels (`mindspore.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
                Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
                config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked),
                the loss is only computed for the tokens with labels

        Returns:
            Union[Tuple, MaskedLMOutput]

        <Tip warning={true}>

        This example uses a false checkpoint since we don't have any available pretrained model for the masked language
        modeling task with the Reformer architecture.

        </Tip>

        Example:
            ```python
            >>> import torch
            >>> from transformers import AutoTokenizer, ReformerForMaskedLM
            ...
            >>> tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-reformer")
            >>> model = ReformerForMaskedLM.from_pretrained("hf-internal-testing/tiny-random-reformer")
            ...
            >>> # add mask_token
            >>> tokenizer.add_special_tokens({"mask_token": "[MASK]"})  # doctest: +IGNORE_RESULT
            >>> inputs = tokenizer("The capital of France is [MASK].", return_tensors="pt")
            ...
            >>> # resize model's embedding matrix
            >>> model.resize_token_embeddings(new_num_tokens=model.config.vocab_size + 1)  # doctest: +IGNORE_RESULT
            ...
            >>> with torch.no_grad():
            ...     logits = model(**inputs).logits
            ...
            >>> # retrieve index of [MASK]
            >>> mask_token_index = (inputs.input_ids == tokenizer.mask_token_id)[0].nonzero(as_tuple=True)[0]
            ...
            >>> predicted_token_id = logits[0, mask_token_index].argmax(axis=-1)
            >>> predicted_token = tokenizer.decode(predicted_token_id)
            ```

            ```python
            >>> labels = tokenizer("The capital of France is Paris.", return_tensors="pt")["input_ids"]
            >>> # mask labels of non-[MASK] tokens
            >>> labels = torch.where(
            ...     inputs.input_ids == tokenizer.mask_token_id, labels[:, : inputs["input_ids"].shape[-1]], -100
            ... )
            ...
            >>> outputs = model(**inputs, labels=labels)
            >>> loss = round(outputs.loss.item(), 2)
            ```
        """
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        reformer_outputs = self.reformer(
            input_ids,
            position_ids=position_ids,
            attention_mask=attention_mask,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            num_hashes=num_hashes,
            use_cache=False,  # no causal mask
            output_hidden_states=output_hidden_states,
            output_attentions=output_attentions,
            return_dict=return_dict,
        )

        sequence_output = reformer_outputs[0]
        logits = self.lm_head(sequence_output)

        masked_lm_loss = None
        if labels is not None:
            masked_lm_loss = ops.cross_entropy(logits.view(-1, self.config.vocab_size), labels.view(-1))

        if not return_dict:
            output = (logits,) + reformer_outputs[1:]
            return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output

        return MaskedLMOutput(
            loss=masked_lm_loss,
            logits=logits,
            hidden_states=reformer_outputs.hidden_states,
            attentions=reformer_outputs.attentions,
        )

mindnlp.transformers.models.reformer.modeling_reformer.ReformerForMaskedLM.__init__(config)

Initializes an instance of the ReformerForMaskedLM class.

PARAMETER	DESCRIPTION
self	The instance of the class.
config	An instance of the configuration class that specifies the model configuration.

RETURNS	DESCRIPTION
	None

RAISES	DESCRIPTION
AssertionError	If config.is_decoder is set to True. To use the ReformerForMaskedLM, config.is_decoder should be set to False for bi-directional self-attention.

Source code in mindnlp/transformers/models/reformer/modeling_reformer.py

def __init__(self, config):

    """
    Initializes an instance of the ReformerForMaskedLM class.

    Args:
        self: The instance of the class.
        config: An instance of the configuration class that specifies the model configuration.

    Returns:
        None

    Raises:
        AssertionError: If `config.is_decoder` is set to True. To use the ReformerForMaskedLM, `config.is_decoder`
            should be set to False for bi-directional self-attention.

    """
    super().__init__(config)
    assert not config.is_decoder, (
        "If you want to use `ReformerForMaskedLM` make sure `config.is_decoder=False` for bi-directional"
        " self-attention."
    )
    self.reformer = ReformerModel(config)
    self.lm_head = ReformerOnlyLMHead(config)

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

mindnlp.transformers.models.reformer.modeling_reformer.ReformerForMaskedLM.forward(input_ids=None, position_ids=None, attention_mask=None, head_mask=None, inputs_embeds=None, num_hashes=None, labels=None, output_hidden_states=None, output_attentions=None, return_dict=None)

PARAMETER	DESCRIPTION
labels	Labels for computing the masked language modeling loss. Indices should be in [-100, 0, ..., config.vocab_size] (see input_ids docstring) Tokens with indices set to -100 are ignored (masked), the loss is only computed for the tokens with labels TYPE: `mindspore.Tensor` of shape `(batch_size, sequence_length)`, *optional* DEFAULT: None

RETURNS	DESCRIPTION
Union[Tuple, MaskedLMOutput]	Union[Tuple, MaskedLMOutput]

This example uses a false checkpoint since we don't have any available pretrained model for the masked language modeling task with the Reformer architecture.

Example

>>> import torch
>>> from transformers import AutoTokenizer, ReformerForMaskedLM
...
>>> tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-reformer")
>>> model = ReformerForMaskedLM.from_pretrained("hf-internal-testing/tiny-random-reformer")
...
>>> # add mask_token
>>> tokenizer.add_special_tokens({"mask_token": "[MASK]"})  # doctest: +IGNORE_RESULT
>>> inputs = tokenizer("The capital of France is [MASK].", return_tensors="pt")
...
>>> # resize model's embedding matrix
>>> model.resize_token_embeddings(new_num_tokens=model.config.vocab_size +<