Skip to content

jamba

mindnlp.transformers.models.jamba.configuration_jamba

Jamba model configuration

mindnlp.transformers.models.jamba.configuration_jamba.JambaConfig

Bases: PretrainedConfig

This is the configuration class to store the configuration of a [JambaModel]. It is used to instantiate a Jamba 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 jamba-small architecture. ai21labs/jamba-small Configuration objects inherit from [PretrainedConfig] and can be used to control the model outputs. Read the documentation from [PretrainedConfig] for more information.

PARAMETER DESCRIPTION
vocab_size

Vocabulary size of the Jamba model. Defines the number of different tokens that can be represented by the inputs_ids passed when calling [JambaModel]

TYPE: `int`, *optional*, defaults to 65536 DEFAULT: 65536

tie_word_embeddings

Whether the model's input and output word embeddings should be tied. Note that this is only relevant if the model has a output word embedding layer.

TYPE: `bool`, *optional*, defaults to `False` DEFAULT: False

hidden_size

Dimension of the hidden representations.

TYPE: `int`, *optional*, defaults to 4096 DEFAULT: 4096

intermediate_size

Dimension of the MLP representations.

TYPE: `int`, *optional*, defaults to 14336 DEFAULT: 14336

num_hidden_layers

Number of hidden layers in the Transformer encoder.

TYPE: `int`, *optional*, defaults to 32 DEFAULT: 32

num_attention_heads

Number of attention heads for each attention layer in the Transformer encoder.

TYPE: `int`, *optional*, defaults to 32 DEFAULT: 32

num_key_value_heads

This is the number of key_value heads that should be used to implement Grouped Query Attention. If num_key_value_heads=num_attention_heads, the model will use Multi Head Attention (MHA), if num_key_value_heads=1, the model will use Multi Query Attention (MQA) otherwise GQA is used. When converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be forwarded by meanpooling all the original heads within that group. For more details checkout this paper. If it is not specified, will default to 8.

TYPE: `int`, *optional*, defaults to 8 DEFAULT: 8

hidden_act

The non-linear activation function (function or string) in the decoder.

TYPE: `str` or `function`, *optional*, defaults to `"silu"` DEFAULT: 'silu'

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

rms_norm_eps

The epsilon used by the rms normalization layers.

TYPE: `float`, *optional*, defaults to 1e-06 DEFAULT: 1e-06

use_cache

Whether or not the model should return the last key/values attentions (not used by all models). Only relevant if config.is_decoder=True.

TYPE: `bool`, *optional*, defaults to `True` DEFAULT: True

calc_logits_for_entire_prompt

Whether or not to calculate logits for entire prompt during generation. If False, only the logits of the last prompt token will be calculated, which are the only logits needed for generation. For long sequences, the logits for the entire sequence may use a lot of memory so setting calc_logits_for_entire_prompt=False will reduce memory footprint significantly. Note: some generation features may not be available if this is set to False.

TYPE: `bool`, *optional*, defaults to `False` DEFAULT: False

output_router_logits

Whether or not the router logits should be returned by the model. Enabling this will also allow the model to output the auxiliary loss. See here for more details

TYPE: `bool`, *optional*, defaults to `False` DEFAULT: False

router_aux_loss_coef

The aux loss factor for the total loss.

TYPE: `float`, *optional*, defaults to 0.001 DEFAULT: 0.001

pad_token_id

The id of the padding token.

TYPE: `int`, *optional*, defaults to 0 DEFAULT: 0

bos_token_id

The id of the "beginning-of-sequence" token.

TYPE: `int`, *optional*, defaults to 1 DEFAULT: 1

eos_token_id

The id of the "end-of-sequence" token.

TYPE: `int`, *optional*, defaults to 2 DEFAULT: 2

sliding_window

Sliding window attention window size. If not specified, will default to None.

TYPE: `int`, *optional* DEFAULT: None

n_ctx

This value doesn't have any real effect. The maximum sequence length that this model is intended to be used with. It can be used with longer sequences, but performance may degrade.

TYPE: `int`, *optional*, defaults to 262144 DEFAULT: 262144

attention_dropout

The dropout ratio for the attention probabilities.

TYPE: `float`, *optional*, defaults to 0.0 DEFAULT: 0.0

num_experts_per_tok

The number of experts to root per-token, can be also interpreted as the top-p routing parameter

TYPE: `int`, *optional*, defaults to 2 DEFAULT: 2

num_experts

Number of experts per Sparse MLP layer.

TYPE: `int`, *optional*, defaults to 16 DEFAULT: 16

expert_layer_period

Once in this many layers, we will have an expert layer

TYPE: `int`, *optional*, defaults to 2 DEFAULT: 2

expert_layer_offset

The first layer index that contains an expert mlp layer

TYPE: `int`, *optional*, defaults to 1 DEFAULT: 1

attn_layer_period

Once in this many layers, we will have a vanilla attention layer

TYPE: `int`, *optional*, defaults to 8 DEFAULT: 8

attn_layer_offset

The first layer index that contains a vanilla attention mlp layer

TYPE: `int`, *optional*, defaults to 4 DEFAULT: 4

use_mamba_kernels

Flag indicating whether or not to use the fast mamba kernels. These are available only if mamba-ssm and causal-conv1d are installed, and the mamba modules are running on a CUDA device. Raises ValueError if True and kernels are not available

TYPE: `bool`, *optional*, defaults to `True` DEFAULT: True

mamba_d_state

The dimension the mamba state space latents

TYPE: `int`, *optional*, defaults to 16 DEFAULT: 16

mamba_d_conv

The size of the mamba convolution kernel

TYPE: `int`, *optional*, defaults to 4 DEFAULT: 4

mamba_expand

Expanding factor (relative to hidden_size) used to determine the mamba intermediate size

TYPE: `int`, *optional*, defaults to 2 DEFAULT: 2

mamba_dt_rank

Rank of the the mamba discretization projection matrix. "auto" means that it will default to math.ceil(self.hidden_size / 16)

TYPE: `Union[int,str]`, *optional*, defaults to `"auto"` DEFAULT: 'auto'

mamba_conv_bias

Flag indicating whether or not to use bias in the convolution layer of the mamba mixer block.

TYPE: `bool`, *optional*, defaults to `True` DEFAULT: True

mamba_proj_bias

Flag indicating whether or not to use bias in the input and output projections (["in_proj", "out_proj"]) of the mamba mixer block

TYPE: `bool`, *optional*, defaults to `False` DEFAULT: False

mamba_inner_layernorms

Flag indicating whether or not to apply layernorms to internal mamba activations

TYPE: `bool`, *optional*, defaults to `True` DEFAULT: True

Source code in mindnlp/transformers/models/jamba/configuration_jamba.py 
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
class JambaConfig(PretrainedConfig):
    r"""
    This is the configuration class to store the configuration of a [`JambaModel`]. It is used to instantiate a
    Jamba 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 jamba-small architecture.
    [ai21labs/jamba-small](https://huggingface.co/ai21labs/Jamba-v0.1)
    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
    documentation from [`PretrainedConfig`] for more information.

    Args:
        vocab_size (`int`, *optional*, defaults to 65536):
            Vocabulary size of the Jamba model. Defines the number of different tokens that can be represented by the
            `inputs_ids` passed when calling [`JambaModel`]
        tie_word_embeddings (`bool`, *optional*, defaults to `False`):
            Whether the model's input and output word embeddings should be tied. Note that this is only relevant if the
            model has a output word embedding layer.
        hidden_size (`int`, *optional*, defaults to 4096):
            Dimension of the hidden representations.
        intermediate_size (`int`, *optional*, defaults to 14336):
            Dimension of the MLP representations.
        num_hidden_layers (`int`, *optional*, defaults to 32):
            Number of hidden layers in the Transformer encoder.
        num_attention_heads (`int`, *optional*, defaults to 32):
            Number of attention heads for each attention layer in the Transformer encoder.
        num_key_value_heads (`int`, *optional*, defaults to 8):
            This is the number of key_value heads that should be used to implement Grouped Query Attention. If
            `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
            `num_key_value_heads=1`, the model will use Multi Query Attention (MQA) otherwise GQA is used. When
            converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be forwarded
            by meanpooling all the original heads within that group. For more details checkout [this
            paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to `8`.
        hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
            The non-linear activation function (function or string) in the decoder.
        initializer_range (`float`, *optional*, defaults to 0.02):
            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
        rms_norm_eps (`float`, *optional*, defaults to 1e-06):
            The epsilon used by the rms normalization layers.
        use_cache (`bool`, *optional*, defaults to `True`):
            Whether or not the model should return the last key/values attentions (not used by all models). Only
            relevant if `config.is_decoder=True`.
        calc_logits_for_entire_prompt (`bool`, *optional*, defaults to `False`):
            Whether or not to calculate logits for entire prompt during generation. If `False`, only the logits of the
            last prompt token will be calculated, which are the only logits needed for generation. For long sequences,
            the logits for the entire sequence may use a lot of memory so setting `calc_logits_for_entire_prompt=False`
            will reduce memory footprint significantly.
            Note: some generation features may not be available if this is set to `False`.
        output_router_logits (`bool`, *optional*, defaults to `False`):
            Whether or not the router logits should be returned by the model. Enabling this will also
            allow the model to output the auxiliary loss. See [here]() for more details
        router_aux_loss_coef (`float`, *optional*, defaults to 0.001):
            The aux loss factor for the total loss.
        pad_token_id (`int`, *optional*, defaults to 0):
            The id of the padding token.
        bos_token_id (`int`, *optional*, defaults to 1):
            The id of the "beginning-of-sequence" token.
        eos_token_id (`int`, *optional*, defaults to 2):
            The id of the "end-of-sequence" token.
        sliding_window (`int`, *optional*):
            Sliding window attention window size. If not specified, will default to `None`.
        n_ctx (`int`, *optional*, defaults to 262144):
            This value doesn't have any real effect. The maximum sequence length that this model is intended to be
            used with. It can be used with longer sequences, but performance may degrade.
        attention_dropout (`float`, *optional*, defaults to 0.0):
            The dropout ratio for the attention probabilities.
        num_experts_per_tok (`int`, *optional*, defaults to 2):
            The number of experts to root per-token, can be also interpreted as the `top-p` routing
            parameter
        num_experts (`int`, *optional*, defaults to 16):
            Number of experts per Sparse MLP layer.
        expert_layer_period (`int`, *optional*, defaults to 2):
            Once in this many layers, we will have an expert layer
        expert_layer_offset (`int`, *optional*, defaults to 1):
            The first layer index that contains an expert mlp layer
        attn_layer_period (`int`, *optional*, defaults to 8):
            Once in this many layers, we will have a vanilla attention layer
        attn_layer_offset (`int`, *optional*, defaults to 4):
            The first layer index that contains a vanilla attention mlp layer
        use_mamba_kernels (`bool`, *optional*, defaults to `True`):
            Flag indicating whether or not to use the fast mamba kernels. These are available only if `mamba-ssm` and
            `causal-conv1d` are installed, and the mamba modules are running on a CUDA device. Raises ValueError if
            `True` and kernels are not available
        mamba_d_state (`int`, *optional*, defaults to 16):
            The dimension the mamba state space latents
        mamba_d_conv (`int`, *optional*, defaults to 4):
            The size of the mamba convolution kernel
        mamba_expand (`int`, *optional*, defaults to 2):
            Expanding factor (relative to hidden_size) used to determine the mamba intermediate size
        mamba_dt_rank (`Union[int,str]`, *optional*, defaults to `"auto"`):
            Rank of the the mamba discretization projection matrix.
            `"auto"` means that it will default to `math.ceil(self.hidden_size / 16)`
        mamba_conv_bias (`bool`, *optional*, defaults to `True`):
            Flag indicating whether or not to use bias in the convolution layer of the mamba mixer block.
        mamba_proj_bias (`bool`, *optional*, defaults to `False`):
            Flag indicating whether or not to use bias in the input and output projections
            (["in_proj", "out_proj"]) of the mamba mixer block
        mamba_inner_layernorms (`bool`, *optional*, defaults to `True`):
            Flag indicating whether or not to apply layernorms to internal mamba activations
    """
    model_type = "jamba"
    keys_to_ignore_at_inference = ["past_key_values"]

    def __init__(
            self,
            vocab_size=65536,
            tie_word_embeddings=False,
            hidden_size=4096,
            intermediate_size=14336,
            num_hidden_layers=32,
            num_attention_heads=32,
            num_key_value_heads=8,
            hidden_act="silu",
            initializer_range=0.02,
            rms_norm_eps=1e-6,
            use_cache=True,
            calc_logits_for_entire_prompt=False,
            output_router_logits=False,
            router_aux_loss_coef=0.001,
            pad_token_id=0,
            bos_token_id=1,
            eos_token_id=2,
            sliding_window=None,
            n_ctx=262144,
            attention_dropout=0.0,
            num_experts_per_tok=2,
            num_experts=16,
            expert_layer_period=2,
            expert_layer_offset=1,
            attn_layer_period=8,
            attn_layer_offset=4,
            use_mamba_kernels=True,
            mamba_d_state=16,
            mamba_d_conv=4,
            mamba_expand=2,
            mamba_dt_rank="auto",
            mamba_conv_bias=True,
            mamba_proj_bias=False,
            mamba_inner_layernorms=True,
            **kwargs,
    ):
        """
        Initializes a new instance of the JambaConfig class.

        Args:
            self: The object instance.
            vocab_size (int, optional): The size of the vocabulary. Default is 65536.
            tie_word_embeddings (bool, optional): Whether to tie the word embeddings. Default is False.
            hidden_size (int, optional): The size of the hidden layers. Default is 4096.
            intermediate_size (int, optional): The size of the intermediate layers. Default is 14336.
            num_hidden_layers (int, optional): The number of hidden layers. Default is 32.
            num_attention_heads (int, optional): The number of attention heads. Default is 32.
            num_key_value_heads (int, optional): The number of key-value heads. Default is 8.
            hidden_act (str, optional): The activation function for the hidden layers. Default is 'silu'.
            initializer_range (float, optional): The range for weight initialization. Default is 0.02.
            rms_norm_eps (float, optional): The epsilon value for RMS normalization. Default is 1e-06.
            use_cache (bool, optional): Whether to use cache for attention layers. Default is True.
            calc_logits_for_entire_prompt (bool, optional): Whether to calculate logits for the entire prompt.
                Default is False.
            output_router_logits (bool, optional): Whether to output router logits. Default is False.
            router_aux_loss_coef (float, optional): The coefficient for the router auxiliary loss. Default is 0.001.
            pad_token_id (int, optional): The token ID for padding. Default is 0.
            bos_token_id (int, optional): The token ID for the beginning of sentence. Default is 1.
            eos_token_id (int, optional): The token ID for the end of sentence. Default is 2.
            sliding_window (None or int, optional): The size of the sliding window. Default is None.
            n_ctx (int, optional): The size of the context window. Default is 262144.
            attention_dropout (float, optional): The dropout rate for attention layers. Default is 0.0.
            num_experts_per_tok (int, optional): The number of experts per token. Default is 2.
            num_experts (int, optional): The total number of experts. Default is 16.
            expert_layer_period (int, optional): The period for expert layers. Default is 2.
            expert_layer_offset (int, optional): The offset for expert layers. Default is 1.
            attn_layer_period (int, optional): The period for attention layers. Default is 8.
            attn_layer_offset (int, optional): The offset for attention layers. Default is 4.
            use_mamba_kernels (bool, optional): Whether to use Mamba kernels. Default is True.
            mamba_d_state (int, optional): The state dimension for Mamba. Default is 16.
            mamba_d_conv (int, optional): The convolutional dimension for Mamba. Default is 4.
            mamba_expand (int, optional): The expansion factor for Mamba. Default is 2.
            mamba_dt_rank (int or 'auto', optional): The rank for Mamba's data tensors. Default is 'auto'.
            mamba_conv_bias (bool, optional): Whether to include biases in Mamba's convolution layers. Default is True.
            mamba_proj_bias (bool, optional): Whether to include biases in Mamba's projection layers. Default is False.
            mamba_inner_layernorms (bool, optional): Whether to use inner layer normalization in Mamba. Default is True.

        Returns:
            None

        Raises:
            None
        """
        self.vocab_size = vocab_size
        self.tie_word_embeddings = tie_word_embeddings
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.sliding_window = sliding_window
        self.n_ctx = n_ctx
        self.attention_dropout = attention_dropout

        # for backward compatibility
        if num_key_value_heads is None:
            num_key_value_heads = num_attention_heads

        self.num_key_value_heads = num_key_value_heads
        self.hidden_act = hidden_act
        self.initializer_range = initializer_range
        self.rms_norm_eps = rms_norm_eps

        self.use_cache = use_cache
        self.calc_logits_for_entire_prompt = calc_logits_for_entire_prompt
        self.output_router_logits = output_router_logits
        self.router_aux_loss_coef = router_aux_loss_coef

        self.num_experts_per_tok = num_experts_per_tok
        self.num_experts = num_experts
        self.expert_layer_period = expert_layer_period
        self.expert_layer_offset = expert_layer_offset
        self.attn_layer_period = attn_layer_period
        self.attn_layer_offset = attn_layer_offset

        self.use_mamba_kernels = use_mamba_kernels
        self.mamba_d_state = mamba_d_state
        self.mamba_d_conv = mamba_d_conv
        self.mamba_expand = mamba_expand
        self.mamba_dt_rank = math.ceil(self.hidden_size / 16) if mamba_dt_rank == "auto" else mamba_dt_rank
        self.mamba_conv_bias = mamba_conv_bias
        self.mamba_proj_bias = mamba_proj_bias
        self.mamba_inner_layernorms = mamba_inner_layernorms

        super().__init__(
            pad_token_id=pad_token_id,
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            tie_word_embeddings=tie_word_embeddings,
            **kwargs,
        )

mindnlp.transformers.models.jamba.configuration_jamba.JambaConfig.__init__(vocab_size=65536, tie_word_embeddings=False, hidden_size=4096, intermediate_size=14336, num_hidden_layers=32, num_attention_heads=32, num_key_value_heads=8, hidden_act='silu', initializer_range=0.02, rms_norm_eps=1e-06, use_cache=True, calc_logits_for_entire_prompt=False, output_router_logits=False, router_aux_loss_coef=0.001, pad_token_id=0, bos_token_id=1, eos_token_id=2, sliding_window=None, n_ctx=262144, attention_dropout=0.0, num_experts_per_tok=2, num_experts=16, expert_layer_period=2, expert_layer_offset=1, attn_layer_period=8, attn_layer_offset=4, use_mamba_kernels=True, mamba_d_state=16, mamba_d_conv=4, mamba_expand=2, mamba_dt_rank='auto', mamba_conv_bias=True, mamba_proj_bias=False, mamba_inner_layernorms=True, **kwargs)

Initializes a new instance of the JambaConfig class.

PARAMETER DESCRIPTION
self

The object instance.
vocab_size

The size of the vocabulary. Default is 65536.

TYPE: int DEFAULT: 65536

tie_word_embeddings

Whether to tie the word embeddings. Default is False.

TYPE: bool DEFAULT: False

hidden_size

The size of the hidden layers. Default is 4096.

TYPE: int DEFAULT: 4096

intermediate_size

The size of the intermediate layers. Default is 14336.

TYPE: int DEFAULT: 14336

num_hidden_layers

The number of hidden layers. Default is 32.

TYPE: int DEFAULT: 32

num_attention_heads

The number of attention heads. Default is 32.

TYPE: int DEFAULT: 32

num_key_value_heads

The number of key-value heads. Default is 8.

TYPE: int DEFAULT: 8

hidden_act

The activation function for the hidden layers. Default is 'silu'.

TYPE: str DEFAULT: 'silu'

initializer_range

The range for weight initialization. Default is 0.02.

TYPE: float DEFAULT: 0.02

rms_norm_eps

The epsilon value for RMS normalization. Default is 1e-06.

TYPE: float DEFAULT: 1e-06

use_cache

Whether to use cache for attention layers. Default is True.

TYPE: bool DEFAULT: True

calc_logits_for_entire_prompt

Whether to calculate logits for the entire prompt. Default is False.

TYPE: bool DEFAULT: False

output_router_logits

Whether to output router logits. Default is False.

TYPE: bool DEFAULT: False

router_aux_loss_coef

The coefficient for the router auxiliary loss. Default is 0.001.

TYPE: float DEFAULT: 0.001

pad_token_id

The token ID for padding. Default is 0.

TYPE: int DEFAULT: 0

bos_token_id

The token ID for the beginning of sentence. Default is 1.

TYPE: int DEFAULT: 1

eos_token_id

The token ID for the end of sentence. Default is 2.

TYPE: int DEFAULT: 2

sliding_window

The size of the sliding window. Default is None.

TYPE: None or int DEFAULT: None

n_ctx

The size of the context window. Default is 262144.

TYPE: int DEFAULT: 262144

attention_dropout

The dropout rate for attention layers. Default is 0.0.

TYPE: float DEFAULT: 0.0

num_experts_per_tok

The number of experts per token. Default is 2.

TYPE: int DEFAULT: 2

num_experts

The total number of experts. Default is 16.

TYPE: int DEFAULT: 16

expert_layer_period

The period for expert layers. Default is 2.

TYPE: int DEFAULT: 2

expert_layer_offset

The offset for expert layers. Default is 1.

TYPE: int DEFAULT: 1

attn_layer_period

The period for attention layers. Default is 8.

TYPE: int DEFAULT: 8

attn_layer_offset

The offset for attention layers. Default is 4.

TYPE: int DEFAULT: 4

use_mamba_kernels

Whether to use Mamba kernels. Default is True.

TYPE: bool DEFAULT: True

mamba_d_state

The state dimension for Mamba. Default is 16.

TYPE: int DEFAULT: 16

mamba_d_conv

The convolutional dimension for Mamba. Default is 4.

TYPE: int DEFAULT: 4

mamba_expand

The expansion factor for Mamba. Default is 2.

TYPE: int DEFAULT: 2

mamba_dt_rank

The rank for Mamba's data tensors. Default is 'auto'.

TYPE: int or auto DEFAULT: 'auto'

mamba_conv_bias

Whether to include biases in Mamba's convolution layers. Default is True.

TYPE: bool DEFAULT: True

mamba_proj_bias

Whether to include biases in Mamba's projection layers. Default is False.

TYPE: bool DEFAULT: False

mamba_inner_layernorms

Whether to use inner layer normalization in Mamba. Default is True.

TYPE: bool DEFAULT: True

RETURNS DESCRIPTION

None
Source code in mindnlp/transformers/models/jamba/configuration_jamba.py 
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
def __init__(
        self,
        vocab_size=65536,
        tie_word_embeddings=False,
        hidden_size=4096,
        intermediate_size=14336,
        num_hidden_layers=32,
        num_attention_heads=32,
        num_key_value_heads=8,
        hidden_act="silu",
        initializer_range=0.02,
        rms_norm_eps=1e-6,
        use_cache=True,
        calc_logits_for_entire_prompt=False,
        output_router_logits=False,
        router_aux_loss_coef=0.001,
        pad_token_id=0,
        bos_token_id=1,
        eos_token_id=2,
        sliding_window=None,
        n_ctx=262144,
        attention_dropout=0.0,
        num_experts_per_tok=2,
        num_experts=16,
        expert_layer_period=2,
        expert_layer_offset=1,
        attn_layer_period=8,
        attn_layer_offset=4,
        use_mamba_kernels=True,
        mamba_d_state=16,
        mamba_d_conv=4,
        mamba_expand=2,
        mamba_dt_rank="auto",
        mamba_conv_bias=True,
        mamba_proj_bias=False,
        mamba_inner_layernorms=True,
        **kwargs,
):
    """
    Initializes a new instance of the JambaConfig class.

    Args:
        self: The object instance.
        vocab_size (int, optional): The size of the vocabulary. Default is 65536.
        tie_word_embeddings (bool, optional): Whether to tie the word embeddings. Default is False.
        hidden_size (int, optional): The size of the hidden layers. Default is 4096.
        intermediate_size (int, optional): The size of the intermediate layers. Default is 14336.
        num_hidden_layers (int, optional): The number of hidden layers. Default is 32.
        num_attention_heads (int, optional): The number of attention heads. Default is 32.
        num_key_value_heads (int, optional): The number of key-value heads. Default is 8.
        hidden_act (str, optional): The activation function for the hidden layers. Default is 'silu'.
        initializer_range (float, optional): The range for weight initialization. Default is 0.02.
        rms_norm_eps (float, optional): The epsilon value for RMS normalization. Default is 1e-06.
        use_cache (bool, optional): Whether to use cache for attention layers. Default is True.
        calc_logits_for_entire_prompt (bool, optional): Whether to calculate logits for the entire prompt.
            Default is False.
        output_router_logits (bool, optional): Whether to output router logits. Default is False.
        router_aux_loss_coef (float, optional): The coefficient for the router auxiliary loss. Default is 0.001.
        pad_token_id (int, optional): The token ID for padding. Default is 0.
        bos_token_id (int, optional): The token ID for the beginning of sentence. Default is 1.
        eos_token_id (int, optional): The token ID for the end of sentence. Default is 2.
        sliding_window (None or int, optional): The size of the sliding window. Default is None.
        n_ctx (int, optional): The size of the context window. Default is 262144.
        attention_dropout (float, optional): The dropout rate for attention layers. Default is 0.0.
        num_experts_per_tok (int, optional): The number of experts per token. Default is 2.
        num_experts (int, optional): The total number of experts. Default is 16.
        expert_layer_period (int, optional): The period for expert layers. Default is 2.
        expert_layer_offset (int, optional): The offset for expert layers. Default is 1.
        attn_layer_period (int, optional): The period for attention layers. Default is 8.
        attn_layer_offset (int, optional): The offset for attention layers. Default is 4.
        use_mamba_kernels (bool, optional): Whether to use Mamba kernels. Default is True.
        mamba_d_state (int, optional): The state dimension for Mamba. Default is 16.
        mamba_d_conv (int, optional): The convolutional dimension for Mamba. Default is 4.
        mamba_expand (int, optional): The expansion factor for Mamba. Default is 2.
        mamba_dt_rank (int or 'auto', optional): The rank for Mamba's data tensors. Default is 'auto'.
        mamba_conv_bias (bool, optional): Whether to include biases in Mamba's convolution layers. Default is True.
        mamba_proj_bias (bool, optional): Whether to include biases in Mamba's projection layers. Default is False.
        mamba_inner_layernorms (bool, optional): Whether to use inner layer normalization in Mamba. Default is True.

    Returns:
        None

    Raises:
        None
    """
    self.vocab_size = vocab_size
    self.tie_word_embeddings = tie_word_embeddings
    self.hidden_size = hidden_size
    self.intermediate_size = intermediate_size
    self.num_hidden_layers = num_hidden_layers
    self.num_attention_heads = num_attention_heads
    self.sliding_window = sliding_window
    self.n_ctx = n_ctx
    self.attention_dropout = attention_dropout

    # for backward compatibility
    if num_key_value_heads is None:
        num_key_value_heads = num_attention_heads

    self.num_key_value_heads = num_key_value_heads
    self.hidden_act = hidden_act
    self.initializer_range = initializer_range
    self.rms_norm_eps = rms_norm_eps

    self.use_cache = use_cache
    self.calc_logits_for_entire_prompt = calc_logits_for_entire_prompt
    self.output_router_logits = output_router_logits
    self.router_aux_loss_coef = router_aux_loss_coef

    self.num_experts_per_tok = num_experts_per_tok
    self.num_experts = num_experts
    self.expert_layer_period = expert_layer_period
    self.expert_layer_offset = expert_layer_offset
    self.attn_layer_period = attn_layer_period
    self.attn_layer_offset = attn_layer_offset

    self.use_mamba_kernels = use_mamba_kernels
    self.mamba_d_state = mamba_d_state
    self.mamba_d_conv = mamba_d_conv
    self.mamba_expand = mamba_expand
    self.mamba_dt_rank = math.ceil(self.hidden_size / 16) if mamba_dt_rank == "auto" else mamba_dt_rank
    self.mamba_conv_bias = mamba_conv_bias
    self.mamba_proj_bias = mamba_proj_bias
    self.mamba_inner_layernorms = mamba_inner_layernorms

    super().__init__(
        pad_token_id=pad_token_id,
        bos_token_id=bos_token_id,
        eos_token_id=eos_token_id,
        tie_word_embeddings=tie_word_embeddings,
        **kwargs,
    )

mindnlp.transformers.models.jamba.modeling_jamba.JambaModel

Bases: JambaPreTrainedModel

Transformer decoder consisting of config.num_hidden_layers layers. Each layer is a [JambaDecoderLayer]

PARAMETER DESCRIPTION
config

JambaConfig

TYPE: JambaConfig

Source code in mindnlp/transformers/models/jamba/modeling_jamba.py 
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
1472
1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
1583
1584
1585
1586
1587
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
1602
1603
1604
1605
1606
1607
1608
1609
1610
1611
1612
1613
1614
1615
1616
1617
1618
1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
1634
1635
1636
1637
1638
1639
1640
1641
1642
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655
1656
1657
1658
1659
1660
1661
1662
1663
1664
1665
1666
1667
1668
1669
1670
1671
1672
1673
1674
1675
1676
1677
1678
1679
1680
1681
1682
1683
1684
1685
1686
1687
1688
1689
1690
1691
1692
1693
1694
1695
1696
1697
1698
1699
1700
1701
1702
1703
1704
1705
class JambaModel(JambaPreTrainedModel):
    """
    Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`JambaDecoderLayer`]

    Args:
        config: JambaConfig
    """
    def __init__(self, config: JambaConfig):

        """
        Initializes a JambaModel instance with the provided configuration.

        Args:
            self (JambaModel): The instance of the JambaModel class.
            config (JambaConfig): An instance of JambaConfig containing configuration parameters for the model.
                The configuration should include:

                - pad_token_id (int): The index of the padding token.
                - vocab_size (int): The size of the vocabulary.
                - hidden_size (int): The size of the hidden layer.
                - num_hidden_layers (int): The total number of hidden layers in the model.
                - attn_layer_offset (int): The offset for the attention layer.
                - attn_layer_period (int): The period for the attention layer.
                - expert_layer_offset (int): The offset for the expert layer.
                - expert_layer_period (int): The period for the expert layer.
                - num_experts (int): The number of experts in the model.

        Returns:
            None.

        Raises:
            ValueError:
                - If at least one layer in the decoder is not an attention layer.
                - If at least one layer in the decoder is not a Mamba layer.
                - If the Mamba state size is equal to the convolution size in the Mamba layer.
        """
        super().__init__(config)
        self.padding_idx = config.pad_token_id
        self.vocab_size = config.vocab_size

        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)

        # init each model layer, decide if it's mamba/attention and has experts or not
        decoder_layers = []
        for i in range(config.num_hidden_layers):
            is_attn = (i - self.config.attn_layer_offset) % self.config.attn_layer_period == 0
            is_expert = (i - self.config.expert_layer_offset) % self.config.expert_layer_period == 0

            num_experts = self.config.num_experts if is_expert else 1
            if is_attn:
                decoder_layers.append(JambaAttentionDecoderLayer(config, num_experts=num_experts, layer_idx=i))
            else:
                decoder_layers.append(JambaMambaDecoderLayer(config, num_experts=num_experts, layer_idx=i))

        if not any(isinstance(layer, JambaAttentionDecoderLayer) for layer in decoder_layers):
            raise ValueError("At least one layer in the decoder must be an attention layer")
        self._attn_layer_index = [isinstance(layer, JambaAttentionDecoderLayer) for layer in decoder_layers].index(
            True
        )

        if not any(isinstance(layer, JambaMambaDecoderLayer) for layer in decoder_layers):
            raise ValueError("At least one layer in the decoder must be a Mamba layer")
        self._mamba_layer_index = [isinstance(layer, JambaMambaDecoderLayer) for layer in decoder_layers].index(True)

        if (
                decoder_layers[self._mamba_layer_index].mamba.ssm_state_size
                == decoder_layers[self._mamba_layer_index].mamba.conv_kernel_size
        ):
            raise ValueError("Mamba state size and convolution size must be different")

        self.layers = nn.ModuleList(decoder_layers)

        self.final_layernorm = JambaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)

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

    def get_input_embeddings(self):

        """
        Method to retrieve the input embeddings from the JambaModel instance.

        Args:
            self (JambaModel): The current instance of the JambaModel class.
                This parameter is required to access the embed_tokens attribute.

        Returns:
            None: This method returns the embed_tokens attribute of the JambaModel instance.
                The embed_tokens attribute represents the input embeddings used by the model.

        Raises:
            This method does not raise any exceptions.
        """
        return self.embed_tokens

    def set_input_embeddings(self, value):

        """
        Set the input embeddings for the JambaModel.

        Args:
            self (JambaModel): The instance of the JambaModel class.
            value: The input embeddings to be set for the model.
               This can be a tensor, array, or any compatible data structure.
               It represents the embeddings to be assigned to the model for input processing.

        Returns:
            None.

        Raises:
            None.
        """
        self.embed_tokens = value

    # Ignore copy
    def forward(
            self,
            input_ids: mindspore.Tensor = None,
            attention_mask: Optional[mindspore.Tensor] = None,
            position_ids: Optional[mindspore.Tensor] = None,
            past_key_values: Optional[Union[List[mindspore.Tensor], HybridMambaAttentionDynamicCache]] = None,
            inputs_embeds: Optional[mindspore.Tensor] = None,
            use_cache: Optional[bool] = None,
            output_attentions: Optional[bool] = None,
            output_hidden_states: Optional[bool] = None,
            output_router_logits: Optional[bool] = None,
            return_dict: Optional[bool] = None,
    ) -> Union[Tuple, MoeModelOutputWithPast]:

        """
        This method 'forward' in the class 'JambaModel' forwards the model by processing input data through
        the layers of the model.

        Args:
            self: The instance of the class.
            input_ids (mindspore.Tensor): The input tensor containing token ids for the input sequence.
            attention_mask (Optional[mindspore.Tensor]): An optional tensor specifying the attention mask to be applied.
            position_ids (Optional[mindspore.Tensor]): An optional tensor containing positional ids for the input sequence.
            past_key_values (Optional[Union[List[mindspore.Tensor], HybridMambaAttentionDynamicCache]]):
                Optional past key values for attention mechanism.
            inputs_embeds (Optional[mindspore.Tensor]): Optional tensor containing input embeddings.
            use_cache (Optional[bool]): Optional boolean specifying whether to use cache for the model.
            output_attentions (Optional[bool]): Optional boolean specifying whether to output attentions.
            output_hidden_states (Optional[bool]): Optional boolean specifying whether to output hidden states.
            output_router_logits (Optional[bool]): Optional boolean specifying whether to output router logits.
            return_dict (Optional[bool]): Optional boolean specifying whether to return the output as a dictionary.

        Returns:
            Union[Tuple, MoeModelOutputWithPast]: The return value can either be a tuple containing relevant outputs
                or an instance of MoeModelOutputWithPast class.

        Raises:
            ValueError: Raised if both input_ids and inputs_embeds are specified at the same time,
                or if neither input_ids nor inputs_embeds are specified.
            Warning: Raised as a warning if 'use_cache=True' is incompatible with gradient checkpointing,
                and it automatically sets 'use_cache=False'.
        """
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_router_logits = (
            output_router_logits if output_router_logits is not None else self.config.output_router_logits
        )
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        use_cache = use_cache if use_cache is not None else self.config.use_cache

        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

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

        past_key_values_length = 0

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

        if use_cache:
            if isinstance(past_key_values, Cache) and not isinstance(
                    past_key_values, HybridMambaAttentionDynamicCache
            ):
                past_key_values = HybridMambaAttentionDynamicCache.from_legacy_cache(past_key_values.to_legacy_cache())
            use_legacy_cache = not isinstance(past_key_values, HybridMambaAttentionDynamicCache)
            if use_legacy_cache:
                past_key_values = HybridMambaAttentionDynamicCache.from_legacy_cache(past_key_values)
            past_key_values_length = past_key_values.get_usable_length(seq_length, self._attn_layer_index)

        if position_ids is None:
            position_ids = ops.arange(
                past_key_values_length, seq_length + past_key_values_length, dtype=mindspore.int64
            )
            position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
        else:
            position_ids = position_ids.view(-1, seq_length).long()

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

        # 4d mask is passed through the layers
        attention_mask = _prepare_4d_causal_attention_mask(
            attention_mask,
            (batch_size, seq_length),
            inputs_embeds,
            past_key_values_length,
            sliding_window=self.config.sliding_window,
        )

        hidden_states = inputs_embeds

        # decoder layers
        all_hidden_states = () if output_hidden_states else None
        all_self_attns = () if output_attentions else None
        all_router_logits = () if output_router_logits else None
        next_decoder_cache = None

        for decoder_layer in self.layers:
            if output_hidden_states:
                all_hidden_states += (hidden_states,)

            if self.gradient_checkpointing and self.training:
                layer_outputs = self._gradient_checkpointing_func(
                    decoder_layer.__call__,
                    hidden_states,
                    attention_mask,
                    position_ids,
                    past_key_values,
                    output_attentions,
                    output_router_logits,
                    use_cache,
                )
            else:
                layer_outputs = decoder_layer(
                    hidden_states,
                    attention_mask=attention_mask,
                    position_ids=position_ids,
                    past_key_value=past_key_values,
                    output_attentions=output_attentions,
                    output_router_logits=output_router_logits,
                    use_cache=use_cache,
                )

            hidden_states = layer_outputs[0]

            if use_cache:
                next_decoder_cache = layer_outputs[2 if output_attentions else 1]

            if output_attentions:
                all_self_attns += (layer_outputs[1],)

            if output_router_logits:
                all_router_logits += (layer_outputs[-1],)

        hidden_states = self.final_layernorm(hidden_states)

        # add hidden states from the last decoder layer
        if output_hidden_states:
            all_hidden_states += (hidden_states,)

        next_cache = None
        if use_cache:
            next_cache = next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache

        if not return_dict:
            return tuple(
                v
                for v in [hidden_states, next_cache, all_hidden_states, all_self_attns, all_router_logits]
                if v is not None
            )
        return MoeModelOutputWithPast(
            last_hidden_state=hidden_states,
            past_key_values=next_cache,
            hidden_states=all_hidden_states,
            attentions=all_self_attns,
            router_logits=all_router_logits,
        )

mindnlp.transformers.models.jamba.modeling_jamba.JambaModel.__init__(config)

Initializes a JambaModel instance with the provided configuration.

PARAMETER DESCRIPTION
self

The instance of the JambaModel class.

TYPE: JambaModel

config

An instance of JambaConfig containing configuration parameters for the model. The configuration should include:

  • pad_token_id (int): The index of the padding token.
  • vocab_size (int): The size of the vocabulary.
  • hidden_size (int): The size of the hidden layer.
  • num_hidden_layers (int): The total number of hidden layers in the model.
  • attn_layer_offset (int): The offset for the attention layer.
  • attn_layer_period (int): The period for the attention layer.
  • expert_layer_offset (int): The offset for the expert layer.
  • expert_layer_period (int): The period for the expert layer.
  • num_experts (int): The number of experts in the model.

TYPE: JambaConfig

RETURNS DESCRIPTION

None.
RAISES DESCRIPTION
ValueError
  • If at least one layer in the decoder is not an attention layer.
  • If at least one layer in the decoder is not a Mamba layer.
  • If the Mamba state size is equal to the convolution size in the Mamba layer.
Source code in mindnlp/transformers/models/jamba/modeling_jamba.py 
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
1472
1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
def __init__(self, config: JambaConfig):

    """
    Initializes a JambaModel instance with the provided configuration.

    Args:
        self (JambaModel): The instance of the JambaModel class.
        config (JambaConfig): An instance of JambaConfig containing configuration parameters for the model.
            The configuration should include:

            - pad_token_id (int): The index of the padding token.
            - vocab_size (int): The size of the vocabulary.
            - hidden_size (int): The size of the hidden layer.
            - num_hidden_layers (int): The total number of hidden layers in the model.
            - attn_layer_offset (int): The offset for the attention layer.
            - attn_layer_period (int): The period for the attention layer.
            - expert_layer_offset (int): The offset for the expert layer.
            - expert_layer_period (int): The period for the expert layer.
            - num_experts (int): The number of experts in the model.

    Returns:
        None.

    Raises:
        ValueError:
            - If at least one layer in the decoder is not an attention layer.
            - If at least one layer in the decoder is not a Mamba layer.
            - If the Mamba state size is equal to the convolution size in the Mamba layer.
    """
    super().__init__(config)
    self.padding_idx = config.pad_token_id
    self.vocab_size = config.vocab_size

    self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)

    # init each model layer, decide if it's mamba/attention and has experts or not
    decoder_layers = []
    for i in range(config.num_hidden_layers):
        is_attn = (i - self.config.attn_layer_offset) % self.config.attn_layer_period == 0
        is_expert = (i - self.config.expert_layer_offset) % self.config.expert_layer_period == 0

        num_experts = self.config.num_experts if is_expert else 1
        if is_attn:
            decoder_layers.append(JambaAttentionDecoderLayer(config, num_experts=num_experts, layer_idx=i))
        else:
            decoder_layers.append(JambaMambaDecoderLayer(config, num_experts=num_experts, layer_idx=i))

    if not any(isinstance(layer, JambaAttentionDecoderLayer) for layer in decoder_layers):
        raise ValueError("At least one layer in the decoder must be an attention layer")
    self._attn_layer_index = [isinstance(layer, JambaAttentionDecoderLayer) for layer in decoder_layers].index(
        True
    )

    if not any(isinstance(layer, JambaMambaDecoderLayer) for layer in decoder_layers):
        raise ValueError("At least one layer in the decoder must be a Mamba layer")
    self._mamba_layer_index = [isinstance(layer, JambaMambaDecoderLayer) for layer in decoder_layers].index(True)

    if (
            decoder_layers[self._mamba_layer_index].mamba.ssm_state_size
            == decoder_layers[self._mamba_layer_index].mamba.conv_kernel_size
    ):
        raise ValueError("Mamba state size and convolution size must be different")

    self.layers = nn.ModuleList(decoder_layers)

    self.final_layernorm = JambaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)

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

mindnlp.transformers.models.jamba.modeling_jamba.JambaModel.forward(input_ids=None, attention_mask=None, position_ids=None, past_key_values=None, inputs_embeds=None, use_cache=None, output_attentions=None, output_hidden_states=None, output_router_logits=None, return_dict=None)

This method 'forward' in the class 'JambaModel' forwards the model by processing input data through the layers of the model.

PARAMETER DESCRIPTION
self

The instance of the class.
input_ids

The input tensor containing token ids for the input sequence.

TYPE: Tensor DEFAULT: None

attention_mask

An optional tensor specifying the attention mask to be applied.

TYPE: Optional[Tensor] DEFAULT: None

position_ids

An optional tensor containing positional ids for the input sequence.

TYPE: Optional[Tensor] DEFAULT: None

past_key_values

Optional past key values for attention mechanism.

TYPE: Optional[Union[List[Tensor], HybridMambaAttentionDynamicCache]] DEFAULT: None

inputs_embeds

Optional tensor containing input embeddings.

TYPE: Optional[Tensor] DEFAULT: None

use_cache

Optional boolean specifying whether to use cache for the model.

TYPE: Optional[bool] DEFAULT: None

output_attentions

Optional boolean specifying whether to output attentions.

TYPE: Optional[bool] DEFAULT: None

output_hidden_states

Optional boolean specifying whether to output hidden states.

TYPE: Optional[bool] DEFAULT: None

output_router_logits

Optional boolean specifying whether to output router logits.

TYPE: Optional[bool] DEFAULT: None

return_dict

Optional boolean specifying whether to return the output as a dictionary.

TYPE: Optional[bool] DEFAULT: None

RETURNS DESCRIPTION
Union[Tuple, MoeModelOutputWithPast]

Union[Tuple, MoeModelOutputWithPast]: The return value can either be a tuple containing relevant outputs or an instance of MoeModelOutputWithPast class.
RAISES DESCRIPTION
ValueError

Raised if both input_ids and inputs_embeds are specified at the same time, or if neither input_ids nor inputs_embeds are specified.
Warning

Raised as a warning if 'use_cache=True' is incompatible with gradient checkpointing, and it automatically sets 'use_cache=False'.
Source code in mindnlp/transformers/models/jamba/modeling_jamba.py 
1536
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
1583
1584
1585
1586
1587
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
1602
1603
1604
1605
1606
1607
1608
1609
1610
1611
1612
1613
1614
1615
1616
1617
1618
1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
1634
1635
1636
1637
1638
1639
1640
1641
1642
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655
1656
1657
1658
1659
1660
1661
1662
1663
1664
1665
1666
1667
1668
1669
1670
1671
1672
1673
1674
1675
1676
1677
1678
1679
1680
1681
1682
1683
1684
1685
1686
1687
1688
1689
1690
1691
1692
1693
1694
1695
1696
1697
1698
1699
1700
1701
1702
1703
1704
1705
def forward(
        self,
        input_ids: mindspore.Tensor = None,
        attention_mask: Optional[mindspore.Tensor] = None,
        position_ids: Optional[mindspore.Tensor] = None,
        past_key_values: Optional[Union[List[mindspore.Tensor], HybridMambaAttentionDynamicCache]] = None,
        inputs_embeds: Optional[mindspore.Tensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        output_router_logits: Optional[bool] = None,
        return_dict: Optional[bool] = None,
) -> Union[Tuple, MoeModelOutputWithPast]:

    """
    This method 'forward' in the class 'JambaModel' forwards the model by processing input data through
    the layers of the model.

    Args:
        self: The instance of the class.
        input_ids (mindspore.Tensor): The input tensor containing token ids for the input sequence.
        attention_mask (Optional[mindspore.Tensor]): An optional tensor specifying the attention mask to be applied.
        position_ids (Optional[mindspore.Tensor]): An optional tensor containing positional ids for the input sequence.
        past_key_values (Optional[Union[List[mindspore.Tensor], HybridMambaAttentionDynamicCache]]):
            Optional past key values for attention mechanism.
        inputs_embeds (Optional[mindspore.Tensor]): Optional tensor containing input embeddings.
        use_cache (Optional[bool]): Optional boolean specifying whether to use cache for the model.
        output_attentions (Optional[bool]): Optional boolean specifying whether to output attentions.
        output_hidden_states (Optional[bool]): Optional boolean specifying whether to output hidden states.
        output_router_logits (Optional[bool]): Optional boolean specifying whether to output router logits.
        return_dict (Optional[bool]): Optional boolean specifying whether to return the output as a dictionary.

    Returns:
        Union[Tuple, MoeModelOutputWithPast]: The return value can either be a tuple containing relevant outputs
            or an instance of MoeModelOutputWithPast class.

    Raises:
        ValueError: Raised if both input_ids and inputs_embeds are specified at the same time,
            or if neither input_ids nor inputs_embeds are specified.
        Warning: Raised as a warning if 'use_cache=True' is incompatible with gradient checkpointing,
            and it automatically sets 'use_cache=False'.
    """
    output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
    output_router_logits = (
        output_router_logits if output_router_logits is not None else self.config.output_router_logits
    )
    output_hidden_states = (
        output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
    )
    use_cache = use_cache if use_cache is not None else self.config.use_cache

    return_dict = return_dict if return_dict is not None else self.config.use_return_dict

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

    past_key_values_length = 0

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

    if use_cache:
        if isinstance(past_key_values, Cache) and not isinstance(
                past_key_values, HybridMambaAttentionDynamicCache
        ):
            past_key_values = HybridMambaAttentionDynamicCache.from_legacy_cache(past_key_values.to_legacy_cache())
        use_legacy_cache = not isinstance(past_key_values, HybridMambaAttentionDynamicCache)
        if use_legacy_cache:
            past_key_values = HybridMambaAttentionDynamicCache.from_legacy_cache(past_key_values)
        past_key_values_length = past_key_values.get_usable_length(seq_length, self._attn_layer_index)

    if position_ids is None:
        position_ids = ops.arange(
            past_key_values_length, seq_length + past_key_values_length, dtype=mindspore.int64
        )
        position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
    else:
        position_ids = position_ids.view(-1, seq_length).long()

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

    # 4d mask is passed through the layers
    attention_mask = _prepare_4d_causal_attention_mask(
        attention_mask,
        (batch_size, seq_length),
        inputs_embeds,
        past_key_values_length,
        sliding_window=self.config.sliding_window,
    )

    hidden_states = inputs_embeds

    # decoder layers
    all_hidden_states = () if output_hidden_states else None
    all_self_attns = () if output_attentions else None
    all_router_logits = () if output_router_logits else None
    next_decoder_cache = None

    for decoder_layer in self.layers:
        if output_hidden_states:
            all_hidden_states += (hidden_states,)

        if self.gradient_checkpointing and self.training:
            layer_outputs = self._gradient_checkpointing_func(
                decoder_layer.__call__,
                hidden_states,
                attention_mask,
                position_ids,
                past_key_values,
                output_attentions,
                output_router_logits,
                use_cache,
            )
        else:
            layer_outputs = decoder_layer(
                hidden_states,
                attention_mask=attention_mask,
                position_ids=position_ids,
                past_key_value=past_key_values,
                output_attentions=output_attentions,
                output_router_logits=output_router_logits,
                use_cache=use_cache,
            )

        hidden_states = layer_outputs[0]

        if use_cache:
            next_decoder_cache = layer_outputs[2 if output_attentions else 1]

        if output_attentions:
            all_self_attns += (layer_outputs[1],)

        if output_router_logits:
            all_router_logits += (layer_outputs[-1],)

    hidden_states = self.final_layernorm(hidden_states)

    # add hidden states from the last decoder layer
    if output_hidden_states:
        all_hidden_states += (hidden_states,)

    next_cache = None
    if use_cache:
        next_cache = next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache

    if not return_dict:
        return tuple(
            v
            for v in [hidden_states, next_cache, all_hidden_states, all_self_attns, all_router_logits]
            if v is not None
        )
    return MoeModelOutputWithPast(
        last_hidden_state=hidden_states,
        past_key_values=next_cache,
        hidden_states=all_hidden_states,
        attentions=all_self_attns,
        router_logits=all_router_logits,
    )

mindnlp.transformers.models.jamba.modeling_jamba.JambaModel.get_input_embeddings()

Method to retrieve the input embeddings from the JambaModel instance.

PARAMETER DESCRIPTION
self

The current instance of the JambaModel class. This parameter is required to access the embed_tokens attribute.

TYPE: JambaModel

RETURNS DESCRIPTION
None

This method returns the embed_tokens attribute of the JambaModel instance. The embed_tokens attribute represents the input embeddings used by the model.
Source code in mindnlp/transformers/models/jamba/modeling_jamba.py 
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
def get_input_embeddings(self):

    """
    Method to retrieve the input embeddings from the JambaModel instance.

    Args:
        self (JambaModel): The current instance of the JambaModel class.
            This parameter is required to access the embed_tokens attribute.

    Returns:
        None: This method returns the embed_tokens attribute of the JambaModel instance.
            The embed_tokens attribute represents the input embeddings used by the model.

    Raises:
        This method does not raise any exceptions.
    """
    return self.embed_tokens

mindnlp.transformers.models.jamba.modeling_jamba.JambaModel.set_input_embeddings(value)

Set the input embeddings for the JambaModel.

PARAMETER DESCRIPTION
self

The instance of the JambaModel class.

TYPE: JambaModel

value

The input embeddings to be set for the model. This can be a tensor, array, or any compatible data structure. It represents the embeddings to be assigned to the model for input processing.
RETURNS DESCRIPTION

None.
Source code in mindnlp/transformers/models/jamba/modeling_jamba.py 
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
def set_input_embeddings(self, value):

    """
    Set the input embeddings for the JambaModel.

    Args:
        self (JambaModel): The instance of the JambaModel class.
        value: The input embeddings to be set for the model.
           This can be a tensor, array, or any compatible data structure.
           It represents the embeddings to be assigned to the model for input processing.

    Returns:
        None.

    Raises:
        None.
    """
    self.embed_tokens = value

mindnlp.transformers.models.jamba.modeling_jamba.JambaForCausalLM

Bases: JambaPreTrainedModel

This class represents a Jamba model for causal language modeling tasks. It is a subclass of JambaPreTrainedModel.

The JambaForCausalLM class encapsulates the architecture and functionality of the Jamba model for generating text. It includes methods for initializing the model, getting and setting input and output embeddings, setting the decoder, and forwarding the model.

ATTRIBUTE DESCRIPTION
model

The Jamba model used for text generation.

TYPE: JambaModel

vocab_size

The size of the vocabulary.

TYPE: int

lm_head

The linear layer for generating the next token in the sequence.

TYPE: Linear

router_aux_loss_coef

The coefficient for the auxiliary loss used in load balancing.

TYPE: float

num_experts

The number of experts used in load balancing.

TYPE: int

num_experts_per_tok

The number of experts per token used in load balancing.

TYPE: int

METHOD DESCRIPTION
__init__

Initializes the JambaForCausalLM instance with the given configuration.
get_input_embeddings

Returns the input embeddings of the model.
set_input_embeddings

Sets the input embeddings of the model.
get_output_embeddings

Returns the output embeddings of the model.
set_output_embeddings

Sets the output embeddings of the model.
set_decoder

Sets the decoder of the model.
get_decoder

Returns the decoder of the model.
forward

Constructs the model for generating text and returns the outputs.
prepare_inputs_for_generation

Prepares the inputs for text generation by reordering the cache and updating the position ids.

Please refer to the source code for more details on the implementation of each method.

Source code in mindnlp/transformers/models/jamba/modeling_jamba.py 
1709
1710
1711
1712
1713
1714
1715
1716
1717
1718
1719
1720
1721
1722
1723
1724
1725
1726
1727
1728
1729
1730
1731
1732
1733
1734
1735
1736
1737
1738
1739
1740
1741
1742
1743
1744
1745
1746
1747
1748
1749
1750
1751
1752
1753
1754
1755
1756
1757
1758
1759
1760
1761
1762
1763
1764
1765
1766
1767
1768
1769
1770
1771
1772
1773
1774
1775
1776
1777
1778
1779
1780
1781
1782
1783
1784
1785
1786
1787
1788
1789
1790
1791
1792
1793
1794
1795
1796
1797
1798
1799
1800
1801
1802
1803
1804
1805
1806
1807
1808
1809
1810
1811
1812
1813
1814
1815
1816
1817
1818
1819
1820
1821
1822
1823
1824
1825
1826
1827
1828
1829
1830
1831
1832
1833
1834
1835
1836
1837
1838
1839
1840
1841
1842
1843
1844
1845
1846
1847
1848
1849
1850
1851
1852
1853
1854
1855
1856
1857
1858
1859
1860
1861
1862
1863
1864
1865
1866
1867
1868
1869
1870
1871
1872
1873
1874
1875
1876
1877
1878
1879
1880
1881
1882
1883
1884
1885
1886
1887
1888
1889
1890
1891
1892
1893
1894
1895
1896
1897
1898
1899
1900
1901
1902
1903
1904
1905
1906
1907
1908
1909
1910
1911
1912
1913
1914
1915
1916
1917
1918
1919
1920
1921
1922
1923
1924
1925
1926
1927
1928
1929
1930
1931
1932
1933
1934
1935
1936
1937
1938
1939
1940
1941
1942
1943
1944
1945
1946
1947
1948
1949
1950
1951
1952
1953
1954
1955
1956
1957
1958
1959
1960
1961
1962
1963
1964
1965
1966
1967
1968
1969
1970
1971
1972
1973
1974
1975
1976
1977
1978
1979
1980
1981
1982
1983
1984
1985
1986
1987
1988
1989
1990
1991
1992
1993
1994
1995
1996
1997
1998
1999
2000
2001
2002
2003
2004
2005
2006
2007
2008
2009
2010
2011
2012
2013
2014
2015
2016
2017
2018
2019
2020
2021
2022
2023
2024
2025
2026
2027
2028
2029
2030
2031
2032
2033
2034
2035
2036
2037
2038
2039
2040
2041
2042
2043
2044
2045
2046
2047
2048
2049
2050
2051
2052
2053
2054
2055
2056
2057
2058
2059
2060
2061
2062
2063
2064
2065
2066
2067
2068
2069
2070
2071
2072
2073
2074
2075
2076
2077
2078
2079
2080
2081
2082
2083
2084
class JambaForCausalLM(JambaPreTrainedModel):

    """
    This class represents a Jamba model for causal language modeling tasks. It is a subclass of JambaPreTrainedModel.

    The JambaForCausalLM class encapsulates the architecture and functionality of the Jamba model for generating text.
    It includes methods for initializing the model, getting and setting input and output
    embeddings, setting the decoder, and forwarding the model.

    Attributes:
        model (JambaModel): The Jamba model used for text generation.
        vocab_size (int): The size of the vocabulary.
        lm_head (nn.Linear): The linear layer for generating the next token in the sequence.
        router_aux_loss_coef (float): The coefficient for the auxiliary loss used in load balancing.
        num_experts (int): The number of experts used in load balancing.
        num_experts_per_tok (int): The number of experts per token used in load balancing.

    Methods:
        __init__: Initializes the JambaForCausalLM instance with the given configuration.
        get_input_embeddings: Returns the input embeddings of the model.
        set_input_embeddings: Sets the input embeddings of the model.
        get_output_embeddings: Returns the output embeddings of the model.
        set_output_embeddings: Sets the output embeddings of the model.
        set_decoder: Sets the decoder of the model.
        get_decoder: Returns the decoder of the model.
        forward: Constructs the model for generating text and returns the outputs.
        prepare_inputs_for_generation: Prepares the inputs for text generation by reordering the cache and
            updating the position ids.

    Please refer to the source code for more details on the implementation of each method.
    """
    _tied_weights_keys = ["lm_head.weight"]

    def __init__(self, config: JambaConfig):

        """
        Initializes an instance of the JambaForCausalLM class.

        Args:
            self: The instance of the class.
            config (JambaConfig): An instance of JambaConfig containing the configuration parameters for the model.

        Returns:
            None.

        Raises:
            None.
        """
        super().__init__(config)
        self.model = JambaModel(config)
        self.vocab_size = config.vocab_size
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.router_aux_loss_coef = config.router_aux_loss_coef
        self.num_experts = config.num_experts
        self.num_experts_per_tok = config.num_experts_per_tok
        # Initialize weights and apply final processing
        self.post_init()

    def get_input_embeddings(self):

        """
        This method retrieves the input embeddings from the JambaForCausalLM model.

        Args:
            self: An instance of the JambaForCausalLM class.

        Returns:
            embed_tokens: The method returns the embed_tokens attribute of the model.

        Raises:
            None.
        """
        return self.model.embed_tokens

    def set_input_embeddings(self, value):

        """
            Set the input embeddings for the JambaForCausalLM model.

            Args:
                self (JambaForCausalLM): The instance of the JambaForCausalLM class.
                value (object): The input embeddings to be set.

            Returns:
                None.

            Raises:
                None.
            """
        self.model.embed_tokens = value

    def get_output_embeddings(self):

        """
        Returns the output embeddings of the JambaForCausalLM model.

        Args:
            self: An instance of the JambaForCausalLM class.

        Returns:
            None.

        Raises:
            None.
        """
        return self.lm_head

    def set_output_embeddings(self, new_embeddings):

        """
        Method to set new output embeddings for the JambaForCausalLM model.

        Args:
            self (JambaForCausalLM): The instance of the JambaForCausalLM class.
            new_embeddings (Any): The new embeddings to be set as the output embeddings for the model.
              This can be of any type.

        Returns:
            None.

        Raises:
            None.
        """
        self.lm_head = new_embeddings

    def set_decoder(self, decoder):

        """
        Sets the decoder for the JambaForCausalLM model.

        Args:
            self (JambaForCausalLM): The instance of the JambaForCausalLM class.
            decoder: The decoder object to be set for the JambaForCausalLM model.

        Returns:
            None.

        Raises:
            None.
        """
        self.model = decoder

    def get_decoder(self):

        """
        This method returns the decoder model for the JambaForCausalLM class.

        Args:
            self: The instance of the JambaForCausalLM class.

        Returns:
            The decoder model associated with the instance of the JambaForCausalLM class.

        Raises:
            None.
        """
        return self.model

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

            calc_logits_for_entire_prompt (`bool`, *optional*):
                Whether or not to calculate the logits for the entire prompt, or just the last token. Only last token
                logits are needed for generation, and calculating them only for that token can save memory,
                which becomes pretty significant for long sequences.

        Returns:
            Union[Tuple, MoeCausalLMOutputWithPast]
        ```"""
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_router_logits = (
            output_router_logits if output_router_logits is not None else self.config.output_router_logits
        )

        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
        outputs = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            output_router_logits=output_router_logits,
            return_dict=return_dict,
        )

        hidden_states = outputs[0]
        if calc_logits_for_entire_prompt:
            logits = self.lm_head(hidden_states)
        else:
            logits = self.lm_head(hidden_states[..., -1:, :])
        logits = logits.float()

        loss = None
        if labels is not None:
            # Shift so that tokens < n predict n
            shift_logits = logits[..., :-1, :]
            shift_labels = labels[..., 1:]
            # Flatten the tokens
            shift_logits = shift_logits.view(-1, self.config.vocab_size)
            shift_labels = shift_labels.view(-1)
            # Enable model parallelism
            loss = ops.cross_entropy(shift_logits, shift_labels)

        aux_loss = None
        if output_router_logits:
            aux_loss = load_balancing_loss_func(
                outputs.router_logits if return_dict else outputs[-1],
                self.num_experts,
                self.num_experts_per_tok,
                attention_mask,
            )
            if labels is not None:
                loss += self.router_aux_loss_coef * aux_loss.to(loss.device)  # make sure to reside in the same device

        if not return_dict:
            output = (logits,) + outputs[1:]
            if output_router_logits:
                output = (aux_loss,) + output
            return (loss,) + output if loss is not None else output

        return MoeCausalLMOutputWithPast(
            loss=loss,
            aux_loss=aux_loss,
            logits=logits,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
            router_logits=outputs.router_logits,
        )

    def prepare_inputs_for_generation(
            self,
            input_ids,
            past_key_values=None,
            attention_mask=None,
            inputs_embeds=None,
            output_router_logits=False,
            **kwargs,
    ):

        """
        Prepare inputs for generation in the JambaForCausalLM class.

        Args:
            self (object): The instance of the JambaForCausalLM class.
            input_ids (torch.Tensor): The input token IDs for the generation process.
            past_key_values (Union[Tuple, Cache]): The past key values from previous generations,
                which can be a Tuple or a Cache object.
            attention_mask (torch.Tensor): Optional tensor indicating the attention mask for the input.
            inputs_embeds (torch.Tensor): Optional tensor containing the embeddings for the input tokens.
            output_router_logits (bool): Flag indicating whether to calculate output router logits.

        Returns:
            dict: A dictionary containing the prepared model inputs for generation, including input_ids, position_ids,
                past_key_values, use_cache, attention_mask, output_router_logits, and calc_logits_for_entire_prompt.

        Raises:
            ValueError: If the shape of past_key_values is not as expected.
            AttributeError: If an attribute error occurs during the method execution.
            RuntimeError: If a runtime error occurs during the method execution.
        """
        # Omit tokens covered by past_key_values
        if past_key_values is not None:
            # the cache may be in the stardard format (e.g. in contrastive search), convert to Jamba's format if needed
            if isinstance(past_key_values, Tuple):
                if past_key_values[self.model._mamba_layer_index][0].shape[2] > 1:
                    past_key_values = self._convert_to_jamba_cache(past_key_values)

            if isinstance(past_key_values, Cache):
                if not isinstance(past_key_values, HybridMambaAttentionDynamicCache):
                    past_key_values = HybridMambaAttentionDynamicCache.from_legacy_cache(
                        past_key_values.to_legacy_cache()
                    )
                cache_length = past_key_values.get_seq_length()
                past_length = past_key_values.seen_tokens
                max_cache_length = past_key_values.get_max_length()
            else:
                cache_length = past_length = past_key_values[self.model._attn_layer_index][0].shape[2]
                max_cache_length = None

            # Keep only the unprocessed tokens:
            # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
            # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as
            # input)
            if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]:
                input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]
            # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
            # input_ids based on the past_length.
            elif past_length < input_ids.shape[1]:
                input_ids = input_ids[:, past_length:]
            # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens.

            # If we are about to go beyond the maximum cache length, we need to crop the input attention mask.
            if (
                    max_cache_length is not None
                    and attention_mask is not None
                    and cache_length + input_ids.shape[1] > max_cache_length
            ):
                attention_mask = attention_mask[:, -max_cache_length:]

        position_ids = kwargs.get("position_ids", None)
        if attention_mask is not None and position_ids is None:
            # create position_ids on the fly for batch generation
            position_ids = attention_mask.long().cumsum(-1) - 1
            position_ids.masked_fill_(attention_mask == 0, 1)
            if past_key_values:
                position_ids = position_ids[:, -input_ids.shape[1] :]

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

        model_inputs.update(
            {
                "position_ids": position_ids,
                "past_key_values": past_key_values,
                "use_cache": kwargs.get("use_cache"),
                "attention_mask": attention_mask,
                "output_router_logits": output_router_logits,
                "calc_logits_for_entire_prompt": self.config.calc_logits_for_entire_prompt,
            }
        )
        return model_inputs

    @staticmethod
    def _reorder_cache(past_key_values, beam_idx):

        """
        Reorders the cache based on the provided beam index.

        Args:
            past_key_values (tuple): A tuple of past key-value states for each layer in the model.
            beam_idx (torch.Tensor): A tensor containing indices to reorder the past key-value states.

        Returns:
            tuple: A tuple of reordered past key-value states for each layer in the model.

        Raises:
            None
        """
        reordered_past = ()
        for layer_past in past_key_values:
            reordered_past += (
                tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
            )
        return reordered_past

mindnlp.transformers.models.jamba.modeling_jamba.JambaForCausalLM.__init__(config)

Initializes an instance of the JambaForCausalLM class.

PARAMETER DESCRIPTION
self

The instance of the class.
config

An instance of JambaConfig containing the configuration parameters for the model.

TYPE: JambaConfig

RETURNS DESCRIPTION

None.
Source code in mindnlp/transformers/models/jamba/modeling_jamba.py 
1742
1743
1744
1745
1746
1747
1748
1749
1750
1751
1752
1753
1754
1755
1756
1757
1758
1759
1760
1761
1762
1763
1764
1765
def __init__(self, config: JambaConfig):

    """
    Initializes an instance of the JambaForCausalLM class.

    Args:
        self: The instance of the class.
        config (JambaConfig): An instance of JambaConfig containing the configuration parameters for the model.

    Returns:
        None.

    Raises:
        None.
    """
    super().__init__(config)
    self.model = JambaModel(config)
    self.vocab_size = config.vocab_size
    self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
    self.router_aux_loss_coef = config.router_aux_loss_coef
    self.num_experts = config.num_experts
    self.num_experts_per_tok = config.num_experts_per_tok
    # Initialize weights and apply final processing
    self.post_init()

mindnlp.transformers.models.jamba.modeling_jamba.JambaForCausalLM.forward(input_ids=None, attention_mask=None, position_ids=None, past_key_values=None, inputs_embeds=None, labels=None, use_cache=None, output_attentions=None, output_hidden_states=None, output_router_logits=None, return_dict=None, calc_logits_for_entire_prompt=True)

PARAMETER DESCRIPTION
labels

Labels for computing the masked language modeling loss. Indices should either be in [0, ..., config.vocab_size] or -100 (see input_ids docstring). Tokens with indices set to -100 are ignored (masked), the loss is only computed for the tokens with labels in [0, ..., config.vocab_size].

TYPE: `mindspore.Tensor` of shape `(batch_size, sequence_length)`, *optional* DEFAULT: None

calc_logits_for_entire_prompt

Whether or not to calculate the logits for the entire prompt, or just the last token. Only last token logits are needed for generation, and calculating them only for that token can save memory, which becomes pretty significant for long sequences.

TYPE: `bool`, *optional* DEFAULT: True

RETURNS DESCRIPTION
Union[Tuple, MoeCausalLMOutputWithPast]

Union[Tuple, MoeCausalLMOutputWithPast]

```

Source code in mindnlp/transformers/models/jamba/modeling_jamba.py 
1868
1869
1870
1871
1872
1873
1874
1875
1876
1877
1878
1879
1880
1881
1882
1883
1884
1885
1886
1887
1888
1889
1890
1891
1892
1893
1894
1895
1896
1897
1898
1899
1900
1901
1902
1903
1904
1905
1906
1907
1908
1909
1910
1911
1912
1913
1914
1915
1916
1917
1918
1919
1920
1921
1922
1923
1924
1925
1926
1927
1928
1929
1930
1931
1932
1933
1934
1935
1936
1937
1938
1939
1940
1941
1942
1943
1944
1945
1946
1947
1948
1949
1950
1951
1952
1953
1954
1955
1956
1957
1958
1959
1960
1961
1962
1963
1964
1965
def forward(
        self,
        input_ids: mindspore.Tensor = None,
        attention_mask: Optional[mindspore.Tensor] = None,
        position_ids: Optional[mindspore.Tensor] = None,
        past_key_values: Optional[List[mindspore.Tensor]] = None,
        inputs_embeds: Optional[mindspore.Tensor] = None,
        labels: Optional[mindspore.Tensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        output_router_logits: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        calc_logits_for_entire_prompt: Optional[bool] = True,
) -> Union[Tuple, MoeCausalLMOutputWithPast]:
    r"""
    Args:
        labels (`mindspore.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.

        calc_logits_for_entire_prompt (`bool`, *optional*):
            Whether or not to calculate the logits for the entire prompt, or just the last token. Only last token
            logits are needed for generation, and calculating them only for that token can save memory,
            which becomes pretty significant for long sequences.

    Returns:
        Union[Tuple, MoeCausalLMOutputWithPast]
    ```"""
    output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
    output_router_logits = (
        output_router_logits if output_router_logits is not None else self.config.output_router_logits
    )

    output_hidden_states = (
        output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
    )
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict

    # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
    outputs = self.model(
        input_ids=input_ids,
        attention_mask=attention_mask,
        position_ids=position_ids,
        past_key_values=past_key_values,
        inputs_embeds=inputs_embeds,
        use_cache=use_cache,
        output_attentions=output_attentions,
        output_hidden_states=output_hidden_states,
        output_router_logits=output_router_logits,
        return_dict=return_dict,
    )

    hidden_states = outputs[0]
    if calc_logits_for_entire_prompt:
        logits = self.lm_head(hidden_states)
    else:
        logits = self.lm_head(hidden_states[..., -1:, :])
    logits = logits.float()

    loss = None
    if labels is not None:
        # Shift so that tokens < n predict n
        shift_logits = logits[..., :-1, :]
        shift_labels = labels[..., 1:]
        # Flatten the tokens
        shift_logits = shift_logits.view(-1, self.config.vocab_size)
        shift_labels = shift_labels.view(-1)
        # Enable model parallelism
        loss = ops.cross_entropy(shift_logits, shift_labels)

    aux_loss = None
    if output_router_logits:
        aux_loss = load_balancing_loss_func(
            outputs.router_logits if return_dict else outputs[-1],
            self.num_experts,
            self.num_experts_per_tok,
            attention_mask,
        )
        if labels is not None:
            loss += self.router_aux_loss_coef * aux_loss.to(loss.device)  # make sure to reside in the same device

    if not return_dict:
        output = (logits,) + outputs[1:]
        if output_router_logits:
            output = (aux_loss,) + output
        return (loss,) + output if loss is not None else output

    return MoeCausalLMOutputWithPast(
        loss=loss,
        aux_loss=aux_loss,
        logits=logits,
        past_key_values=outputs.past_key_values,
        hidden_states=outputs.hidden_states,
        attentions=outputs.attentions,
        router_logits=outputs.router_logits,
    )

mindnlp.transformers.models.jamba.modeling_jamba.JambaForCausalLM.get_decoder()

This method returns the decoder model for the JambaForCausalLM class.

PARAMETER DESCRIPTION
self

The instance of the JambaForCausalLM class.
RETURNS DESCRIPTION

The decoder model associated with the instance of the JambaForCausalLM class.
Source code in mindnlp/transformers/models/jamba/modeling_jamba.py 
1851
1852
1853
1854
1855
1856
1857
1858
1859
1860
1861
1862
1863
1864
1865
def get_decoder(self):

    """
    This method returns the decoder model for the JambaForCausalLM class.

    Args:
        self: The instance of the JambaForCausalLM class.

    Returns:
        The decoder model associated with the instance of the JambaForCausalLM class.

    Raises:
        None.
    """
    return self.model

mindnlp.transformers.models.jamba.modeling_jamba.JambaForCausalLM.get_input_embeddings()

This method retrieves the input embeddings from the JambaForCausalLM model.

PARAMETER DESCRIPTION
self

An instance of the JambaForCausalLM class.
RETURNS DESCRIPTION
embed_tokens

The method returns the embed_tokens attribute of the model.
Source code in mindnlp/transformers/models/jamba/modeling_jamba.py 
1767
1768
1769
1770
1771
1772
1773
1774
1775
1776
1777
1778
1779
1780
1781
def get_input_embeddings(self):

    """
    This method retrieves the input embeddings from the JambaForCausalLM model.

    Args:
        self: An instance of the JambaForCausalLM class.

    Returns:
        embed_tokens: The method returns the embed_tokens attribute of the model.

    Raises:
        None.
    """
    return self.model.embed_tokens

mindnlp.transformers.models.jamba.modeling_jamba.JambaForCausalLM.get_output_embeddings()

Returns the output embeddings of the JambaForCausalLM model.

PARAMETER DESCRIPTION
self

An instance of the JambaForCausalLM class.
RETURNS DESCRIPTION

None.
Source code in mindnlp/transformers/models/jamba/modeling_jamba.py 
1800
1801
1802
1803
1804
1805
1806
1807
1808
1809
1810
1811
1812
1813
1814
def get_output_embeddings(self):

    """
    Returns the output embeddings of the JambaForCausalLM model.

    Args:
        self: An instance of the JambaForCausalLM class.

    Returns:
        None.

    Raises:
        None.
    """
    return self.lm_head

mindnlp.transformers.models.jamba.modeling_jamba.JambaForCausalLM.prepare_inputs_for_generation(input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, output_router_logits=False, **kwargs)

Prepare inputs for generation in the JambaForCausalLM class.

PARAMETER DESCRIPTION
self

The instance of the JambaForCausalLM class.

TYPE: object

input_ids

The input token IDs for the generation process.

TYPE: Tensor

past_key_values

The past key values from previous generations, which can be a Tuple or a Cache object.

TYPE: Union[Tuple, Cache] DEFAULT: None

attention_mask

Optional tensor indicating the attention mask for the input.

TYPE: Tensor DEFAULT: None

inputs_embeds

Optional tensor containing the embeddings for the input tokens.

TYPE: Tensor DEFAULT: None

output_router_logits

Flag indicating whether to calculate output router logits.

TYPE: bool DEFAULT: False

RETURNS DESCRIPTION
dict

A dictionary containing the prepared model inputs for generation, including input_ids, position_ids, past_key_values, use_cache, attention_mask, output_router_logits, and calc_logits_for_entire_prompt.
RAISES DESCRIPTION
ValueError

If the shape of past_key_values is not as expected.
AttributeError

If an attribute error occurs during the method execution.
RuntimeError

If a runtime error occurs during the method execution.
Source code in mindnlp/transformers/models/jamba/modeling_jamba.py 
1967
1968
1969
1970
1971
1972
1973
1974
1975
1976
1977
1978
1979
1980
1981
1982
1983
1984
1985
1986
1987
1988
1989
1990
1991
1992
1993
1994
1995
1996
1997
1998
1999
2000
2001
2002
2003
2004
2005
2006
2007
2008
2009
2010
2011
2012
2013
2014
2015
2016
2017
2018
2019
2020
2021
2022
2023
2024
2025
2026
2027
2028
2029
2030
2031
2032
2033
2034
2035
2036
2037
2038
2039
2040
2041
2042
2043
2044
2045
2046
2047
2048
2049
2050
2051
2052
2053
2054
2055
2056
2057
2058
2059
2060
2061
def prepare_inputs_for_generation(
        self,
        input_ids,
        past_key_values=None,
        attention_mask=None,
        inputs_embeds=None,
        output_router_logits=False,
        **kwargs,
):

    """
    Prepare inputs for generation in the JambaForCausalLM class.

    Args:
        self (object): The instance of the JambaForCausalLM class.
        input_ids (torch.Tensor): The input token IDs for the generation process.
        past_key_values (Union[Tuple, Cache]): The past key values from previous generations,
            which can be a Tuple or a Cache object.
        attention_mask (torch.Tensor): Optional tensor indicating the attention mask for the input.
        inputs_embeds (torch.Tensor): Optional tensor containing the embeddings for the input tokens.
        output_router_logits (bool): Flag indicating whether to calculate output router logits.

    Returns:
        dict: A dictionary containing the prepared model inputs for generation, including input_ids, position_ids,
            past_key_values, use_cache, attention_mask, output_router_logits, and calc_logits_for_entire_prompt.

    Raises:
        ValueError: If the shape of past_key_values is not as expected.
        AttributeError: If an attribute error occurs during the method execution.
        RuntimeError: If a runtime error occurs during the method execution.
    """
    # Omit tokens covered by past_key_values
    if past_key_values is not None:
        # the cache may be in the stardard format (e.g. in contrastive search), convert to Jamba's format if needed
        if isinstance(past_key_values, Tuple):
            if past_key_values[self.model._mamba_layer_index][0].shape[2] > 1:
                past_key_values = self._convert_to_jamba_cache(past_key_values)

        if isinstance(past_key_values, Cache):
            if not isinstance(past_key_values, HybridMambaAttentionDynamicCache):
                past_key_values = HybridMambaAttentionDynamicCache.from_legacy_cache(
                    past_key_values.to_legacy_cache()
                )
            cache_length = past_key_values.get_seq_length()
            past_length = past_key_values.seen_tokens
            max_cache_length = past_key_values.get_max_length()
        else:
            cache_length = past_length = past_key_values[self.model._attn_layer_index][0].shape[2]
            max_cache_length = None

        # Keep only the unprocessed tokens:
        # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
        # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as
        # input)
        if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]:
            input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]
        # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
        # input_ids based on the past_length.
        elif past_length < input_ids.shape[1]:
            input_ids = input_ids[:, past_length:]
        # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens.

        # If we are about to go beyond the maximum cache length, we need to crop the input attention mask.
        if (
                max_cache_length is not None
                and attention_mask is not None
                and cache_length + input_ids.shape[1] > max_cache_length
        ):
            attention_mask = attention_mask[:, -max_cache_length:]

    position_ids = kwargs.get("position_ids", None)
    if attention_mask is not None and position_ids is None:
        # create position_ids on the fly for batch generation
        position_ids = attention_mask.long().cumsum(-1) - 1
        position_ids.masked_fill_(attention_mask == 0, 1)
        if past_key_values:
            position_ids = position_ids[:, -input_ids.shape[1] :]

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

    model_inputs.update(
        {
            "position_ids": position_ids,
            "past_key_values": past_key_values,
            "use_cache": kwargs.get("use_cache"),
            "attention_mask": attention_mask,
            "output_router_logits": output_router_logits,
            "calc_logits_for_entire_prompt": self.config.calc_logits_for_entire_prompt,
        }
    )
    return model_inputs

mindnlp.transformers.models.jamba.modeling_jamba.JambaForCausalLM.set_decoder(decoder)

Sets the decoder for the JambaForCausalLM model.

PARAMETER DESCRIPTION
self

The instance of the JambaForCausalLM class.

TYPE: JambaForCausalLM

decoder

The decoder object to be set for the JambaForCausalLM model.
RETURNS DESCRIPTION

None.
Source code in mindnlp/transformers/models/jamba/modeling_jamba.py 
1834
1835
1836
1837
1838
1839
1840
1841
1842
1843
1844
1845
1846
1847
1848
1849
def set_decoder(self, decoder):

    """
    Sets the decoder for the JambaForCausalLM model.

    Args:
        self (JambaForCausalLM): The instance of the JambaForCausalLM class.
        decoder: The decoder object to be set for the JambaForCausalLM model.

    Returns:
        None.

    Raises:
        None.
    """
    self.model = decoder

mindnlp.transformers.models.jamba.modeling_jamba.JambaForCausalLM.set_input_embeddings(value)

Set the input embeddings for the JambaForCausalLM model.

PARAMETER DESCRIPTION
self

The instance of the JambaForCausalLM class.

TYPE: JambaForCausalLM

value

The input embeddings to be set.

TYPE: object

RETURNS DESCRIPTION

None.
Source code in mindnlp/transformers/models/jamba/modeling_jamba.py 
1783
1784
1785
1786
1787
1788
1789
1790
1791
1792
1793
1794
1795
1796
1797
1798
def set_input_embeddings(self, value):

    """
        Set the input embeddings for the JambaForCausalLM model.

        Args:
            self (JambaForCausalLM): The instance of the JambaForCausalLM class.
            value (object): The input embeddings to be set.

        Returns:
            None.

        Raises:
            None.
        """
    self.model.embed_tokens = value

mindnlp.transformers.models.jamba.modeling_jamba.JambaForCausalLM.set_output_embeddings(new_embeddings)

Method to set new output embeddings for the JambaForCausalLM model.

PARAMETER DESCRIPTION
self

The instance of the JambaForCausalLM class.

TYPE: JambaForCausalLM

new_embeddings

The new embeddings to be set as the output embeddings for the model. This can be of any type.

TYPE: Any

RETURNS DESCRIPTION

None.
Source code in mindnlp/transformers/models/jamba/modeling_jamba.py 
1816
1817
1818
1819
1820
1821
1822
1823
1824
1825
1826
1827
1828
1829
1830
1831
1832
def set_output_embeddings(self, new_embeddings):

    """
    Method to set new output embeddings for the JambaForCausalLM model.

    Args:
        self (JambaForCausalLM): The instance of the JambaForCausalLM class.
        new_embeddings (Any): The new embeddings to be set as the output embeddings for the model.
          This can be of any type.

    Returns:
        None.

    Raises:
        None.
    """
    self.lm_head = new_embeddings

mindnlp.transformers.models.jamba.modeling_jamba.JambaForSequenceClassification

Bases: JambaPreTrainedModel

JambaForSequenceClassification is a class that represents a sequence classification model based on the Jamba architecture.

This class extends JambaPreTrainedModel and includes methods for initializing the model, getting and setting input embeddings, and forwarding the sequence classification output.

The forward method takes input_ids, attention_mask, position_ids, past_key_values, inputs_embeds, labels, and various optional arguments to generate the sequence classifier output.

It calculates the loss based on the labels provided and handles different types of classification problems such as regression, single-label classification, and multi-label classification.

The class provides flexibility in handling return types and outputs a SequenceClassifierOutputWithPast object if return_dict is set to True.

Source code in mindnlp/transformers/models/jamba/modeling_jamba.py 
2088
2089
2090
2091
2092
2093
2094
2095
2096
2097
2098
2099
2100
2101
2102
2103
2104
2105
2106
2107
2108
2109
2110
2111
2112
2113
2114
2115
2116
2117
2118
2119
2120
2121
2122
2123
2124
2125
2126
2127
2128
2129
2130
2131
2132
2133
2134
2135
2136
2137
2138
2139
2140
2141
2142
2143
2144
2145
2146
2147
2148
2149
2150
2151
2152
2153
2154
2155
2156
2157
2158
2159
2160
2161
2162
2163
2164
2165
2166
2167
2168
2169
2170
2171
2172
2173
2174
2175
2176
2177
2178
2179
2180
2181
2182
2183
2184
2185
2186
2187
2188
2189
2190
2191
2192
2193
2194
2195
2196
2197
2198
2199
2200
2201
2202
2203
2204
2205
2206
2207
2208
2209
2210
2211
2212
2213
2214
2215
2216
2217
2218
2219
2220
2221
2222
2223
2224
2225
2226
2227
2228
2229
2230
2231
2232
2233
2234
2235
2236
2237
2238
2239
2240
2241
2242
2243
2244
2245
2246
2247
2248
2249
class JambaForSequenceClassification(JambaPreTrainedModel):

    """
    JambaForSequenceClassification is a class that represents a sequence classification model based on the
    Jamba architecture.

    This class extends JambaPreTrainedModel and includes methods for initializing the model,
    getting and setting input embeddings, and forwarding the sequence classification output.

    The forward method takes input_ids, attention_mask, position_ids, past_key_values, inputs_embeds, labels,
    and various optional arguments to generate the sequence classifier output.

    It calculates the loss based on the labels provided and handles different types of classification problems
    such as regression, single-label classification, and multi-label classification.

    The class provides flexibility in handling return types and outputs a SequenceClassifierOutputWithPast object
    if return_dict is set to True.
    """
    def __init__(self, config):

        """
        Initializes a new instance of the JambaForSequenceClassification class.

        Args:
            self: The object itself.
            config: An instance of the JambaConfig class that contains the configuration settings for the Jamba model.

        Returns:
            None.

        Raises:
            None.
        """
        super().__init__(config)
        self.num_labels = config.num_labels
        self.model = JambaModel(config)
        self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)

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

    def get_input_embeddings(self):

        """
        Method to retrieve input embeddings from the model for JambaForSequenceClassification.

        Args:
            self: The instance of the JambaForSequenceClassification class.

        Returns:
            The input embeddings obtained from the model's embed_tokens attribute.

        Raises:
            None.
        """
        return self.model.embed_tokens

    def set_input_embeddings(self, value):

        """
        Set the input embeddings for the JambaForSequenceClassification model.

        Args:
            self (JambaForSequenceClassification): The instance of the JambaForSequenceClassification class.
            value (torch.Tensor): The input embeddings to be set for the model.
                Should be a torch.Tensor representing the embeddings to be used.

        Returns:
            None.

        Raises:
            None.
        """
        self.model.embed_tokens = value

    def forward(
            self,
            input_ids: mindspore.Tensor = None,
            attention_mask: Optional[mindspore.Tensor] = None,
            position_ids: Optional[mindspore.Tensor] = None,
            past_key_values: Optional[List[mindspore.Tensor]] = None,
            inputs_embeds: Optional[mindspore.Tensor] = None,
            labels: Optional[mindspore.Tensor] = None,
            use_cache: Optional[bool] = None,
            output_attentions: Optional[bool] = None,
            output_hidden_states: Optional[bool] = None,
            return_dict: Optional[bool] = None,
    ) -> Union[Tuple, SequenceClassifierOutputWithPast]:
        r"""
        Args:
            labels (`mindspore.Tensor` of shape `(batch_size,)`, *optional*):
                Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
                config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
                `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
        """
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        transformer_outputs = self.model(
            input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        hidden_states = transformer_outputs[0]
        logits = self.score(hidden_states)

        if input_ids is not None:
            batch_size = input_ids.shape[0]
        else:
            batch_size = inputs_embeds.shape[0]

        if self.config.pad_token_id is None and batch_size != 1:
            raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
        if self.config.pad_token_id is None:
            sequence_lengths = -1
        else:
            if input_ids is not None:
                # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility
                sequence_lengths = ops.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1
                sequence_lengths = sequence_lengths % input_ids.shape[-1]
                sequence_lengths = sequence_lengths.to(logits.device)
            else:
                sequence_lengths = -1

        pooled_logits = logits[ops.arange(batch_size), sequence_lengths]

        loss = None
        if labels is not None:
            labels = labels.to(logits.device)
            if self.config.problem_type is None:
                if self.num_labels == 1:
                    self.config.problem_type = "regression"
                elif self.num_labels > 1 and labels.dtype in (mindspore.int64, mindspore.int32):
                    self.config.problem_type = "single_label_classification"
                else:
                    self.config.problem_type = "multi_label_classification"

            if self.config.problem_type == "regression":
                if self.num_labels == 1:
                    loss = ops.mse_loss(pooled_logits.squeeze(), labels.squeeze())
                else:
                    loss = ops.mse_loss(pooled_logits, labels)
            elif self.config.problem_type == "single_label_classification":
                loss = ops.cross_entropy(pooled_logits.view(-1, self.num_labels), labels.view(-1))
            elif self.config.problem_type == "multi_label_classification":
                loss = ops.binary_cross_entropy_with_logits(pooled_logits, labels)
        if not return_dict:
            output = (pooled_logits,) + transformer_outputs[1:]
            return ((loss,) + output) if loss is not None else output

        return SequenceClassifierOutputWithPast(
            loss=loss,
            logits=pooled_logits,
            past_key_values=transformer_outputs.past_key_values,
            hidden_states=transformer_outputs.hidden_states,
            attentions=transformer_outputs.attentions,
        )

mindnlp.transformers.models.jamba.modeling_jamba.JambaForSequenceClassification.__init__(config)

Initializes a new instance of the JambaForSequenceClassification class.

PARAMETER DESCRIPTION
self

The object itself.
config

An instance of the JambaConfig class that contains the configuration settings for the Jamba model.
RETURNS DESCRIPTION

None.
Source code in mindnlp/transformers/models/jamba/modeling_jamba.py 
2106
2107
2108
2109
2110
2111
2112
2113
2114
2115
2116
2117
2118
2119
2120
2121
2122
2123
2124
2125
2126
2127
def __init__(self, config):

    """
    Initializes a new instance of the JambaForSequenceClassification class.

    Args:
        self: The object itself.
        config: An instance of the JambaConfig class that contains the configuration settings for the Jamba model.

    Returns:
        None.

    Raises:
        None.
    """
    super().__init__(config)
    self.num_labels = config.num_labels
    self.model = JambaModel(config)
    self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)

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

mindnlp.transformers.models.jamba.modeling_jamba.JambaForSequenceClassification.forward(input_ids=None, attention_mask=None, position_ids=None, past_key_values=None, inputs_embeds=None, labels=None, use_cache=None, output_attentions=None, output_hidden_states=None, return_dict=None)

PARAMETER DESCRIPTION
labels

Labels for computing the sequence classification/regression loss. Indices should be in [0, ..., config.num_labels - 1]. If config.num_labels == 1 a regression loss is computed (Mean-Square loss), If config.num_labels > 1 a classification loss is computed (Cross-Entropy).

TYPE: `mindspore.Tensor` of shape `(batch_size,)`, *optional* DEFAULT: None

Source code in mindnlp/transformers/models/jamba/modeling_jamba.py 
2163
2164
2165
2166
2167
2168
2169
2170
2171
2172
2173
2174
2175
2176
2177
2178
2179
2180
2181
2182
2183
2184
2185
2186
2187
2188
2189
2190
2191
2192
2193
2194
2195
2196
2197
2198
2199
2200
2201
2202
2203
2204
2205
2206
2207
2208
2209
2210
2211
2212
2213
2214
2215
2216
2217
2218
2219
2220
2221
2222
2223
2224
2225
2226
2227
2228
2229
2230
2231
2232
2233
2234
2235
2236
2237
2238
2239
2240
2241
2242
2243
2244
2245
2246
2247
2248
2249
def forward(
        self,
        input_ids: mindspore.Tensor = None,
        attention_mask: Optional[mindspore.Tensor] = None,
        position_ids: Optional[mindspore.Tensor] = None,
        past_key_values: Optional[List[mindspore.Tensor]] = None,
        inputs_embeds: Optional[mindspore.Tensor] = None,
        labels: Optional[mindspore.Tensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
) -> Union[Tuple, SequenceClassifierOutputWithPast]:
    r"""
    Args:
        labels (`mindspore.Tensor` of shape `(batch_size,)`, *optional*):
            Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
    """
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict

    transformer_outputs = self.model(
        input_ids,
        attention_mask=attention_mask,
        position_ids=position_ids,
        past_key_values=past_key_values,
        inputs_embeds=inputs_embeds,
        use_cache=use_cache,
        output_attentions=output_attentions,
        output_hidden_states=output_hidden_states,
        return_dict=return_dict,
    )
    hidden_states = transformer_outputs[0]
    logits = self.score(hidden_states)

    if input_ids is not None:
        batch_size = input_ids.shape[0]
    else:
        batch_size = inputs_embeds.shape[0]

    if self.config.pad_token_id is None and batch_size != 1:
        raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
    if self.config.pad_token_id is None:
        sequence_lengths = -1
    else:
        if input_ids is not None:
            # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility
            sequence_lengths = ops.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1
            sequence_lengths = sequence_lengths % input_ids.shape[-1]
            sequence_lengths = sequence_lengths.to(logits.device)
        else:
            sequence_lengths = -1

    pooled_logits = logits[ops.arange(batch_size), sequence_lengths]

    loss = None
    if labels is not None:
        labels = labels.to(logits.device)
        if self.config.problem_type is None:
            if self.num_labels == 1:
                self.config.problem_type = "regression"
            elif self.num_labels > 1 and labels.dtype in (mindspore.int64, mindspore.int32):
                self.config.problem_type = "single_label_classification"
            else:
                self.config.problem_type = "multi_label_classification"

        if self.config.problem_type == "regression":
            if self.num_labels == 1:
                loss = ops.mse_loss(pooled_logits.squeeze(), labels.squeeze())
            else:
                loss = ops.mse_loss(pooled_logits, labels)
        elif self.config.problem_type == "single_label_classification":
            loss = ops.cross_entropy(pooled_logits.view(-1, self.num_labels), labels.view(-1))
        elif self.config.problem_type == "multi_label_classification":
            loss = ops.binary_cross_entropy_with_logits(pooled_logits, labels)
    if not return_dict:
        output = (pooled_logits,) + transformer_outputs[1:]
        return ((loss,) + output) if loss is not None else output

    return SequenceClassifierOutputWithPast(
        loss=loss,
        logits=pooled_logits,
        past_key_values=transformer_outputs.past_key_values,
        hidden_states=transformer_outputs.hidden_states,
        attentions=transformer_outputs.attentions,
    )

mindnlp.transformers.models.jamba.modeling_jamba.JambaForSequenceClassification.get_input_embeddings()

Method to retrieve input embeddings from the model for JambaForSequenceClassification.

PARAMETER DESCRIPTION
self

The instance of the JambaForSequenceClassification class.
RETURNS DESCRIPTION

The input embeddings obtained from the model's embed_tokens attribute.
Source code in mindnlp/transformers/models/jamba/modeling_jamba.py 
2129
2130
2131
2132
2133
2134
2135
2136
2137
2138
2139
2140
2141
2142
2143
def get_input_embeddings(self):

    """
    Method to retrieve input embeddings from the model for JambaForSequenceClassification.

    Args:
        self: The instance of the JambaForSequenceClassification class.

    Returns:
        The input embeddings obtained from the model's embed_tokens attribute.

    Raises:
        None.
    """
    return self.model.embed_tokens

mindnlp.transformers.models.jamba.modeling_jamba.JambaForSequenceClassification.set_input_embeddings(value)

Set the input embeddings for the JambaForSequenceClassification model.

PARAMETER DESCRIPTION
self

The instance of the JambaForSequenceClassification class.

TYPE: JambaForSequenceClassification

value

The input embeddings to be set for the model. Should be a torch.Tensor representing the embeddings to be used.

TYPE: Tensor

RETURNS DESCRIPTION

None.
Source code in mindnlp/transformers/models/jamba/modeling_jamba.py 
2145
2146
2147
2148
2149
2150
2151
2152
2153
2154
2155
2156
2157
2158
2159
2160
2161
def set_input_embeddings(self, value):

    """
    Set the input embeddings for the JambaForSequenceClassification model.

    Args:
        self (JambaForSequenceClassification): The instance of the JambaForSequenceClassification class.
        value (torch.Tensor): The input embeddings to be set for the model.
            Should be a torch.Tensor representing the embeddings to be used.

    Returns:
        None.

    Raises:
        None.
    """
    self.model.embed_tokens = value