Skip to content

megatron_bert

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert

MindSpore MegatronBERT model.

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertAttention

Bases: Module

This class represents the attention mechanism used in Megatron-BERT models. It is a part of the Megatron-BERT architecture and is responsible for performing self-attention operations.

The MegatronBertAttention class inherits from the nn.Module class.

ATTRIBUTE DESCRIPTION
ln

Layer normalization module used in the attention mechanism.

TYPE: LayerNorm

self

Self-attention module responsible for computing attention scores.

TYPE: MegatronBertSelfAttention

output

Output module that combines attention output with the input hidden states.

TYPE: MegatronBertSelfOutput

pruned_heads

A set of pruned attention heads.

TYPE: set

METHOD DESCRIPTION
__init__

Initializes the MegatronBertAttention instance.
prune_heads

Prunes the specified attention heads.
forward

Performs the attention mechanism computation.
Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
class MegatronBertAttention(nn.Module):

    """
    This class represents the attention mechanism used in Megatron-BERT models. It is a part of the Megatron-BERT
    architecture and is responsible for performing self-attention operations.

    The MegatronBertAttention class inherits from the nn.Module class.

    Attributes:
        ln (nn.LayerNorm): Layer normalization module used in the attention mechanism.
        self (MegatronBertSelfAttention): Self-attention module responsible for computing attention scores.
        output (MegatronBertSelfOutput): Output module that combines attention output with the input hidden states.
        pruned_heads (set): A set of pruned attention heads.

    Methods:
        __init__: Initializes the MegatronBertAttention instance.
        prune_heads: Prunes the specified attention heads.
        forward: Performs the attention mechanism computation.

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

        Args:
            self (MegatronBertAttention): The current instance of the class.
            config (object):
                The configuration object containing the hyperparameters for the attention mechanism.

                - hidden_size (int): The size of the hidden state.
                - layer_norm_eps (float): The epsilon value for layer normalization.

        Returns:
            None

        Raises:
            None
        """
        super().__init__()
        self.ln = nn.LayerNorm([config.hidden_size], eps=config.layer_norm_eps)
        self.self = MegatronBertSelfAttention(config)
        self.output = MegatronBertSelfOutput(config)
        self.pruned_heads = set()

    def prune_heads(self, heads):
        """
        This method 'prune_heads' is defined within the 'MegatronBertAttention' class. It prunes specific attention
        heads from the self-attention mechanism based on the provided 'heads' parameter.

        Args:
            self: Represents the instance of the MegatronBertAttention class.
                It is used to access the attributes and methods of the class.
            heads: A list that contains the indices of the attention heads to be pruned.
                These indices correspond to the specific attention heads that should be removed from the self-attention
                mechanism.

        Returns:
            None: However, it modifies the internal state of the MegatronBertAttention instance by pruning the specified
                attention heads from the self-attention mechanism.

        Raises:
            None:
                However, potential exceptions that might occur during the execution could include:

                - TypeError: If the input parameters are not of the expected types.
                - IndexError: If there are issues with accessing elements within the 'heads' list or other data structures.
                - ValueError: If there are inconsistencies or unexpected values encountered during the pruning process.
        """
        if len(heads) == 0:
            return
        heads, index = find_pruneable_heads_and_indices(
            heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads
        )

        # Prune linear layers
        self.self.query = prune_linear_layer(self.self.query, index)
        self.self.key = prune_linear_layer(self.self.key, index)
        self.self.value = prune_linear_layer(self.self.value, index)
        self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)

        # Update hyper params and store pruned heads
        self.self.num_attention_heads = self.self.num_attention_heads - len(heads)
        self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads
        self.pruned_heads = self.pruned_heads.union(heads)

    def forward(
        self,
        hidden_states: mindspore.Tensor,
        attention_mask: Optional[mindspore.Tensor] = None,
        head_mask: Optional[mindspore.Tensor] = None,
        encoder_hidden_states: Optional[mindspore.Tensor] = None,
        encoder_attention_mask: Optional[mindspore.Tensor] = None,
        past_key_value: Optional[Tuple[Tuple[mindspore.Tensor]]] = None,
        output_attentions: Optional[bool] = False,
    ) -> Tuple[mindspore.Tensor]:
        """
        Args:
            self: The instance of the MegatronBertAttention class.
            hidden_states (mindspore.Tensor): The input hidden states for the attention mechanism.
            attention_mask (Optional[mindspore.Tensor]): Optional tensor specifying which elements should be attended to.
            head_mask (Optional[mindspore.Tensor]): Optional tensor for masking individual attention heads.
            encoder_hidden_states (Optional[mindspore.Tensor]): Optional tensor representing the hidden states of the encoder.
            encoder_attention_mask (Optional[mindspore.Tensor]): Optional tensor specifying which elements of the encoder
                hidden states should be attended to.
            past_key_value (Optional[Tuple[Tuple[mindspore.Tensor]]]): Optional tuple of past key and value tensors for
                fast decoding.
            output_attentions (Optional[bool]): Optional flag indicating whether to return attentions.

        Returns:
            Tuple[mindspore.Tensor]: A tuple containing the attention output and additional outputs from
                the attention mechanism.

        Raises:
            ValueError: If the input tensors have incompatible shapes or types.
            TypeError: If the input parameters are not of the expected types.
            RuntimeError: If there is an issue during the attention computation process.
        """
        ln_outputs = self.ln(hidden_states)
        self_outputs = self.self(
            ln_outputs,
            attention_mask,
            head_mask,
            encoder_hidden_states,
            encoder_attention_mask,
            past_key_value,
            output_attentions,
        )
        attention_output = self.output(self_outputs[0], hidden_states)
        outputs = (attention_output,) + self_outputs[1:]  # add attentions if we output them
        return outputs

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertAttention.__init__(config)

Initializes an instance of the MegatronBertAttention class.

PARAMETER DESCRIPTION
self

The current instance of the class.

TYPE: MegatronBertAttention

config

The configuration object containing the hyperparameters for the attention mechanism.

  • hidden_size (int): The size of the hidden state.
  • layer_norm_eps (float): The epsilon value for layer normalization.

TYPE: object

RETURNS DESCRIPTION

None
Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
def __init__(self, config):
    """
    Initializes an instance of the MegatronBertAttention class.

    Args:
        self (MegatronBertAttention): The current instance of the class.
        config (object):
            The configuration object containing the hyperparameters for the attention mechanism.

            - hidden_size (int): The size of the hidden state.
            - layer_norm_eps (float): The epsilon value for layer normalization.

    Returns:
        None

    Raises:
        None
    """
    super().__init__()
    self.ln = nn.LayerNorm([config.hidden_size], eps=config.layer_norm_eps)
    self.self = MegatronBertSelfAttention(config)
    self.output = MegatronBertSelfOutput(config)
    self.pruned_heads = set()

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertAttention.forward(hidden_states, attention_mask=None, head_mask=None, encoder_hidden_states=None, encoder_attention_mask=None, past_key_value=None, output_attentions=False)

PARAMETER DESCRIPTION
self

The instance of the MegatronBertAttention class.
hidden_states

The input hidden states for the attention mechanism.

TYPE: Tensor

attention_mask

Optional tensor specifying which elements should be attended to.

TYPE: Optional[Tensor] DEFAULT: None

head_mask

Optional tensor for masking individual attention heads.

TYPE: Optional[Tensor] DEFAULT: None

encoder_hidden_states

Optional tensor representing the hidden states of the encoder.

TYPE: Optional[Tensor] DEFAULT: None

encoder_attention_mask

Optional tensor specifying which elements of the encoder hidden states should be attended to.

TYPE: Optional[Tensor] DEFAULT: None

past_key_value

Optional tuple of past key and value tensors for fast decoding.

TYPE: Optional[Tuple[Tuple[Tensor]]] DEFAULT: None

output_attentions

Optional flag indicating whether to return attentions.

TYPE: Optional[bool] DEFAULT: False

RETURNS DESCRIPTION
Tuple[Tensor]

Tuple[mindspore.Tensor]: A tuple containing the attention output and additional outputs from the attention mechanism.
RAISES DESCRIPTION
ValueError

If the input tensors have incompatible shapes or types.
TypeError

If the input parameters are not of the expected types.
RuntimeError

If there is an issue during the attention computation process.
Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
def forward(
    self,
    hidden_states: mindspore.Tensor,
    attention_mask: Optional[mindspore.Tensor] = None,
    head_mask: Optional[mindspore.Tensor] = None,
    encoder_hidden_states: Optional[mindspore.Tensor] = None,
    encoder_attention_mask: Optional[mindspore.Tensor] = None,
    past_key_value: Optional[Tuple[Tuple[mindspore.Tensor]]] = None,
    output_attentions: Optional[bool] = False,
) -> Tuple[mindspore.Tensor]:
    """
    Args:
        self: The instance of the MegatronBertAttention class.
        hidden_states (mindspore.Tensor): The input hidden states for the attention mechanism.
        attention_mask (Optional[mindspore.Tensor]): Optional tensor specifying which elements should be attended to.
        head_mask (Optional[mindspore.Tensor]): Optional tensor for masking individual attention heads.
        encoder_hidden_states (Optional[mindspore.Tensor]): Optional tensor representing the hidden states of the encoder.
        encoder_attention_mask (Optional[mindspore.Tensor]): Optional tensor specifying which elements of the encoder
            hidden states should be attended to.
        past_key_value (Optional[Tuple[Tuple[mindspore.Tensor]]]): Optional tuple of past key and value tensors for
            fast decoding.
        output_attentions (Optional[bool]): Optional flag indicating whether to return attentions.

    Returns:
        Tuple[mindspore.Tensor]: A tuple containing the attention output and additional outputs from
            the attention mechanism.

    Raises:
        ValueError: If the input tensors have incompatible shapes or types.
        TypeError: If the input parameters are not of the expected types.
        RuntimeError: If there is an issue during the attention computation process.
    """
    ln_outputs = self.ln(hidden_states)
    self_outputs = self.self(
        ln_outputs,
        attention_mask,
        head_mask,
        encoder_hidden_states,
        encoder_attention_mask,
        past_key_value,
        output_attentions,
    )
    attention_output = self.output(self_outputs[0], hidden_states)
    outputs = (attention_output,) + self_outputs[1:]  # add attentions if we output them
    return outputs

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertAttention.prune_heads(heads)

This method 'prune_heads' is defined within the 'MegatronBertAttention' class. It prunes specific attention heads from the self-attention mechanism based on the provided 'heads' parameter.

PARAMETER DESCRIPTION
self

Represents the instance of the MegatronBertAttention class. It is used to access the attributes and methods of the class.
heads

A list that contains the indices of the attention heads to be pruned. These indices correspond to the specific attention heads that should be removed from the self-attention mechanism.
RETURNS DESCRIPTION
None

However, it modifies the internal state of the MegatronBertAttention instance by pruning the specified attention heads from the self-attention mechanism.
RAISES DESCRIPTION
None

However, potential exceptions that might occur during the execution could include:

  • TypeError: If the input parameters are not of the expected types.
  • IndexError: If there are issues with accessing elements within the 'heads' list or other data structures.
  • ValueError: If there are inconsistencies or unexpected values encountered during the pruning process.
Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
def prune_heads(self, heads):
    """
    This method 'prune_heads' is defined within the 'MegatronBertAttention' class. It prunes specific attention
    heads from the self-attention mechanism based on the provided 'heads' parameter.

    Args:
        self: Represents the instance of the MegatronBertAttention class.
            It is used to access the attributes and methods of the class.
        heads: A list that contains the indices of the attention heads to be pruned.
            These indices correspond to the specific attention heads that should be removed from the self-attention
            mechanism.

    Returns:
        None: However, it modifies the internal state of the MegatronBertAttention instance by pruning the specified
            attention heads from the self-attention mechanism.

    Raises:
        None:
            However, potential exceptions that might occur during the execution could include:

            - TypeError: If the input parameters are not of the expected types.
            - IndexError: If there are issues with accessing elements within the 'heads' list or other data structures.
            - ValueError: If there are inconsistencies or unexpected values encountered during the pruning process.
    """
    if len(heads) == 0:
        return
    heads, index = find_pruneable_heads_and_indices(
        heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads
    )

    # Prune linear layers
    self.self.query = prune_linear_layer(self.self.query, index)
    self.self.key = prune_linear_layer(self.self.key, index)
    self.self.value = prune_linear_layer(self.self.value, index)
    self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)

    # Update hyper params and store pruned heads
    self.self.num_attention_heads = self.self.num_attention_heads - len(heads)
    self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads
    self.pruned_heads = self.pruned_heads.union(heads)

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertEmbeddings

Bases: Module

Construct the embeddings from word, position and token_type embeddings.

Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
 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
class MegatronBertEmbeddings(nn.Module):
    """Construct the embeddings from word, position and token_type embeddings."""
    def __init__(self, config):
        """
        Initialize the MegatronBertEmbeddings class.

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

                - Type: Object
                - Purpose: Contains various configuration parameters such as vocab_size, hidden_size,
                max_position_embeddings, type_vocab_size, pad_token_id, hidden_dropout_prob, and position_embedding_type.
                - Restrictions: Must be a valid configuration object.

        Returns:
            None

        Raises:
            None
        """
        super().__init__()
        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
        self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
        self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)

        # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
        # any TensorFlow checkpoint file

        # In Megatron, layer-norm is applied after the 1st dropout.
        # self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dropout = nn.Dropout(p=config.hidden_dropout_prob)

        # position_ids (1, len position emb) is contiguous in memory and exported when serialized
        self.position_ids = ops.arange(config.max_position_embeddings).broadcast_to((1, -1))
        self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")

    def forward(
        self,
        input_ids: Optional[mindspore.Tensor] = None,
        token_type_ids: Optional[mindspore.Tensor] = None,
        position_ids: Optional[mindspore.Tensor] = None,
        inputs_embeds: Optional[mindspore.Tensor] = None,
        past_key_values_length: int = 0,
    ) -> mindspore.Tensor:
        """
        Construct embeddings for the MegatronBertEmbeddings class.

        Args:
            self: The instance of the MegatronBertEmbeddings class.
            input_ids (Optional[mindspore.Tensor]): The input token IDs. Default is None.
            token_type_ids (Optional[mindspore.Tensor]): The token type IDs. Default is None.
            position_ids (Optional[mindspore.Tensor]): The position IDs. Default is None.
            inputs_embeds (Optional[mindspore.Tensor]): The embedded input tokens. Default is None.
            past_key_values_length (int): The length of past key values. Default is 0.

        Returns:
            mindspore.Tensor: The forwarded embeddings.

        Raises:
            None.
        """
        if input_ids is not None:
            input_shape = input_ids.shape
        else:
            input_shape = inputs_embeds.shape[:-1]

        seq_length = input_shape[1]

        if position_ids is None:
            position_ids = self.position_ids[:, past_key_values_length : seq_length + past_key_values_length]

        if token_type_ids is None:
            token_type_ids = ops.zeros(input_shape, dtype=mindspore.int64)

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

        embeddings = inputs_embeds + token_type_embeddings
        if self.position_embedding_type == "absolute":
            position_embeddings = self.position_embeddings(position_ids)
            embeddings += position_embeddings

        # Megatron BERT moves that layer norm after the drop-out (and to each layer).
        # embeddings = self.LayerNorm(embeddings)
        embeddings = self.dropout(embeddings)
        return embeddings

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertEmbeddings.__init__(config)

Initialize the MegatronBertEmbeddings class.

PARAMETER DESCRIPTION
self

The instance of the class.
config

An object containing configuration parameters for the embeddings.

  • Type: Object
  • Purpose: Contains various configuration parameters such as vocab_size, hidden_size, max_position_embeddings, type_vocab_size, pad_token_id, hidden_dropout_prob, and position_embedding_type.
  • Restrictions: Must be a valid configuration object.
RETURNS DESCRIPTION

None
Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
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
def __init__(self, config):
    """
    Initialize the MegatronBertEmbeddings class.

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

            - Type: Object
            - Purpose: Contains various configuration parameters such as vocab_size, hidden_size,
            max_position_embeddings, type_vocab_size, pad_token_id, hidden_dropout_prob, and position_embedding_type.
            - Restrictions: Must be a valid configuration object.

    Returns:
        None

    Raises:
        None
    """
    super().__init__()
    self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
    self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
    self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)

    # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
    # any TensorFlow checkpoint file

    # In Megatron, layer-norm is applied after the 1st dropout.
    # self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
    self.dropout = nn.Dropout(p=config.hidden_dropout_prob)

    # position_ids (1, len position emb) is contiguous in memory and exported when serialized
    self.position_ids = ops.arange(config.max_position_embeddings).broadcast_to((1, -1))
    self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertEmbeddings.forward(input_ids=None, token_type_ids=None, position_ids=None, inputs_embeds=None, past_key_values_length=0)

Construct embeddings for the MegatronBertEmbeddings class.

PARAMETER DESCRIPTION
self

The instance of the MegatronBertEmbeddings class.
input_ids

The input token IDs. Default is None.

TYPE: Optional[Tensor] DEFAULT: None

token_type_ids

The token type IDs. Default is None.

TYPE: Optional[Tensor] DEFAULT: None

position_ids

The position IDs. Default is None.

TYPE: Optional[Tensor] DEFAULT: None

inputs_embeds

The embedded input tokens. Default is None.

TYPE: Optional[Tensor] DEFAULT: None

past_key_values_length

The length of past key values. Default is 0.

TYPE: int DEFAULT: 0

RETURNS DESCRIPTION
Tensor

mindspore.Tensor: The forwarded embeddings.
Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
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
def forward(
    self,
    input_ids: Optional[mindspore.Tensor] = None,
    token_type_ids: Optional[mindspore.Tensor] = None,
    position_ids: Optional[mindspore.Tensor] = None,
    inputs_embeds: Optional[mindspore.Tensor] = None,
    past_key_values_length: int = 0,
) -> mindspore.Tensor:
    """
    Construct embeddings for the MegatronBertEmbeddings class.

    Args:
        self: The instance of the MegatronBertEmbeddings class.
        input_ids (Optional[mindspore.Tensor]): The input token IDs. Default is None.
        token_type_ids (Optional[mindspore.Tensor]): The token type IDs. Default is None.
        position_ids (Optional[mindspore.Tensor]): The position IDs. Default is None.
        inputs_embeds (Optional[mindspore.Tensor]): The embedded input tokens. Default is None.
        past_key_values_length (int): The length of past key values. Default is 0.

    Returns:
        mindspore.Tensor: The forwarded embeddings.

    Raises:
        None.
    """
    if input_ids is not None:
        input_shape = input_ids.shape
    else:
        input_shape = inputs_embeds.shape[:-1]

    seq_length = input_shape[1]

    if position_ids is None:
        position_ids = self.position_ids[:, past_key_values_length : seq_length + past_key_values_length]

    if token_type_ids is None:
        token_type_ids = ops.zeros(input_shape, dtype=mindspore.int64)

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

    embeddings = inputs_embeds + token_type_embeddings
    if self.position_embedding_type == "absolute":
        position_embeddings = self.position_embeddings(position_ids)
        embeddings += position_embeddings

    # Megatron BERT moves that layer norm after the drop-out (and to each layer).
    # embeddings = self.LayerNorm(embeddings)
    embeddings = self.dropout(embeddings)
    return embeddings

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertEncoder

Bases: Module

The MegatronBertEncoder class represents a transformer encoder for Megatron-BERT. It inherits from nn.Module and is responsible for encoding input sequences using multiple layers of transformer blocks. The class provides methods for forwarding the encoder and performing forward pass computations, including handling gradient checkpointing and caching for efficient training and inference.

ATTRIBUTE DESCRIPTION
config

The configuration parameters for the encoder.
layer

A list of MegatronBertLayer instances representing the stacked transformer layers in the encoder.
ln

A LayerNorm instance for layer normalization.
gradient_checkpointing

A boolean indicating whether gradient checkpointing is enabled.
METHOD DESCRIPTION
__init__

Initializes the MegatronBertEncoder with the provided configuration.
forward

Constructs the encoder and performs forward pass computations, optionally returning hidden states, attentions, and cross-attentions based on the specified parameters.

The forward method handles the processing of input hidden states, attention masks, head masks, encoder hidden states, encoder attention masks, past key values, and caching options. It iterates through the stacked transformer layers, applying gradient checkpointing if enabled, and computes the final hidden states with layer normalization. Additionally, it returns the output as a BaseModelOutputWithPastAndCrossAttentions object if return_dict is True.

Note

The MegatronBertEncoder class is designed for use in the Megatron-BERT architecture and is designed to work in conjunction with other components such as MegatronBertLayer and LayerNorm for efficient transformer-based encoding.

Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
 933
 934
 935
 936
 937
 938
 939
 940
 941
 942
 943
 944
 945
 946
 947
 948
 949
 950
 951
 952
 953
 954
 955
 956
 957
 958
 959
 960
 961
 962
 963
 964
 965
 966
 967
 968
 969
 970
 971
 972
 973
 974
 975
 976
 977
 978
 979
 980
 981
 982
 983
 984
 985
 986
 987
 988
 989
 990
 991
 992
 993
 994
 995
 996
 997
 998
 999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
class MegatronBertEncoder(nn.Module):

    """
    The MegatronBertEncoder class represents a transformer encoder for Megatron-BERT. It inherits from nn.Module and
    is responsible for encoding input sequences using multiple layers of transformer blocks. The class provides methods
    for forwarding the encoder and performing forward pass computations, including handling gradient checkpointing
    and caching for efficient training and inference.

    Attributes:
        config: The configuration parameters for the encoder.
        layer: A list of MegatronBertLayer instances representing the stacked transformer layers in the encoder.
        ln: A LayerNorm instance for layer normalization.
        gradient_checkpointing: A boolean indicating whether gradient checkpointing is enabled.

    Methods:
        __init__: Initializes the MegatronBertEncoder with the provided configuration.
        forward: Constructs the encoder and performs forward pass computations, optionally returning hidden states,
            attentions, and cross-attentions based on the specified parameters.

    The forward method handles the processing of input hidden states, attention masks, head masks, encoder hidden
    states, encoder attention masks, past key values, and caching options. It iterates through the stacked transformer
    layers, applying gradient checkpointing if enabled, and computes the final hidden states with layer normalization.
    Additionally, it returns the output as a BaseModelOutputWithPastAndCrossAttentions object if return_dict is True.

    Note:
        The MegatronBertEncoder class is designed for use in the Megatron-BERT architecture and is designed to work in
        conjunction with other components such as MegatronBertLayer and LayerNorm for efficient transformer-based
        encoding.
    """
    def __init__(self, config):
        """
        Initializes a new instance of the MegatronBertEncoder class.

        Args:
            self: The instance of the MegatronBertEncoder class.
            config (object): An object containing the configuration parameters for the MegatronBertEncoder.
                It should include the following attributes:

                - num_hidden_layers (int): The number of hidden layers in the encoder.
                - hidden_size (int): The size of the hidden layers.
                - layer_norm_eps (float): The epsilon value for layer normalization.

        Returns:
            None.

        Raises:
            None.
        """
        super().__init__()
        self.config = config
        self.layer = nn.ModuleList([MegatronBertLayer(config) for _ in range(config.num_hidden_layers)])

        # The final layer norm. We removed the 1st LN, moved LN to each hidden layer and this one
        # is simply the final LN (Transformer's BERT has it attached to each hidden layer).
        self.ln = nn.LayerNorm([config.hidden_size], eps=config.layer_norm_eps)
        self.gradient_checkpointing = False

    def forward(
        self,
        hidden_states: mindspore.Tensor,
        attention_mask: Optional[mindspore.Tensor] = None,
        head_mask: Optional[mindspore.Tensor] = None,
        encoder_hidden_states: Optional[mindspore.Tensor] = None,
        encoder_attention_mask: Optional[mindspore.Tensor] = None,
        past_key_values: Optional[Tuple[Tuple[mindspore.Tensor]]] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = False,
        output_hidden_states: Optional[bool] = False,
        return_dict: Optional[bool] = True,
    ) -> Union[Tuple, BaseModelOutputWithPastAndCrossAttentions]:
        '''
        Constructs the MegatronBertEncoder.

        Args:
            self (MegatronBertEncoder): The instance of MegatronBertEncoder.
            hidden_states (mindspore.Tensor): The hidden states of the input sequence.
                Shape: (batch_size, sequence_length, hidden_size).
            attention_mask (Optional[mindspore.Tensor]): The attention mask tensor.
                Shape: (batch_size, sequence_length) or (batch_size, sequence_length, sequence_length). Defaults to None.
            head_mask (Optional[mindspore.Tensor]): The head mask tensor.
                Shape: (num_heads,) or (num_layers, num_heads) or (batch_size, num_layers, num_heads) or
                (batch_size, num_heads, sequence_length, sequence_length). Defaults to None.
            encoder_hidden_states (Optional[mindspore.Tensor]): The hidden states of the encoder sequence.
                Shape: (batch_size, encoder_sequence_length, hidden_size). Defaults to None.
            encoder_attention_mask (Optional[mindspore.Tensor]): The attention mask tensor for the encoder.
                Shape: (batch_size, encoder_sequence_length) or (batch_size, encoder_sequence_length,
                encoder_sequence_length). Defaults to None.
            past_key_values (Optional[Tuple[Tuple[mindspore.Tensor]]]): The past key value tensors. Defaults to None.
            use_cache (Optional[bool]): Whether to use cache. Defaults to None.
            output_attentions (Optional[bool]): Whether to output attentions. Defaults to False.
            output_hidden_states (Optional[bool]): Whether to output hidden states. Defaults to False.
            return_dict (Optional[bool]): Whether to return a dictionary as the output. Defaults to True.

        Returns:
            Union[Tuple, BaseModelOutputWithPastAndCrossAttentions]: The output of the MegatronBertEncoder.
                It can be either a tuple of tensors or an instance of BaseModelOutputWithPastAndCrossAttentions.

        Raises:
            None

        '''
        if self.gradient_checkpointing and self.training:
            if use_cache:
                logger.warning_once(
                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
                )
                use_cache = False
        all_hidden_states = () if output_hidden_states else None
        all_self_attentions = () if output_attentions else None
        all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None

        next_decoder_cache = () if use_cache else None
        for i, layer_module in enumerate(self.layer):
            if output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)

            layer_head_mask = head_mask[i] if head_mask is not None else None
            past_key_value = past_key_values[i] if past_key_values is not None else None

            if self.gradient_checkpointing and self.training:
                layer_outputs = self._gradient_checkpointing_func(
                    layer_module.__call__,
                    hidden_states,
                    attention_mask,
                    layer_head_mask,
                    encoder_hidden_states,
                    encoder_attention_mask,
                    past_key_value,
                    output_attentions,
                )
            else:
                layer_outputs = layer_module(
                    hidden_states,
                    attention_mask,
                    layer_head_mask,
                    encoder_hidden_states,
                    encoder_attention_mask,
                    past_key_value,
                    output_attentions,
                )

            # Because we moved the layer-norm at the end of the hidden layer, we have non-normali-
            # zed data here. If that's really needed, we must apply LN to match Transformer's BERT.

            hidden_states = layer_outputs[0]
            if use_cache:
                next_decoder_cache += (layer_outputs[-1],)
            if output_attentions:
                all_self_attentions = all_self_attentions + (layer_outputs[1],)
                if self.config.add_cross_attention:
                    all_cross_attentions = all_cross_attentions + (layer_outputs[2],)

        # Finalize the hidden states.
        hidden_states = self.ln(hidden_states)

        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

        if not return_dict:
            return tuple(
                v
                for v in [
                    hidden_states,
                    next_decoder_cache,
                    all_hidden_states,
                    all_self_attentions,
                    all_cross_attentions,
                ]
                if v is not None
            )
        return BaseModelOutputWithPastAndCrossAttentions(
            last_hidden_state=hidden_states,
            past_key_values=next_decoder_cache,
            hidden_states=all_hidden_states,
            attentions=all_self_attentions,
            cross_attentions=all_cross_attentions,
        )

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertEncoder.__init__(config)

Initializes a new instance of the MegatronBertEncoder class.

PARAMETER DESCRIPTION
self

The instance of the MegatronBertEncoder class.
config

An object containing the configuration parameters for the MegatronBertEncoder. It should include the following attributes:

  • num_hidden_layers (int): The number of hidden layers in the encoder.
  • hidden_size (int): The size of the hidden layers.
  • layer_norm_eps (float): The epsilon value for layer normalization.

TYPE: object

RETURNS DESCRIPTION

None.
Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
def __init__(self, config):
    """
    Initializes a new instance of the MegatronBertEncoder class.

    Args:
        self: The instance of the MegatronBertEncoder class.
        config (object): An object containing the configuration parameters for the MegatronBertEncoder.
            It should include the following attributes:

            - num_hidden_layers (int): The number of hidden layers in the encoder.
            - hidden_size (int): The size of the hidden layers.
            - layer_norm_eps (float): The epsilon value for layer normalization.

    Returns:
        None.

    Raises:
        None.
    """
    super().__init__()
    self.config = config
    self.layer = nn.ModuleList([MegatronBertLayer(config) for _ in range(config.num_hidden_layers)])

    # The final layer norm. We removed the 1st LN, moved LN to each hidden layer and this one
    # is simply the final LN (Transformer's BERT has it attached to each hidden layer).
    self.ln = nn.LayerNorm([config.hidden_size], eps=config.layer_norm_eps)
    self.gradient_checkpointing = False

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertEncoder.forward(hidden_states, attention_mask=None, head_mask=None, encoder_hidden_states=None, encoder_attention_mask=None, past_key_values=None, use_cache=None, output_attentions=False, output_hidden_states=False, return_dict=True)

Constructs the MegatronBertEncoder.

PARAMETER DESCRIPTION
self

The instance of MegatronBertEncoder.

TYPE: MegatronBertEncoder

hidden_states

The hidden states of the input sequence. Shape: (batch_size, sequence_length, hidden_size).

TYPE: Tensor

attention_mask

The attention mask tensor. Shape: (batch_size, sequence_length) or (batch_size, sequence_length, sequence_length). Defaults to None.

TYPE: Optional[Tensor] DEFAULT: None

head_mask

The head mask tensor. Shape: (num_heads,) or (num_layers, num_heads) or (batch_size, num_layers, num_heads) or (batch_size, num_heads, sequence_length, sequence_length). Defaults to None.

TYPE: Optional[Tensor] DEFAULT: None

encoder_hidden_states

The hidden states of the encoder sequence. Shape: (batch_size, encoder_sequence_length, hidden_size). Defaults to None.

TYPE: Optional[Tensor] DEFAULT: None

encoder_attention_mask

The attention mask tensor for the encoder. Shape: (batch_size, encoder_sequence_length) or (batch_size, encoder_sequence_length, encoder_sequence_length). Defaults to None.

TYPE: Optional[Tensor] DEFAULT: None

past_key_values

The past key value tensors. Defaults to None.

TYPE: Optional[Tuple[Tuple[Tensor]]] DEFAULT: None

use_cache

Whether to use cache. Defaults to None.

TYPE: Optional[bool] DEFAULT: None

output_attentions

Whether to output attentions. Defaults to False.

TYPE: Optional[bool] DEFAULT: False

output_hidden_states

Whether to output hidden states. Defaults to False.

TYPE: Optional[bool] DEFAULT: False

return_dict

Whether to return a dictionary as the output. Defaults to True.

TYPE: Optional[bool] DEFAULT: True

RETURNS DESCRIPTION
Union[Tuple, BaseModelOutputWithPastAndCrossAttentions]

Union[Tuple, BaseModelOutputWithPastAndCrossAttentions]: The output of the MegatronBertEncoder. It can be either a tuple of tensors or an instance of BaseModelOutputWithPastAndCrossAttentions.
Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
 990
 991
 992
 993
 994
 995
 996
 997
 998
 999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
def forward(
    self,
    hidden_states: mindspore.Tensor,
    attention_mask: Optional[mindspore.Tensor] = None,
    head_mask: Optional[mindspore.Tensor] = None,
    encoder_hidden_states: Optional[mindspore.Tensor] = None,
    encoder_attention_mask: Optional[mindspore.Tensor] = None,
    past_key_values: Optional[Tuple[Tuple[mindspore.Tensor]]] = None,
    use_cache: Optional[bool] = None,
    output_attentions: Optional[bool] = False,
    output_hidden_states: Optional[bool] = False,
    return_dict: Optional[bool] = True,
) -> Union[Tuple, BaseModelOutputWithPastAndCrossAttentions]:
    '''
    Constructs the MegatronBertEncoder.

    Args:
        self (MegatronBertEncoder): The instance of MegatronBertEncoder.
        hidden_states (mindspore.Tensor): The hidden states of the input sequence.
            Shape: (batch_size, sequence_length, hidden_size).
        attention_mask (Optional[mindspore.Tensor]): The attention mask tensor.
            Shape: (batch_size, sequence_length) or (batch_size, sequence_length, sequence_length). Defaults to None.
        head_mask (Optional[mindspore.Tensor]): The head mask tensor.
            Shape: (num_heads,) or (num_layers, num_heads) or (batch_size, num_layers, num_heads) or
            (batch_size, num_heads, sequence_length, sequence_length). Defaults to None.
        encoder_hidden_states (Optional[mindspore.Tensor]): The hidden states of the encoder sequence.
            Shape: (batch_size, encoder_sequence_length, hidden_size). Defaults to None.
        encoder_attention_mask (Optional[mindspore.Tensor]): The attention mask tensor for the encoder.
            Shape: (batch_size, encoder_sequence_length) or (batch_size, encoder_sequence_length,
            encoder_sequence_length). Defaults to None.
        past_key_values (Optional[Tuple[Tuple[mindspore.Tensor]]]): The past key value tensors. Defaults to None.
        use_cache (Optional[bool]): Whether to use cache. Defaults to None.
        output_attentions (Optional[bool]): Whether to output attentions. Defaults to False.
        output_hidden_states (Optional[bool]): Whether to output hidden states. Defaults to False.
        return_dict (Optional[bool]): Whether to return a dictionary as the output. Defaults to True.

    Returns:
        Union[Tuple, BaseModelOutputWithPastAndCrossAttentions]: The output of the MegatronBertEncoder.
            It can be either a tuple of tensors or an instance of BaseModelOutputWithPastAndCrossAttentions.

    Raises:
        None

    '''
    if self.gradient_checkpointing and self.training:
        if use_cache:
            logger.warning_once(
                "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
            )
            use_cache = False
    all_hidden_states = () if output_hidden_states else None
    all_self_attentions = () if output_attentions else None
    all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None

    next_decoder_cache = () if use_cache else None
    for i, layer_module in enumerate(self.layer):
        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

        layer_head_mask = head_mask[i] if head_mask is not None else None
        past_key_value = past_key_values[i] if past_key_values is not None else None

        if self.gradient_checkpointing and self.training:
            layer_outputs = self._gradient_checkpointing_func(
                layer_module.__call__,
                hidden_states,
                attention_mask,
                layer_head_mask,
                encoder_hidden_states,
                encoder_attention_mask,
                past_key_value,
                output_attentions,
            )
        else:
            layer_outputs = layer_module(
                hidden_states,
                attention_mask,
                layer_head_mask,
                encoder_hidden_states,
                encoder_attention_mask,
                past_key_value,
                output_attentions,
            )

        # Because we moved the layer-norm at the end of the hidden layer, we have non-normali-
        # zed data here. If that's really needed, we must apply LN to match Transformer's BERT.

        hidden_states = layer_outputs[0]
        if use_cache:
            next_decoder_cache += (layer_outputs[-1],)
        if output_attentions:
            all_self_attentions = all_self_attentions + (layer_outputs[1],)
            if self.config.add_cross_attention:
                all_cross_attentions = all_cross_attentions + (layer_outputs[2],)

    # Finalize the hidden states.
    hidden_states = self.ln(hidden_states)

    if output_hidden_states:
        all_hidden_states = all_hidden_states + (hidden_states,)

    if not return_dict:
        return tuple(
            v
            for v in [
                hidden_states,
                next_decoder_cache,
                all_hidden_states,
                all_self_attentions,
                all_cross_attentions,
            ]
            if v is not None
        )
    return BaseModelOutputWithPastAndCrossAttentions(
        last_hidden_state=hidden_states,
        past_key_values=next_decoder_cache,
        hidden_states=all_hidden_states,
        attentions=all_self_attentions,
        cross_attentions=all_cross_attentions,
    )

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertForCausalLM

Bases: MegatronBertPreTrainedModel

A class that represents the MegatronBERT model for Causal Language Modeling. This class inherits from MegatronBertPreTrainedModel and provides methods for model initialization, output embeddings, input preparation for generation, cache reordering, and model forwardion. It also includes detailed explanations for the model's input and output parameters, as well as usage examples. The methods within the class enable fine-tuning and using the model for causal language modeling tasks.

Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
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
2085
2086
2087
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
class MegatronBertForCausalLM(MegatronBertPreTrainedModel):

    '''
    A class that represents the MegatronBERT model for Causal Language Modeling. This class inherits from
    MegatronBertPreTrainedModel and provides methods for model initialization, output embeddings, input
    preparation for generation, cache reordering, and model forwardion. It also includes detailed explanations for
    the model's input and output parameters, as well as usage examples. The methods within the class
    enable fine-tuning and using the model for causal language modeling tasks.
    '''
    _tied_weights_keys = ["cls.predictions.decoder"]

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

        Args:
            self: The instance of MegatronBertForCausalLM class.
            config:
                A configuration object containing settings for the model initialization.

                - Type: object
                - Purpose: To configure the model with specific settings.
                - Restrictions: Must be a valid configuration object compatible with the model.

        Returns:
            None.

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

        if not config.is_decoder:
            logger.warning("If you want to use `MegatronBertForCausalLM` as a standalone, add `is_decoder=True.`")

        self.bert = MegatronBertModel(config, add_pooling_layer=False)
        self.cls = MegatronBertOnlyMLMHead(config)

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

    def get_output_embeddings(self):
        """
        Method to retrieve the output embeddings from MegatronBertForCausalLM model.

        Args:
            self (MegatronBertForCausalLM): The instance of the MegatronBertForCausalLM class.
                It represents the model for which the output embeddings are being retrieved.

        Returns:
            None.

        Raises:
            None.
        """
        return self.cls.predictions.decoder

    def set_output_embeddings(self, new_embeddings):
        """
        Set the output embeddings for the MegatronBertForCausalLM model.

        Args:
            self (MegatronBertForCausalLM): The instance of the MegatronBertForCausalLM class.
            new_embeddings (object): The new output embeddings to be set for the model.
                It could be a tensor, array, or any object representing the new embeddings.

        Returns:
            None.

        Raises:
            None.
        """
        self.cls.predictions.decoder = new_embeddings

    def forward(
        self,
        input_ids: Optional[mindspore.Tensor] = None,
        attention_mask: Optional[mindspore.Tensor] = None,
        token_type_ids: Optional[mindspore.Tensor] = None,
        position_ids: Optional[mindspore.Tensor] = None,
        head_mask: Optional[mindspore.Tensor] = None,
        inputs_embeds: Optional[mindspore.Tensor] = None,
        encoder_hidden_states: Optional[mindspore.Tensor] = None,
        encoder_attention_mask: Optional[mindspore.Tensor] = None,
        labels: Optional[mindspore.Tensor] = None,
        past_key_values: Optional[Tuple[Tuple[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, CausalLMOutputWithCrossAttentions]:
        r"""
        Args:
            encoder_hidden_states  (`mindspore.Tensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
                Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
                the model is configured as a decoder.
            encoder_attention_mask (`mindspore.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
                Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
                the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`:

                - 1 for tokens that are **not masked**,
                - 0 for tokens that are **masked**.
            labels (`mindspore.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
                Labels for computing the left-to-right language modeling loss (next word prediction). Indices should be in
                `[-100, 0, ..., config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are
                ignored (masked), the loss is only computed for the tokens with labels n `[0, ..., config.vocab_size]`
            past_key_values (`tuple(tuple(mindspore.Tensor))` of length `config.n_layers` with each tuple having
                4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
                Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.

                If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
                don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
                `decoder_input_ids` of shape `(batch_size, sequence_length)`.
            use_cache (`bool`, *optional*):
                If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
                `past_key_values`).

        Returns:
            Union[Tuple, CausalLMOutputWithCrossAttentions]

        Example:
            ```python
            >>> from transformers import AutoTokenizer, MegatronBertForCausalLM, MegatronBertConfig
            ...
            >>> tokenizer = AutoTokenizer.from_pretrained("nvidia/megatron-bert-cased-345m")
            >>> model = MegatronBertForCausalLM.from_pretrained("nvidia/megatron-bert-cased-345m", is_decoder=True)
            ...
            >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
            >>> outputs = model(**inputs)
            ...
            >>> prediction_logits = outputs.logits
            ```
        """
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        if labels is not None:
            use_cache = False

        outputs = self.bert(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_attention_mask,
            past_key_values=past_key_values,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        sequence_output = outputs[0]
        prediction_scores = self.cls(sequence_output)

        lm_loss = None
        if labels is not None:
            # we are doing next-token prediction; shift prediction scores and input ids by one
            shifted_prediction_scores = prediction_scores[:, :-1, :]
            labels = labels[:, 1:]
            lm_loss = ops.cross_entropy(shifted_prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))

        if not return_dict:
            output = (prediction_scores,) + outputs[2:]
            return ((lm_loss,) + output) if lm_loss is not None else output

        return CausalLMOutputWithCrossAttentions(
            loss=lm_loss,
            logits=prediction_scores,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
            cross_attentions=outputs.cross_attentions,
        )

    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, attention_mask=None, **model_kwargs):
        """
        Prepare inputs for generation.

        Args:
            self (object): The instance of the class.
            input_ids (tensor): The input tensor containing the token ids.
                Its shape should be (batch_size, sequence_length).
            past_key_values (tuple, optional): The past key values if available for autoregressive generation.
                Default is None.
            attention_mask (tensor, optional): The attention mask tensor.
                If not provided, it is initialized with ones of the same shape as input_ids.

        Returns:
            dict: A dictionary containing the prepared input ids, attention mask, and past key values.

        Raises:
            ValueError: If the input_ids shape is invalid for past_key_values removal.
        """
        input_shape = input_ids.shape
        # if model is used as a decoder in encoder-decoder model, the decoder attention mask is created on the fly
        if attention_mask is None:
            attention_mask = input_ids.new_ones(input_shape)

        # cut decoder_input_ids if past_key_values is used
        if past_key_values is not None:
            past_length = past_key_values[0][0].shape[2]

            # Some generation methods already pass only the last input ID
            if input_ids.shape[1] > past_length:
                remove_prefix_length = past_length
            else:
                # Default to old behavior: keep only final ID
                remove_prefix_length = input_ids.shape[1] - 1

            input_ids = input_ids[:, remove_prefix_length:]

        return {"input_ids": input_ids, "attention_mask": attention_mask, "past_key_values": past_key_values}

    def _reorder_cache(self, past_key_values, beam_idx):
        """
        Method to reorder the cache for a MegatronBertForCausalLM model.

        Args:
            self (object): The instance of the MegatronBertForCausalLM class.
            past_key_values (tuple): A tuple containing the past key values from the model.
            beam_idx (tensor): A tensor representing the indices for reordering the cache.

        Returns:
            None.

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

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertForCausalLM.__init__(config)

Initializes an instance of MegatronBertForCausalLM class.

PARAMETER DESCRIPTION
self

The instance of MegatronBertForCausalLM class.
config

A configuration object containing settings for the model initialization.

  • Type: object
  • Purpose: To configure the model with specific settings.
  • Restrictions: Must be a valid configuration object compatible with the model.
RETURNS DESCRIPTION

None.
Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
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
def __init__(self, config):
    """
    Initializes an instance of MegatronBertForCausalLM class.

    Args:
        self: The instance of MegatronBertForCausalLM class.
        config:
            A configuration object containing settings for the model initialization.

            - Type: object
            - Purpose: To configure the model with specific settings.
            - Restrictions: Must be a valid configuration object compatible with the model.

    Returns:
        None.

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

    if not config.is_decoder:
        logger.warning("If you want to use `MegatronBertForCausalLM` as a standalone, add `is_decoder=True.`")

    self.bert = MegatronBertModel(config, add_pooling_layer=False)
    self.cls = MegatronBertOnlyMLMHead(config)

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

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertForCausalLM.forward(input_ids=None, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None, inputs_embeds=None, encoder_hidden_states=None, encoder_attention_mask=None, labels=None, past_key_values=None, use_cache=None, output_attentions=None, output_hidden_states=None, return_dict=None)

PARAMETER DESCRIPTION
encoder_hidden_states

Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if the model is configured as a decoder.

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

encoder_attention_mask

Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in the cross-attention if the model is configured as a decoder. Mask values selected in [0, 1]:

  • 1 for tokens that are not masked,
  • 0 for tokens that are masked.

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

labels

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

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

use_cache

If set to True, past_key_values key value states are returned and can be used to speed up decoding (see past_key_values).

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

RETURNS DESCRIPTION
Union[Tuple, CausalLMOutputWithCrossAttentions]

Union[Tuple, CausalLMOutputWithCrossAttentions]
Example 
>>> from transformers import AutoTokenizer, MegatronBertForCausalLM, MegatronBertConfig
...
>>> tokenizer = AutoTokenizer.from_pretrained("nvidia/megatron-bert-cased-345m")
>>> model = MegatronBertForCausalLM.from_pretrained("nvidia/megatron-bert-cased-345m", is_decoder=True)
...
>>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
>>> outputs = model(**inputs)
...
>>> prediction_logits = outputs.logits
Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
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
2085
2086
2087
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
def forward(
    self,
    input_ids: Optional[mindspore.Tensor] = None,
    attention_mask: Optional[mindspore.Tensor] = None,
    token_type_ids: Optional[mindspore.Tensor] = None,
    position_ids: Optional[mindspore.Tensor] = None,
    head_mask: Optional[mindspore.Tensor] = None,
    inputs_embeds: Optional[mindspore.Tensor] = None,
    encoder_hidden_states: Optional[mindspore.Tensor] = None,
    encoder_attention_mask: Optional[mindspore.Tensor] = None,
    labels: Optional[mindspore.Tensor] = None,
    past_key_values: Optional[Tuple[Tuple[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, CausalLMOutputWithCrossAttentions]:
    r"""
    Args:
        encoder_hidden_states  (`mindspore.Tensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
            Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
            the model is configured as a decoder.
        encoder_attention_mask (`mindspore.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
            Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
            the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`:

            - 1 for tokens that are **not masked**,
            - 0 for tokens that are **masked**.
        labels (`mindspore.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
            Labels for computing the left-to-right language modeling loss (next word prediction). Indices should be in
            `[-100, 0, ..., config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are
            ignored (masked), the loss is only computed for the tokens with labels n `[0, ..., config.vocab_size]`
        past_key_values (`tuple(tuple(mindspore.Tensor))` of length `config.n_layers` with each tuple having
            4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
            Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.

            If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
            don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
            `decoder_input_ids` of shape `(batch_size, sequence_length)`.
        use_cache (`bool`, *optional*):
            If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
            `past_key_values`).

    Returns:
        Union[Tuple, CausalLMOutputWithCrossAttentions]

    Example:
        ```python
        >>> from transformers import AutoTokenizer, MegatronBertForCausalLM, MegatronBertConfig
        ...
        >>> tokenizer = AutoTokenizer.from_pretrained("nvidia/megatron-bert-cased-345m")
        >>> model = MegatronBertForCausalLM.from_pretrained("nvidia/megatron-bert-cased-345m", is_decoder=True)
        ...
        >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
        >>> outputs = model(**inputs)
        ...
        >>> prediction_logits = outputs.logits
        ```
    """
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    if labels is not None:
        use_cache = False

    outputs = self.bert(
        input_ids,
        attention_mask=attention_mask,
        token_type_ids=token_type_ids,
        position_ids=position_ids,
        head_mask=head_mask,
        inputs_embeds=inputs_embeds,
        encoder_hidden_states=encoder_hidden_states,
        encoder_attention_mask=encoder_attention_mask,
        past_key_values=past_key_values,
        use_cache=use_cache,
        output_attentions=output_attentions,
        output_hidden_states=output_hidden_states,
        return_dict=return_dict,
    )

    sequence_output = outputs[0]
    prediction_scores = self.cls(sequence_output)

    lm_loss = None
    if labels is not None:
        # we are doing next-token prediction; shift prediction scores and input ids by one
        shifted_prediction_scores = prediction_scores[:, :-1, :]
        labels = labels[:, 1:]
        lm_loss = ops.cross_entropy(shifted_prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))

    if not return_dict:
        output = (prediction_scores,) + outputs[2:]
        return ((lm_loss,) + output) if lm_loss is not None else output

    return CausalLMOutputWithCrossAttentions(
        loss=lm_loss,
        logits=prediction_scores,
        past_key_values=outputs.past_key_values,
        hidden_states=outputs.hidden_states,
        attentions=outputs.attentions,
        cross_attentions=outputs.cross_attentions,
    )

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertForCausalLM.get_output_embeddings()

Method to retrieve the output embeddings from MegatronBertForCausalLM model.

PARAMETER DESCRIPTION
self

The instance of the MegatronBertForCausalLM class. It represents the model for which the output embeddings are being retrieved.

TYPE: MegatronBertForCausalLM

RETURNS DESCRIPTION

None.
Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
1992
1993
1994
1995
1996
1997
1998
1999
2000
2001
2002
2003
2004
2005
2006
def get_output_embeddings(self):
    """
    Method to retrieve the output embeddings from MegatronBertForCausalLM model.

    Args:
        self (MegatronBertForCausalLM): The instance of the MegatronBertForCausalLM class.
            It represents the model for which the output embeddings are being retrieved.

    Returns:
        None.

    Raises:
        None.
    """
    return self.cls.predictions.decoder

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertForCausalLM.prepare_inputs_for_generation(input_ids, past_key_values=None, attention_mask=None, **model_kwargs)

Prepare inputs for generation.

PARAMETER DESCRIPTION
self

The instance of the class.

TYPE: object

input_ids

The input tensor containing the token ids. Its shape should be (batch_size, sequence_length).

TYPE: tensor

past_key_values

The past key values if available for autoregressive generation. Default is None.

TYPE: tuple DEFAULT: None

attention_mask

The attention mask tensor. If not provided, it is initialized with ones of the same shape as input_ids.

TYPE: tensor DEFAULT: None

RETURNS DESCRIPTION
dict

A dictionary containing the prepared input ids, attention mask, and past key values.
RAISES DESCRIPTION
ValueError

If the input_ids shape is invalid for past_key_values removal.
Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
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
def prepare_inputs_for_generation(self, input_ids, past_key_values=None, attention_mask=None, **model_kwargs):
    """
    Prepare inputs for generation.

    Args:
        self (object): The instance of the class.
        input_ids (tensor): The input tensor containing the token ids.
            Its shape should be (batch_size, sequence_length).
        past_key_values (tuple, optional): The past key values if available for autoregressive generation.
            Default is None.
        attention_mask (tensor, optional): The attention mask tensor.
            If not provided, it is initialized with ones of the same shape as input_ids.

    Returns:
        dict: A dictionary containing the prepared input ids, attention mask, and past key values.

    Raises:
        ValueError: If the input_ids shape is invalid for past_key_values removal.
    """
    input_shape = input_ids.shape
    # if model is used as a decoder in encoder-decoder model, the decoder attention mask is created on the fly
    if attention_mask is None:
        attention_mask = input_ids.new_ones(input_shape)

    # cut decoder_input_ids if past_key_values is used
    if past_key_values is not None:
        past_length = past_key_values[0][0].shape[2]

        # Some generation methods already pass only the last input ID
        if input_ids.shape[1] > past_length:
            remove_prefix_length = past_length
        else:
            # Default to old behavior: keep only final ID
            remove_prefix_length = input_ids.shape[1] - 1

        input_ids = input_ids[:, remove_prefix_length:]

    return {"input_ids": input_ids, "attention_mask": attention_mask, "past_key_values": past_key_values}

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertForCausalLM.set_output_embeddings(new_embeddings)

Set the output embeddings for the MegatronBertForCausalLM model.

PARAMETER DESCRIPTION
self

The instance of the MegatronBertForCausalLM class.

TYPE: MegatronBertForCausalLM

new_embeddings

The new output embeddings to be set for the model. It could be a tensor, array, or any object representing the new embeddings.

TYPE: object

RETURNS DESCRIPTION

None.
Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
2008
2009
2010
2011
2012
2013
2014
2015
2016
2017
2018
2019
2020
2021
2022
2023
def set_output_embeddings(self, new_embeddings):
    """
    Set the output embeddings for the MegatronBertForCausalLM model.

    Args:
        self (MegatronBertForCausalLM): The instance of the MegatronBertForCausalLM class.
        new_embeddings (object): The new output embeddings to be set for the model.
            It could be a tensor, array, or any object representing the new embeddings.

    Returns:
        None.

    Raises:
        None.
    """
    self.cls.predictions.decoder = new_embeddings

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertForMaskedLM

Bases: MegatronBertPreTrainedModel

This class represents a MegatronBert model for Masked Language Modeling (MLM). It inherits from the MegatronBertPreTrainedModel and includes methods for initializing the model, getting and setting output embeddings, forwarding the model, and preparing inputs for generation. The class provides functionality for performing masked language modeling tasks using the MegatronBert model.

ATTRIBUTE DESCRIPTION
config

The configuration for the MegatronBert model.

TYPE: MegatronBertConfig

METHOD DESCRIPTION
__init__

Initializes the MegatronBertForMaskedLM model with the given configuration.
get_output_embeddings

Retrieves the output embeddings of the model.
set_output_embeddings

Sets the output embeddings of the model to the specified new embeddings.
forward

Constructs the model with the given input and optional arguments, and returns the MaskedLMOutput.
prepare_inputs_for_generation

Prepares the input for generation by updating the input_ids and attention_mask for the model.
Note

For consistency, always use triple double quotes around docstrings.

Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
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
2250
2251
2252
2253
2254
2255
2256
2257
2258
2259
2260
2261
2262
2263
2264
2265
2266
2267
2268
2269
2270
2271
2272
2273
2274
2275
2276
2277
2278
2279
2280
2281
2282
2283
2284
2285
2286
2287
2288
2289
2290
2291
2292
2293
2294
2295
2296
2297
2298
2299
2300
2301
2302
2303
2304
2305
2306
2307
2308
2309
2310
2311
2312
2313
2314
2315
2316
2317
2318
2319
2320
2321
2322
2323
2324
2325
2326
2327
2328
2329
2330
2331
2332
2333
2334
2335
2336
2337
2338
2339
2340
2341
2342
2343
2344
2345
2346
2347
2348
2349
2350
2351
2352
2353
2354
2355
2356
2357
2358
2359
2360
2361
2362
2363
2364
2365
2366
2367
2368
2369
2370
class MegatronBertForMaskedLM(MegatronBertPreTrainedModel):

    """
    This class represents a MegatronBert model for Masked Language Modeling (MLM). It inherits from the
    MegatronBertPreTrainedModel and includes methods for initializing the model, getting and setting output
    embeddings, forwarding the model, and preparing inputs for generation. The class provides functionality
    for performing masked language modeling tasks using the MegatronBert model.

    Attributes:
        config (MegatronBertConfig): The configuration for the MegatronBert model.

    Methods:
        __init__: Initializes the MegatronBertForMaskedLM model with the given configuration.
        get_output_embeddings: Retrieves the output embeddings of the model.
        set_output_embeddings: Sets the output embeddings of the model to the specified new embeddings.
        forward: Constructs the model with the given input and optional arguments, and returns the MaskedLMOutput.
        prepare_inputs_for_generation: Prepares the input for generation by updating the input_ids and attention_mask
            for the model.

    Note:
        For consistency, always use triple double quotes around docstrings.
    """
    _tied_weights_keys = ["cls.predictions.decoder"]

    def __init__(self, config):
        """
        Initializes an instance of MegatronBertForMaskedLM.

        Args:
            self: The instance of the class.
            config: A configuration object containing settings for the MegatronBertForMaskedLM model.
                It must have attributes like 'is_decoder', which is a boolean indicating if the model is a decoder.
                The configuration object is used to configure the model's behavior.

        Returns:
            None.

        Raises:
            Warning: If the 'is_decoder' attribute in the config is set to True, a warning message is logged.
            AttributeError: If the config object does not have the required attributes, an AttributeError may be raised.
        """
        super().__init__(config)

        if config.is_decoder:
            logger.warning(
                "If you want to use `MegatronBertForMaskedLM` make sure `config.is_decoder=False` for "
                "bi-directional self-attention."
            )

        self.bert = MegatronBertModel(config, add_pooling_layer=False)
        self.cls = MegatronBertOnlyMLMHead(config)

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

    def get_output_embeddings(self):
        """
        Returns the output embeddings of the MegatronBertForMaskedLM model.

        Args:
            self (MegatronBertForMaskedLM): An instance of the MegatronBertForMaskedLM class.

        Returns:
            None.

        Raises:
            None.
        """
        return self.cls.predictions.decoder

    def set_output_embeddings(self, new_embeddings):
        """
        Sets the output embeddings for the MegatronBertForMaskedLM model.

        Args:
            self (MegatronBertForMaskedLM): An instance of the MegatronBertForMaskedLM class.
            new_embeddings: The new embeddings to be set for the model's output.

        Returns:
            None: This method modifies the model in-place.

        Raises:
            None.

        This method is used to set the output embeddings for the MegatronBertForMaskedLM model. The new embeddings are
        assigned to the model's predictions.decoder attribute, which represents the decoder layer responsible for
        generating output embeddings during inference. The method does not return any value and modifies the model
        directly.
        """
        self.cls.predictions.decoder = new_embeddings

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

        outputs = self.bert(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_attention_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        sequence_output = outputs[0]
        prediction_scores = self.cls(sequence_output)

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

        if not return_dict:
            output = (prediction_scores,) + outputs[2:]
            return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output

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

    def prepare_inputs_for_generation(self, input_ids, attention_mask=None, **model_kwargs):
        """
        Prepare inputs for generation.

        This method prepares input tensors for generation in the MegatronBertForMaskedLM model.

        Args:
            self: (object) The instance of the MegatronBertForMaskedLM class.
            input_ids: (Tensor) The input token IDs. Shape [batch_size, sequence_length].
            attention_mask: (Tensor, optional) The attention mask tensor. Shape [batch_size, sequence_length].

        Returns:
            dict:
                A dictionary containing the prepared input tensors for generation:

                - 'input_ids': (Tensor) The prepared input token IDs with dummy token appended.
                Shape [batch_size, sequence_length + 1].
                - 'attention_mask': (Tensor) The prepared attention mask tensor with an additional column of zeros appended.
                Shape [batch_size, sequence_length + 1].

        Raises:
            ValueError: If the PAD token is not defined for generation.
        """
        input_shape = input_ids.shape
        effective_batch_size = input_shape[0]

        #  add a dummy token
        if self.config.pad_token_id is None:
            raise ValueError("The PAD token should be defined for generation")
        attention_mask = ops.cat([attention_mask, attention_mask.new_zeros((attention_mask.shape[0], 1))], axis=-1)
        dummy_token = ops.full(
            (effective_batch_size, 1), self.config.pad_token_id, dtype=mindspore.int64)
        input_ids = ops.cat([input_ids, dummy_token], axis=1)

        return {"input_ids": input_ids, "attention_mask": attention_mask}

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertForMaskedLM.__init__(config)

Initializes an instance of MegatronBertForMaskedLM.

PARAMETER DESCRIPTION
self

The instance of the class.
config

A configuration object containing settings for the MegatronBertForMaskedLM model. It must have attributes like 'is_decoder', which is a boolean indicating if the model is a decoder. The configuration object is used to configure the model's behavior.
RETURNS DESCRIPTION

None.
RAISES DESCRIPTION
Warning

If the 'is_decoder' attribute in the config is set to True, a warning message is logged.
AttributeError

If the config object does not have the required attributes, an AttributeError may be raised.
Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
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
def __init__(self, config):
    """
    Initializes an instance of MegatronBertForMaskedLM.

    Args:
        self: The instance of the class.
        config: A configuration object containing settings for the MegatronBertForMaskedLM model.
            It must have attributes like 'is_decoder', which is a boolean indicating if the model is a decoder.
            The configuration object is used to configure the model's behavior.

    Returns:
        None.

    Raises:
        Warning: If the 'is_decoder' attribute in the config is set to True, a warning message is logged.
        AttributeError: If the config object does not have the required attributes, an AttributeError may be raised.
    """
    super().__init__(config)

    if config.is_decoder:
        logger.warning(
            "If you want to use `MegatronBertForMaskedLM` make sure `config.is_decoder=False` for "
            "bi-directional self-attention."
        )

    self.bert = MegatronBertModel(config, add_pooling_layer=False)
    self.cls = MegatronBertOnlyMLMHead(config)

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

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertForMaskedLM.forward(input_ids=None, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None, inputs_embeds=None, encoder_hidden_states=None, encoder_attention_mask=None, labels=None, output_attentions=None, output_hidden_states=None, return_dict=None)

PARAMETER DESCRIPTION
labels

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

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

Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
2280
2281
2282
2283
2284
2285
2286
2287
2288
2289
2290
2291
2292
2293
2294
2295
2296
2297
2298
2299
2300
2301
2302
2303
2304
2305
2306
2307
2308
2309
2310
2311
2312
2313
2314
2315
2316
2317
2318
2319
2320
2321
2322
2323
2324
2325
2326
2327
2328
2329
2330
2331
2332
2333
2334
def forward(
    self,
    input_ids: Optional[mindspore.Tensor] = None,
    attention_mask: Optional[mindspore.Tensor] = None,
    token_type_ids: Optional[mindspore.Tensor] = None,
    position_ids: Optional[mindspore.Tensor] = None,
    head_mask: Optional[mindspore.Tensor] = None,
    inputs_embeds: Optional[mindspore.Tensor] = None,
    encoder_hidden_states: Optional[mindspore.Tensor] = None,
    encoder_attention_mask: Optional[mindspore.Tensor] = None,
    labels: Optional[mindspore.Tensor] = None,
    output_attentions: Optional[bool] = None,
    output_hidden_states: Optional[bool] = None,
    return_dict: Optional[bool] = None,
) -> Union[Tuple, MaskedLMOutput]:
    r"""
    Args:
        labels (`mindspore.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
            Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
            config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the
            loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`
    """
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict

    outputs = self.bert(
        input_ids,
        attention_mask=attention_mask,
        token_type_ids=token_type_ids,
        position_ids=position_ids,
        head_mask=head_mask,
        inputs_embeds=inputs_embeds,
        encoder_hidden_states=encoder_hidden_states,
        encoder_attention_mask=encoder_attention_mask,
        output_attentions=output_attentions,
        output_hidden_states=output_hidden_states,
        return_dict=return_dict,
    )

    sequence_output = outputs[0]
    prediction_scores = self.cls(sequence_output)

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

    if not return_dict:
        output = (prediction_scores,) + outputs[2:]
        return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output

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

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertForMaskedLM.get_output_embeddings()

Returns the output embeddings of the MegatronBertForMaskedLM model.

PARAMETER DESCRIPTION
self

An instance of the MegatronBertForMaskedLM class.

TYPE: MegatronBertForMaskedLM

RETURNS DESCRIPTION

None.
Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
2244
2245
2246
2247
2248
2249
2250
2251
2252
2253
2254
2255
2256
2257
def get_output_embeddings(self):
    """
    Returns the output embeddings of the MegatronBertForMaskedLM model.

    Args:
        self (MegatronBertForMaskedLM): An instance of the MegatronBertForMaskedLM class.

    Returns:
        None.

    Raises:
        None.
    """
    return self.cls.predictions.decoder

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertForMaskedLM.prepare_inputs_for_generation(input_ids, attention_mask=None, **model_kwargs)

Prepare inputs for generation.

This method prepares input tensors for generation in the MegatronBertForMaskedLM model.

PARAMETER DESCRIPTION
self

(object) The instance of the MegatronBertForMaskedLM class.
input_ids

(Tensor) The input token IDs. Shape [batch_size, sequence_length].
attention_mask

(Tensor, optional) The attention mask tensor. Shape [batch_size, sequence_length].

DEFAULT: None

RETURNS DESCRIPTION
dict

A dictionary containing the prepared input tensors for generation:

  • 'input_ids': (Tensor) The prepared input token IDs with dummy token appended. Shape [batch_size, sequence_length + 1].
  • 'attention_mask': (Tensor) The prepared attention mask tensor with an additional column of zeros appended. Shape [batch_size, sequence_length + 1].
RAISES DESCRIPTION
ValueError

If the PAD token is not defined for generation.
Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
2336
2337
2338
2339
2340
2341
2342
2343
2344
2345
2346
2347
2348
2349
2350
2351
2352
2353
2354
2355
2356
2357
2358
2359
2360
2361
2362
2363
2364
2365
2366
2367
2368
2369
2370
def prepare_inputs_for_generation(self, input_ids, attention_mask=None, **model_kwargs):
    """
    Prepare inputs for generation.

    This method prepares input tensors for generation in the MegatronBertForMaskedLM model.

    Args:
        self: (object) The instance of the MegatronBertForMaskedLM class.
        input_ids: (Tensor) The input token IDs. Shape [batch_size, sequence_length].
        attention_mask: (Tensor, optional) The attention mask tensor. Shape [batch_size, sequence_length].

    Returns:
        dict:
            A dictionary containing the prepared input tensors for generation:

            - 'input_ids': (Tensor) The prepared input token IDs with dummy token appended.
            Shape [batch_size, sequence_length + 1].
            - 'attention_mask': (Tensor) The prepared attention mask tensor with an additional column of zeros appended.
            Shape [batch_size, sequence_length + 1].

    Raises:
        ValueError: If the PAD token is not defined for generation.
    """
    input_shape = input_ids.shape
    effective_batch_size = input_shape[0]

    #  add a dummy token
    if self.config.pad_token_id is None:
        raise ValueError("The PAD token should be defined for generation")
    attention_mask = ops.cat([attention_mask, attention_mask.new_zeros((attention_mask.shape[0], 1))], axis=-1)
    dummy_token = ops.full(
        (effective_batch_size, 1), self.config.pad_token_id, dtype=mindspore.int64)
    input_ids = ops.cat([input_ids, dummy_token], axis=1)

    return {"input_ids": input_ids, "attention_mask": attention_mask}

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertForMaskedLM.set_output_embeddings(new_embeddings)

Sets the output embeddings for the MegatronBertForMaskedLM model.

PARAMETER DESCRIPTION
self

An instance of the MegatronBertForMaskedLM class.

TYPE: MegatronBertForMaskedLM

new_embeddings

The new embeddings to be set for the model's output.
RETURNS DESCRIPTION
None

This method modifies the model in-place.

This method is used to set the output embeddings for the MegatronBertForMaskedLM model. The new embeddings are assigned to the model's predictions.decoder attribute, which represents the decoder layer responsible for generating output embeddings during inference. The method does not return any value and modifies the model directly.

Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
2259
2260
2261
2262
2263
2264
2265
2266
2267
2268
2269
2270
2271
2272
2273
2274
2275
2276
2277
2278
def set_output_embeddings(self, new_embeddings):
    """
    Sets the output embeddings for the MegatronBertForMaskedLM model.

    Args:
        self (MegatronBertForMaskedLM): An instance of the MegatronBertForMaskedLM class.
        new_embeddings: The new embeddings to be set for the model's output.

    Returns:
        None: This method modifies the model in-place.

    Raises:
        None.

    This method is used to set the output embeddings for the MegatronBertForMaskedLM model. The new embeddings are
    assigned to the model's predictions.decoder attribute, which represents the decoder layer responsible for
    generating output embeddings during inference. The method does not return any value and modifies the model
    directly.
    """
    self.cls.predictions.decoder = new_embeddings

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertForMultipleChoice

Bases: MegatronBertPreTrainedModel

A Python class representing the MegatronBertForMultipleChoice model, which is designed for multiple choice classification tasks. It is a subclass of the MegatronBertPreTrainedModel.

The MegatronBertForMultipleChoice model consists of a MegatronBertModel, a dropout layer, and a classifier. The MegatronBertModel encodes the input sequence using the BERT architecture, while the dropout layer helps prevent overfitting. The classifier then produces logits for each choice in the multiple choice question.

METHOD DESCRIPTION
__init__

Initializes the MegatronBertForMultipleChoice model with the given configuration.
forward

Constructs the model and performs forward pass given the input tensors. It returns the logits for each choice and optionally computes the loss.
ATTRIBUTE DESCRIPTION
bert

The MegatronBertModel used for encoding the input sequence.
dropout

The dropout layer for regularization.
classifier

The linear layer for producing logits.
Note
  • The input tensors should be either mindspore.Tensor objects or None if not applicable.
  • The labels tensor should have shape (batch_size,) and contain indices in [0, ..., num_choices-1].
  • The return_dict argument is optional and defaults to the use_return_dict value from the model configuration.
Example 
>>> config = MegatronBertConfig(...)
>>> model = MegatronBertForMultipleChoice(config)
>>> input_ids = ...
>>> attention_mask = ...
>>> token_type_ids = ...
>>> position_ids = ...
>>> head_mask = ...
>>> inputs_embeds = ...
>>> labels = ...
>>> output_attentions = ...
>>> output_hidden_states = ...
>>> return_dict = ...
>>> logits, loss = model.forward(input_ids, attention_mask, token_type_ids, position_ids, head_mask,
... inputs_embeds, labels, output_attentions, output_hidden_states, return_dict)
Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
2647
2648
2649
2650
2651
2652
2653
2654
2655
2656
2657
2658
2659
2660
2661
2662
2663
2664
2665
2666
2667
2668
2669
2670
2671
2672
2673
2674
2675
2676
2677
2678
2679
2680
2681
2682
2683
2684
2685
2686
2687
2688
2689
2690
2691
2692
2693
2694
2695
2696
2697
2698
2699
2700
2701
2702
2703
2704
2705
2706
2707
2708
2709
2710
2711
2712
2713
2714
2715
2716
2717
2718
2719
2720
2721
2722
2723
2724
2725
2726
2727
2728
2729
2730
2731
2732
2733
2734
2735
2736
2737
2738
2739
2740
2741
2742
2743
2744
2745
2746
2747
2748
2749
2750
2751
2752
2753
2754
2755
2756
2757
2758
2759
2760
2761
2762
2763
2764
2765
2766
2767
2768
2769
2770
2771
2772
2773
2774
2775
2776
2777
2778
2779
class MegatronBertForMultipleChoice(MegatronBertPreTrainedModel):

    """
    A Python class representing the MegatronBertForMultipleChoice model, which is designed for multiple choice
    classification tasks. It is a subclass of the MegatronBertPreTrainedModel.

    The MegatronBertForMultipleChoice model consists of a MegatronBertModel, a dropout layer, and a classifier.
    The MegatronBertModel encodes the input sequence using the BERT architecture, while the dropout layer helps prevent
    overfitting. The classifier then produces logits for each choice in the multiple choice question.

    Methods:
        __init__: Initializes the MegatronBertForMultipleChoice model with the given configuration.
        forward: Constructs the model and performs forward pass given the input tensors. It returns the logits for
            each choice and optionally computes the loss.

    Attributes:
        bert: The MegatronBertModel used for encoding the input sequence.
        dropout: The dropout layer for regularization.
        classifier: The linear layer for producing logits.

    Note:
        - The input tensors should be either `mindspore.Tensor` objects or `None` if not applicable.
        - The `labels` tensor should have shape `(batch_size,)` and contain indices in `[0, ..., num_choices-1]`.
        - The `return_dict` argument is optional and defaults to the `use_return_dict` value from the model configuration.

    Example:
        ```python
        >>> config = MegatronBertConfig(...)
        >>> model = MegatronBertForMultipleChoice(config)
        >>> input_ids = ...
        >>> attention_mask = ...
        >>> token_type_ids = ...
        >>> position_ids = ...
        >>> head_mask = ...
        >>> inputs_embeds = ...
        >>> labels = ...
        >>> output_attentions = ...
        >>> output_hidden_states = ...
        >>> return_dict = ...
        >>> logits, loss = model.forward(input_ids, attention_mask, token_type_ids, position_ids, head_mask,
        ... inputs_embeds, labels, output_attentions, output_hidden_states, return_dict)
        ```
    """
    def __init__(self, config):
        """
        Initializes an instance of the MegatronBertForMultipleChoice class.

        Args:
            self (object): The instance of the class itself.
            config (object): The configuration object containing parameters for model initialization.
                It should have attributes like hidden_dropout_prob, hidden_size, etc.
                This parameter is required for configuring the model.

        Returns:
            None.

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

        self.bert = MegatronBertModel(config)
        self.dropout = nn.Dropout(p=config.hidden_dropout_prob)
        self.classifier = nn.Linear(config.hidden_size, 1)

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

    def forward(
        self,
        input_ids: Optional[mindspore.Tensor] = None,
        attention_mask: Optional[mindspore.Tensor] = None,
        token_type_ids: Optional[mindspore.Tensor] = None,
        position_ids: Optional[mindspore.Tensor] = None,
        head_mask: Optional[mindspore.Tensor] = None,
        inputs_embeds: Optional[mindspore.Tensor] = None,
        labels: Optional[mindspore.Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, MultipleChoiceModelOutput]:
        r"""
        Args:
            labels (`mindspore.Tensor` of shape `(batch_size,)`, *optional*):
                Labels for computing the multiple choice classification loss. Indices should be in `[0, ...,
                num_choices-1]` where `num_choices` is the size of the second dimension of the input tensors. (See
                `input_ids` above)
        """
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        num_choices = input_ids.shape[1] if input_ids is not None else inputs_embeds.shape[1]

        input_ids = input_ids.view(-1, input_ids.shape[-1]) if input_ids is not None else None
        attention_mask = attention_mask.view(-1, attention_mask.shape[-1]) if attention_mask is not None else None
        token_type_ids = token_type_ids.view(-1, token_type_ids.shape[-1]) if token_type_ids is not None else None
        position_ids = position_ids.view(-1, position_ids.shape[-1]) if position_ids is not None else None
        inputs_embeds = (
            inputs_embeds.view(-1, inputs_embeds.shape[-2], inputs_embeds.shape[-1])
            if inputs_embeds is not None
            else None
        )

        outputs = self.bert(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        pooled_output = outputs[1]

        pooled_output = self.dropout(pooled_output)
        logits = self.classifier(pooled_output)
        reshaped_logits = logits.view(-1, num_choices)

        loss = None
        if labels is not None:
            loss = ops.cross_entropy(reshaped_logits, labels)

        if not return_dict:
            output = (reshaped_logits,) + outputs[2:]
            return ((loss,) + output) if loss is not None else output

        return MultipleChoiceModelOutput(
            loss=loss,
            logits=reshaped_logits,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertForMultipleChoice.__init__(config)

Initializes an instance of the MegatronBertForMultipleChoice class.

PARAMETER DESCRIPTION
self

The instance of the class itself.

TYPE: object

config

The configuration object containing parameters for model initialization. It should have attributes like hidden_dropout_prob, hidden_size, etc. This parameter is required for configuring the model.

TYPE: object

RETURNS DESCRIPTION

None.
Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
2690
2691
2692
2693
2694
2695
2696
2697
2698
2699
2700
2701
2702
2703
2704
2705
2706
2707
2708
2709
2710
2711
2712
2713
def __init__(self, config):
    """
    Initializes an instance of the MegatronBertForMultipleChoice class.

    Args:
        self (object): The instance of the class itself.
        config (object): The configuration object containing parameters for model initialization.
            It should have attributes like hidden_dropout_prob, hidden_size, etc.
            This parameter is required for configuring the model.

    Returns:
        None.

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

    self.bert = MegatronBertModel(config)
    self.dropout = nn.Dropout(p=config.hidden_dropout_prob)
    self.classifier = nn.Linear(config.hidden_size, 1)

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

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertForMultipleChoice.forward(input_ids=None, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None, inputs_embeds=None, labels=None, output_attentions=None, output_hidden_states=None, return_dict=None)

PARAMETER DESCRIPTION
labels

Labels for computing the multiple choice classification loss. Indices should be in [0, ..., num_choices-1] where num_choices is the size of the second dimension of the input tensors. (See input_ids above)

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

Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
2715
2716
2717
2718
2719
2720
2721
2722
2723
2724
2725
2726
2727
2728
2729
2730
2731
2732
2733
2734
2735
2736
2737
2738
2739
2740
2741
2742
2743
2744
2745
2746
2747
2748
2749
2750
2751
2752
2753
2754
2755
2756
2757
2758
2759
2760
2761
2762
2763
2764
2765
2766
2767
2768
2769
2770
2771
2772
2773
2774
2775
2776
2777
2778
2779
def forward(
    self,
    input_ids: Optional[mindspore.Tensor] = None,
    attention_mask: Optional[mindspore.Tensor] = None,
    token_type_ids: Optional[mindspore.Tensor] = None,
    position_ids: Optional[mindspore.Tensor] = None,
    head_mask: Optional[mindspore.Tensor] = None,
    inputs_embeds: Optional[mindspore.Tensor] = None,
    labels: Optional[mindspore.Tensor] = None,
    output_attentions: Optional[bool] = None,
    output_hidden_states: Optional[bool] = None,
    return_dict: Optional[bool] = None,
) -> Union[Tuple, MultipleChoiceModelOutput]:
    r"""
    Args:
        labels (`mindspore.Tensor` of shape `(batch_size,)`, *optional*):
            Labels for computing the multiple choice classification loss. Indices should be in `[0, ...,
            num_choices-1]` where `num_choices` is the size of the second dimension of the input tensors. (See
            `input_ids` above)
    """
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    num_choices = input_ids.shape[1] if input_ids is not None else inputs_embeds.shape[1]

    input_ids = input_ids.view(-1, input_ids.shape[-1]) if input_ids is not None else None
    attention_mask = attention_mask.view(-1, attention_mask.shape[-1]) if attention_mask is not None else None
    token_type_ids = token_type_ids.view(-1, token_type_ids.shape[-1]) if token_type_ids is not None else None
    position_ids = position_ids.view(-1, position_ids.shape[-1]) if position_ids is not None else None
    inputs_embeds = (
        inputs_embeds.view(-1, inputs_embeds.shape[-2], inputs_embeds.shape[-1])
        if inputs_embeds is not None
        else None
    )

    outputs = self.bert(
        input_ids,
        attention_mask=attention_mask,
        token_type_ids=token_type_ids,
        position_ids=position_ids,
        head_mask=head_mask,
        inputs_embeds=inputs_embeds,
        output_attentions=output_attentions,
        output_hidden_states=output_hidden_states,
        return_dict=return_dict,
    )

    pooled_output = outputs[1]

    pooled_output = self.dropout(pooled_output)
    logits = self.classifier(pooled_output)
    reshaped_logits = logits.view(-1, num_choices)

    loss = None
    if labels is not None:
        loss = ops.cross_entropy(reshaped_logits, labels)

    if not return_dict:
        output = (reshaped_logits,) + outputs[2:]
        return ((loss,) + output) if loss is not None else output

    return MultipleChoiceModelOutput(
        loss=loss,
        logits=reshaped_logits,
        hidden_states=outputs.hidden_states,
        attentions=outputs.attentions,
    )

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertForNextSentencePrediction

Bases: MegatronBertPreTrainedModel

Represents a MegatronBert model for next sentence prediction.

This class inherits from the MegatronBertPreTrainedModel and provides next sentence prediction functionality using the Megatron BERT model.

The class forwardor initializes the MegatronBertForNextSentencePrediction model with the given configuration.

The forward method takes input tensors and computes the next sentence prediction loss. It returns the next sentence predictor output.

PARAMETER DESCRIPTION
input_ids

The input tensor containing the indices of input sequence tokens in the vocabulary. Defaults to None.

TYPE: Optional[Tensor]

attention_mask

The input tensor containing indices specifying which tokens should be attended to. Defaults to None.

TYPE: Optional[Tensor]

token_type_ids

The input tensor containing the segment token indices to differentiate the sequences. Defaults to None.

TYPE: Optional[Tensor]

position_ids

The input tensor containing the position indices of each input token. Defaults to None.

TYPE: Optional[Tensor]

head_mask

The input tensor containing the mask for the attention heads. Defaults to None.

TYPE: Optional[Tensor]

inputs_embeds

The input tensor containing the embedded inputs. Defaults to None.

TYPE: Optional[Tensor]

labels

The tensor containing the labels for computing the next sequence prediction loss. Defaults to None.

TYPE: Optional[Tensor]

output_attentions

Whether to return attentions. Defaults to None.

TYPE: Optional[bool]

output_hidden_states

Whether to return hidden states. Defaults to None.

TYPE: Optional[bool]

return_dict

Whether to return a dictionary. Defaults to None.

TYPE: Optional[bool]

RETURNS DESCRIPTION

Union[Tuple, NextSentencePredictorOutput]: A tuple containing the next sentence prediction loss and the next sentence predictor output.
RAISES DESCRIPTION
FutureWarning

If the next_sentence_label argument is used, as it is deprecated.
Example 
>>> from transformers import AutoTokenizer, MegatronBertForNextSentencePrediction
...
>>> tokenizer = AutoTokenizer.from_pretrained("nvidia/megatron-bert-cased-345m")
>>> model = MegatronBertForNextSentencePrediction.from_pretrained("nvidia/megatron-bert-cased-345m")
...
>>> prompt = "In Italy, pizza served in formal settings, such as at a restaurant, is presented unsliced."
>>> next_sentence = "The sky is blue due to the shorter wavelength of blue light."
>>> encoding = tokenizer(prompt, next_sentence, return_tensors="pt")
...
>>> outputs = model(**encoding, labels=mindspore.Tensor([1]))
>>> logits = outputs.logits
>>> assert logits[0, 0] < logits[0, 1]  # next sentence was random
Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
2373
2374
2375
2376
2377
2378
2379
2380
2381
2382
2383
2384
2385
2386
2387
2388
2389
2390
2391
2392
2393
2394
2395
2396
2397
2398
2399
2400
2401
2402
2403
2404
2405
2406
2407
2408
2409
2410
2411
2412
2413
2414
2415
2416
2417
2418
2419
2420
2421
2422
2423
2424
2425
2426
2427
2428
2429
2430
2431
2432
2433
2434
2435
2436
2437
2438
2439
2440
2441
2442
2443
2444
2445
2446
2447
2448
2449
2450
2451
2452
2453
2454
2455
2456
2457
2458
2459
2460
2461
2462
2463
2464
2465
2466
2467
2468
2469
2470
2471
2472
2473
2474
2475
2476
2477
2478
2479
2480
2481
2482
2483
2484
2485
2486
2487
2488
2489
2490
2491
2492
2493
2494
2495
2496
2497
2498
2499
2500
2501
2502
2503
2504
2505
2506
2507
2508
2509
2510
2511
2512
2513
2514
2515
2516
2517
2518
2519
2520
2521
2522
2523
2524
2525
2526
2527
2528
2529
2530
2531
class MegatronBertForNextSentencePrediction(MegatronBertPreTrainedModel):

    """
    Represents a MegatronBert model for next sentence prediction.

    This class inherits from the MegatronBertPreTrainedModel and provides next sentence prediction functionality
    using the Megatron BERT model.

    The class forwardor initializes the MegatronBertForNextSentencePrediction model with the given configuration.

    The `forward` method takes input tensors and computes the next sentence prediction loss.
    It returns the next sentence predictor output.

    Args:
        input_ids (Optional[mindspore.Tensor], optional): The input tensor containing the indices of input sequence
            tokens in the vocabulary. Defaults to None.
        attention_mask (Optional[mindspore.Tensor], optional): The input tensor containing indices specifying which
            tokens should be attended to. Defaults to None.
        token_type_ids (Optional[mindspore.Tensor], optional): The input tensor containing the segment token indices
            to differentiate the sequences. Defaults to None.
        position_ids (Optional[mindspore.Tensor], optional): The input tensor containing the position indices of
            each input token. Defaults to None.
        head_mask (Optional[mindspore.Tensor], optional): The input tensor containing the mask for the attention heads.
            Defaults to None.
        inputs_embeds (Optional[mindspore.Tensor], optional): The input tensor containing the embedded inputs.
            Defaults to None.
        labels (Optional[mindspore.Tensor], optional): The tensor containing the labels for computing the next sequence
            prediction loss. Defaults to None.
        output_attentions (Optional[bool], optional): Whether to return attentions. Defaults to None.
        output_hidden_states (Optional[bool], optional): Whether to return hidden states. Defaults to None.
        return_dict (Optional[bool], optional): Whether to return a dictionary. Defaults to None.

    Returns:
        Union[Tuple, NextSentencePredictorOutput]: A tuple containing the next sentence prediction loss and the
            next sentence predictor output.

    Raises:
        FutureWarning: If the `next_sentence_label` argument is used, as it is deprecated.

    Example:
        ```python
        >>> from transformers import AutoTokenizer, MegatronBertForNextSentencePrediction
        ...
        >>> tokenizer = AutoTokenizer.from_pretrained("nvidia/megatron-bert-cased-345m")
        >>> model = MegatronBertForNextSentencePrediction.from_pretrained("nvidia/megatron-bert-cased-345m")
        ...
        >>> prompt = "In Italy, pizza served in formal settings, such as at a restaurant, is presented unsliced."
        >>> next_sentence = "The sky is blue due to the shorter wavelength of blue light."
        >>> encoding = tokenizer(prompt, next_sentence, return_tensors="pt")
        ...
        >>> outputs = model(**encoding, labels=mindspore.Tensor([1]))
        >>> logits = outputs.logits
        >>> assert logits[0, 0] < logits[0, 1]  # next sentence was random
        ```
    """
    def __init__(self, config):
        """
        Initializes an instance of the MegatronBertForNextSentencePrediction class.

        Args:
            self (MegatronBertForNextSentencePrediction): The instance of the class.
            config: The configuration object containing the settings for the model.

        Returns:
            None

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

        self.bert = MegatronBertModel(config)
        self.cls = MegatronBertOnlyNSPHead(config)

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

    def forward(
        self,
        input_ids: Optional[mindspore.Tensor] = None,
        attention_mask: Optional[mindspore.Tensor] = None,
        token_type_ids: Optional[mindspore.Tensor] = None,
        position_ids: Optional[mindspore.Tensor] = None,
        head_mask: Optional[mindspore.Tensor] = None,
        inputs_embeds: Optional[mindspore.Tensor] = None,
        labels: Optional[mindspore.Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        **kwargs,
    ) -> Union[Tuple, NextSentencePredictorOutput]:
        r"""
        Args:
            labels (`mindspore.Tensor` of shape `(batch_size,)`, *optional*):
                Labels for computing the next sequence prediction (classification) loss. Input should be a sequence pair
                (see `input_ids` docstring). Indices should be in `[0, 1]`:

                - 0 indicates sequence B is a continuation of sequence A,
                - 1 indicates sequence B is a random sequence.

        Returns:
            Union[Tuple, NextSentencePredictorOutput]

        Example:
            ```python
            >>> from transformers import AutoTokenizer, MegatronBertForNextSentencePrediction
            ...
            >>> tokenizer = AutoTokenizer.from_pretrained("nvidia/megatron-bert-cased-345m")
            >>> model = MegatronBertForNextSentencePrediction.from_pretrained("nvidia/megatron-bert-cased-345m")
            ...
            >>> prompt = "In Italy, pizza served in formal settings, such as at a restaurant, is presented unsliced."
            >>> next_sentence = "The sky is blue due to the shorter wavelength of blue light."
            >>> encoding = tokenizer(prompt, next_sentence, return_tensors="pt")
            ...
            >>> outputs = model(**encoding, labels=mindspore.Tensor([1]))
            >>> logits = outputs.logits
            >>> assert logits[0, 0] < logits[0, 1]  # next sentence was random
            ```
        """
        if "next_sentence_label" in kwargs:
            warnings.warn(
                "The `next_sentence_label` argument is deprecated and will be removed in a future version, use"
                " `labels` instead.",
                FutureWarning,
            )
            labels = kwargs.pop("next_sentence_label")

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

        outputs = self.bert(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        pooled_output = outputs[1]

        seq_relationship_scores = self.cls(pooled_output)

        next_sentence_loss = None
        if labels is not None:
            next_sentence_loss = ops.cross_entropy(seq_relationship_scores.view(-1, 2), labels.view(-1))

        if not return_dict:
            output = (seq_relationship_scores,) + outputs[2:]
            return ((next_sentence_loss,) + output) if next_sentence_loss is not None else output

        return NextSentencePredictorOutput(
            loss=next_sentence_loss,
            logits=seq_relationship_scores,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertForNextSentencePrediction.__init__(config)

Initializes an instance of the MegatronBertForNextSentencePrediction class.

PARAMETER DESCRIPTION
self

The instance of the class.

TYPE: MegatronBertForNextSentencePrediction

config

The configuration object containing the settings for the model.
RETURNS DESCRIPTION

None
Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
2428
2429
2430
2431
2432
2433
2434
2435
2436
2437
2438
2439
2440
2441
2442
2443
2444
2445
2446
2447
2448
def __init__(self, config):
    """
    Initializes an instance of the MegatronBertForNextSentencePrediction class.

    Args:
        self (MegatronBertForNextSentencePrediction): The instance of the class.
        config: The configuration object containing the settings for the model.

    Returns:
        None

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

    self.bert = MegatronBertModel(config)
    self.cls = MegatronBertOnlyNSPHead(config)

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

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertForNextSentencePrediction.forward(input_ids=None, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None, inputs_embeds=None, labels=None, output_attentions=None, output_hidden_states=None, return_dict=None, **kwargs)

PARAMETER DESCRIPTION
labels

Labels for computing the next sequence prediction (classification) loss. Input should be a sequence pair (see input_ids docstring). Indices should be in [0, 1]:

  • 0 indicates sequence B is a continuation of sequence A,
  • 1 indicates sequence B is a random sequence.

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

RETURNS DESCRIPTION
Union[Tuple, NextSentencePredictorOutput]

Union[Tuple, NextSentencePredictorOutput]
Example 
>>> from transformers import AutoTokenizer, MegatronBertForNextSentencePrediction
...
>>> tokenizer = AutoTokenizer.from_pretrained("nvidia/megatron-bert-cased-345m")
>>> model = MegatronBertForNextSentencePrediction.from_pretrained("nvidia/megatron-bert-cased-345m")
...
>>> prompt = "In Italy, pizza served in formal settings, such as at a restaurant, is presented unsliced."
>>> next_sentence = "The sky is blue due to the shorter wavelength of blue light."
>>> encoding = tokenizer(prompt, next_sentence, return_tensors="pt")
...
>>> outputs = model(**encoding, labels=mindspore.Tensor([1]))
>>> logits = outputs.logits
>>> assert logits[0, 0] < logits[0, 1]  # next sentence was random
Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
2450
2451
2452
2453
2454
2455
2456
2457
2458
2459
2460
2461
2462
2463
2464
2465
2466
2467
2468
2469
2470
2471
2472
2473
2474
2475
2476
2477
2478
2479
2480
2481
2482
2483
2484
2485
2486
2487
2488
2489
2490
2491
2492
2493
2494
2495
2496
2497
2498
2499
2500
2501
2502
2503
2504
2505
2506
2507
2508
2509
2510
2511
2512
2513
2514
2515
2516
2517
2518
2519
2520
2521
2522
2523
2524
2525
2526
2527
2528
2529
2530
2531
def forward(
    self,
    input_ids: Optional[mindspore.Tensor] = None,
    attention_mask: Optional[mindspore.Tensor] = None,
    token_type_ids: Optional[mindspore.Tensor] = None,
    position_ids: Optional[mindspore.Tensor] = None,
    head_mask: Optional[mindspore.Tensor] = None,
    inputs_embeds: Optional[mindspore.Tensor] = None,
    labels: Optional[mindspore.Tensor] = None,
    output_attentions: Optional[bool] = None,
    output_hidden_states: Optional[bool] = None,
    return_dict: Optional[bool] = None,
    **kwargs,
) -> Union[Tuple, NextSentencePredictorOutput]:
    r"""
    Args:
        labels (`mindspore.Tensor` of shape `(batch_size,)`, *optional*):
            Labels for computing the next sequence prediction (classification) loss. Input should be a sequence pair
            (see `input_ids` docstring). Indices should be in `[0, 1]`:

            - 0 indicates sequence B is a continuation of sequence A,
            - 1 indicates sequence B is a random sequence.

    Returns:
        Union[Tuple, NextSentencePredictorOutput]

    Example:
        ```python
        >>> from transformers import AutoTokenizer, MegatronBertForNextSentencePrediction
        ...
        >>> tokenizer = AutoTokenizer.from_pretrained("nvidia/megatron-bert-cased-345m")
        >>> model = MegatronBertForNextSentencePrediction.from_pretrained("nvidia/megatron-bert-cased-345m")
        ...
        >>> prompt = "In Italy, pizza served in formal settings, such as at a restaurant, is presented unsliced."
        >>> next_sentence = "The sky is blue due to the shorter wavelength of blue light."
        >>> encoding = tokenizer(prompt, next_sentence, return_tensors="pt")
        ...
        >>> outputs = model(**encoding, labels=mindspore.Tensor([1]))
        >>> logits = outputs.logits
        >>> assert logits[0, 0] < logits[0, 1]  # next sentence was random
        ```
    """
    if "next_sentence_label" in kwargs:
        warnings.warn(
            "The `next_sentence_label` argument is deprecated and will be removed in a future version, use"
            " `labels` instead.",
            FutureWarning,
        )
        labels = kwargs.pop("next_sentence_label")

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

    outputs = self.bert(
        input_ids,
        attention_mask=attention_mask,
        token_type_ids=token_type_ids,
        position_ids=position_ids,
        head_mask=head_mask,
        inputs_embeds=inputs_embeds,
        output_attentions=output_attentions,
        output_hidden_states=output_hidden_states,
        return_dict=return_dict,
    )

    pooled_output = outputs[1]

    seq_relationship_scores = self.cls(pooled_output)

    next_sentence_loss = None
    if labels is not None:
        next_sentence_loss = ops.cross_entropy(seq_relationship_scores.view(-1, 2), labels.view(-1))

    if not return_dict:
        output = (seq_relationship_scores,) + outputs[2:]
        return ((next_sentence_loss,) + output) if next_sentence_loss is not None else output

    return NextSentencePredictorOutput(
        loss=next_sentence_loss,
        logits=seq_relationship_scores,
        hidden_states=outputs.hidden_states,
        attentions=outputs.attentions,
    )

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertForPreTraining

Bases: MegatronBertPreTrainedModel

The MegatronBertForPreTraining class represents a pre-trained Megatron-BERT model for pre-training tasks. It inherits from the MegatronBertPreTrainedModel class and provides methods for forwarding the model, retrieving and setting output embeddings, and performing pre-training tasks such as masked language modeling and next sentence prediction.

The forward method takes input tensors for various model inputs and optional labels, and returns pre-training outputs including loss, prediction logits, sequence relationship logits, hidden states, and attentions. This method supports both batch and sequence-level losses for masked language modeling and next sentence prediction.

The get_output_embeddings method returns the decoder for predictions, while the set_output_embeddings method allows for updating the decoder with new embeddings.

This class is designed to work with the Megatron-BERT model and is intended to be used for pre-training tasks in natural language processing applications.

Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
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
class MegatronBertForPreTraining(MegatronBertPreTrainedModel):

    """
    The `MegatronBertForPreTraining` class represents a pre-trained Megatron-BERT model for pre-training tasks.
    It inherits from the `MegatronBertPreTrainedModel` class and provides methods for forwarding
    the model, retrieving and setting output embeddings, and performing pre-training tasks such as masked
    language modeling and next sentence prediction.

    The `forward` method takes input tensors for various model inputs and optional labels, and returns pre-training
    outputs including loss, prediction logits, sequence relationship logits, hidden states, and attentions.
    This method supports both batch and sequence-level losses for masked language modeling and next sentence prediction.

    The `get_output_embeddings` method returns the decoder for predictions, while the `set_output_embeddings` method
    allows for updating the decoder with new embeddings.

    This class is designed to work with the Megatron-BERT model and is intended to be used for pre-training tasks in
    natural language processing applications.
    """
    _tied_weights_keys = ["cls.predictions.decoder"]

    def __init__(self, config, add_binary_head=True):
        """
        Initializes a new instance of the MegatronBertForPreTraining class.

        Args:
            self (MegatronBertForPreTraining): The instance of the class.
            config (object): The configuration object containing the model's settings.
            add_binary_head (bool): Indicates whether to add a binary head to the model. Defaults to True.

        Returns:
            None.

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

        self.bert = MegatronBertModel(config)
        self.cls = MegatronBertPreTrainingHeads(config)

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

    def get_output_embeddings(self):
        """
        Returns the output embeddings of the MegatronBertForPreTraining model.

        Args:
            self (MegatronBertForPreTraining): The instance of the MegatronBertForPreTraining class.

        Returns:
            None.

        Raises:
            None.

        """
        return self.cls.predictions.decoder

    def set_output_embeddings(self, new_embeddings):
        """
        Sets the output embeddings of the MegatronBertForPreTraining model.

        Args:
            self (MegatronBertForPreTraining): An instance of the MegatronBertForPreTraining class.
            new_embeddings: The new embeddings to be set for the model's output.
                This should be a tensor of the same shape as the previous embeddings.

        Returns:
            None.

        Raises:
            None.
        """
        self.cls.predictions.decoder = new_embeddings

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

                - 0 indicates sequence B is a continuation of sequence A,
                - 1 indicates sequence B is a random sequence.
            kwargs (`Dict[str, any]`, optional, defaults to *{}*):
                Used to hide legacy arguments that have been deprecated.

        Returns:
            Union[Tuple, MegatronBertForPreTrainingOutput]

        Example:
            ```python
            >>> from transformers import AutoTokenizer, MegatronBertForPreTraining
            ...
            >>> tokenizer = AutoTokenizer.from_pretrained("nvidia/megatron-bert-cased-345m")
            >>> model = MegatronBertForPreTraining.from_pretrained("nvidia/megatron-bert-cased-345m")
            ...
            >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
            >>> outputs = model(**inputs)
            ...
            >>> prediction_logits = outputs.prediction_logits
            >>> seq_relationship_logits = outputs.seq_relationship_logits
            ```
        """
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        outputs = self.bert(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        sequence_output, pooled_output = outputs[:2]
        prediction_scores, seq_relationship_score = self.cls(sequence_output, pooled_output)

        total_loss = None
        if labels is not None and next_sentence_label is not None:
            masked_lm_loss = ops.cross_entropy(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
            next_sentence_loss = ops.cross_entropy(seq_relationship_score.view(-1, 2), next_sentence_label.view(-1))
            total_loss = masked_lm_loss + next_sentence_loss

        if not return_dict:
            output = (prediction_scores, seq_relationship_score) + outputs[2:]
            return ((total_loss,) + output) if total_loss is not None else output

        return MegatronBertForPreTrainingOutput(
            loss=total_loss,
            prediction_logits=prediction_scores,
            seq_relationship_logits=seq_relationship_score,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertForPreTraining.__init__(config, add_binary_head=True)

Initializes a new instance of the MegatronBertForPreTraining class.

PARAMETER DESCRIPTION
self

The instance of the class.

TYPE: MegatronBertForPreTraining

config

The configuration object containing the model's settings.

TYPE: object

add_binary_head

Indicates whether to add a binary head to the model. Defaults to True.

TYPE: bool DEFAULT: True

RETURNS DESCRIPTION

None.
Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
1813
1814
1815
1816
1817
1818
1819
1820
1821
1822
1823
1824
1825
1826
1827
1828
1829
1830
1831
1832
1833
1834
def __init__(self, config, add_binary_head=True):
    """
    Initializes a new instance of the MegatronBertForPreTraining class.

    Args:
        self (MegatronBertForPreTraining): The instance of the class.
        config (object): The configuration object containing the model's settings.
        add_binary_head (bool): Indicates whether to add a binary head to the model. Defaults to True.

    Returns:
        None.

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

    self.bert = MegatronBertModel(config)
    self.cls = MegatronBertPreTrainingHeads(config)

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

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertForPreTraining.forward(input_ids=None, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None, inputs_embeds=None, labels=None, next_sentence_label=None, output_attentions=None, output_hidden_states=None, return_dict=None)

PARAMETER DESCRIPTION
labels

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

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

next_sentence_label

Labels for computing the next sequence prediction (classification) loss. Input should be a sequence pair (see input_ids docstring) Indices should be in [0, 1]:

  • 0 indicates sequence B is a continuation of sequence A,
  • 1 indicates sequence B is a random sequence.

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

kwargs

Used to hide legacy arguments that have been deprecated.

TYPE: `Dict[str, any]`, optional, defaults to *{}*

RETURNS DESCRIPTION
Union[Tuple, MegatronBertForPreTrainingOutput]

Union[Tuple, MegatronBertForPreTrainingOutput]
Example 
>>> from transformers import AutoTokenizer, MegatronBertForPreTraining
...
>>> tokenizer = AutoTokenizer.from_pretrained("nvidia/megatron-bert-cased-345m")
>>> model = MegatronBertForPreTraining.from_pretrained("nvidia/megatron-bert-cased-345m")
...
>>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
>>> outputs = model(**inputs)
...
>>> prediction_logits = outputs.prediction_logits
>>> seq_relationship_logits = outputs.seq_relationship_logits
Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
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
def forward(
    self,
    input_ids: Optional[mindspore.Tensor] = None,
    attention_mask: Optional[mindspore.Tensor] = None,
    token_type_ids: Optional[mindspore.Tensor] = None,
    position_ids: Optional[mindspore.Tensor] = None,
    head_mask: Optional[mindspore.Tensor] = None,
    inputs_embeds: Optional[mindspore.Tensor] = None,
    labels: Optional[mindspore.Tensor] = None,
    next_sentence_label: Optional[mindspore.Tensor] = None,
    output_attentions: Optional[bool] = None,
    output_hidden_states: Optional[bool] = None,
    return_dict: Optional[bool] = None,
) -> Union[Tuple, MegatronBertForPreTrainingOutput]:
    r"""
    Args:
        labels (`mindspore.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
            Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
            config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the
            loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`
        next_sentence_label (`mindspore.Tensor` of shape `(batch_size,)`, *optional*):
            Labels for computing the next sequence prediction (classification) loss. Input should be a sequence pair
            (see `input_ids` docstring) Indices should be in `[0, 1]`:

            - 0 indicates sequence B is a continuation of sequence A,
            - 1 indicates sequence B is a random sequence.
        kwargs (`Dict[str, any]`, optional, defaults to *{}*):
            Used to hide legacy arguments that have been deprecated.

    Returns:
        Union[Tuple, MegatronBertForPreTrainingOutput]

    Example:
        ```python
        >>> from transformers import AutoTokenizer, MegatronBertForPreTraining
        ...
        >>> tokenizer = AutoTokenizer.from_pretrained("nvidia/megatron-bert-cased-345m")
        >>> model = MegatronBertForPreTraining.from_pretrained("nvidia/megatron-bert-cased-345m")
        ...
        >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
        >>> outputs = model(**inputs)
        ...
        >>> prediction_logits = outputs.prediction_logits
        >>> seq_relationship_logits = outputs.seq_relationship_logits
        ```
    """
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict

    outputs = self.bert(
        input_ids,
        attention_mask=attention_mask,
        token_type_ids=token_type_ids,
        position_ids=position_ids,
        head_mask=head_mask,
        inputs_embeds=inputs_embeds,
        output_attentions=output_attentions,
        output_hidden_states=output_hidden_states,
        return_dict=return_dict,
    )

    sequence_output, pooled_output = outputs[:2]
    prediction_scores, seq_relationship_score = self.cls(sequence_output, pooled_output)

    total_loss = None
    if labels is not None and next_sentence_label is not None:
        masked_lm_loss = ops.cross_entropy(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
        next_sentence_loss = ops.cross_entropy(seq_relationship_score.view(-1, 2), next_sentence_label.view(-1))
        total_loss = masked_lm_loss + next_sentence_loss

    if not return_dict:
        output = (prediction_scores, seq_relationship_score) + outputs[2:]
        return ((total_loss,) + output) if total_loss is not None else output

    return MegatronBertForPreTrainingOutput(
        loss=total_loss,
        prediction_logits=prediction_scores,
        seq_relationship_logits=seq_relationship_score,
        hidden_states=outputs.hidden_states,
        attentions=outputs.attentions,
    )

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertForPreTraining.get_output_embeddings()

Returns the output embeddings of the MegatronBertForPreTraining model.

PARAMETER DESCRIPTION
self

The instance of the MegatronBertForPreTraining class.

TYPE: MegatronBertForPreTraining

RETURNS DESCRIPTION

None.
Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
1836
1837
1838
1839
1840
1841
1842
1843
1844
1845
1846
1847
1848
1849
1850
def get_output_embeddings(self):
    """
    Returns the output embeddings of the MegatronBertForPreTraining model.

    Args:
        self (MegatronBertForPreTraining): The instance of the MegatronBertForPreTraining class.

    Returns:
        None.

    Raises:
        None.

    """
    return self.cls.predictions.decoder

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertForPreTraining.set_output_embeddings(new_embeddings)

Sets the output embeddings of the MegatronBertForPreTraining model.

PARAMETER DESCRIPTION
self

An instance of the MegatronBertForPreTraining class.

TYPE: MegatronBertForPreTraining

new_embeddings

The new embeddings to be set for the model's output. This should be a tensor of the same shape as the previous embeddings.
RETURNS DESCRIPTION

None.
Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
1852
1853
1854
1855
1856
1857
1858
1859
1860
1861
1862
1863
1864
1865
1866
1867
def set_output_embeddings(self, new_embeddings):
    """
    Sets the output embeddings of the MegatronBertForPreTraining model.

    Args:
        self (MegatronBertForPreTraining): An instance of the MegatronBertForPreTraining class.
        new_embeddings: The new embeddings to be set for the model's output.
            This should be a tensor of the same shape as the previous embeddings.

    Returns:
        None.

    Raises:
        None.
    """
    self.cls.predictions.decoder = new_embeddings

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertForPreTrainingOutput dataclass

Bases: ModelOutput

Output type of [MegatronBertForPreTraining].

PARAMETER DESCRIPTION
loss

Total loss as the sum of the masked language modeling loss and the next sequence prediction (classification) loss.

TYPE: *optional*, returned when `labels` is provided, `mindspore.Tensor` of shape `(1,)` DEFAULT: None

prediction_logits

Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).

TYPE: `mindspore.Tensor` of shape `(batch_size, sequence_length, config.vocab_size)` DEFAULT: None

seq_relationship_logits

Prediction scores of the next sequence prediction (classification) head (scores of True/False continuation before SoftMax).

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

Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
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
@dataclass
# Copied from transformers.models.bert.modeling_bert.BertForPreTrainingOutput with Bert->MegatronBert
class MegatronBertForPreTrainingOutput(ModelOutput):
    """
    Output type of [`MegatronBertForPreTraining`].

    Args:
        loss (*optional*, returned when `labels` is provided, `mindspore.Tensor` of shape `(1,)`):
            Total loss as the sum of the masked language modeling loss and the next sequence prediction
            (classification) loss.
        prediction_logits (`mindspore.Tensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
        seq_relationship_logits (`mindspore.Tensor` of shape `(batch_size, 2)`):
            Prediction scores of the next sequence prediction (classification) head (scores of True/False continuation
            before SoftMax).
        hidden_states (`tuple(mindspore.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or
            when `config.output_hidden_states=True`):
            Tuple of `mindspore.Tensor` (one for the output of the embeddings + one for the output of each layer) of
            shape `(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        attentions (`tuple(mindspore.Tensor)`, *optional*, returned when `output_attentions=True` is passed or
            when `config.output_attentions=True`):
            Tuple of `mindspore.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
            sequence_length)`.

            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
            heads.
    """
    loss: Optional[mindspore.Tensor] = None
    prediction_logits: mindspore.Tensor = None
    seq_relationship_logits: mindspore.Tensor = None
    hidden_states: Optional[Tuple[mindspore.Tensor]] = None
    attentions: Optional[Tuple[mindspore.Tensor]] = None

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertForQuestionAnswering

Bases: MegatronBertPreTrainedModel

A class representing a Megatron-BERT model for question answering.

This class inherits from the MegatronBertPreTrainedModel class and is specifically designed for question answering tasks. It includes methods for forwarding the model and generating predictions.

ATTRIBUTE DESCRIPTION
num_labels

The number of labels for token classification.

TYPE: int

bert

The Megatron-BERT model.

TYPE: MegatronBertModel

qa_outputs

The dense layer for question answering outputs.

TYPE: Linear

METHOD DESCRIPTION
__init__

Initializes the MegatronBertForQuestionAnswering instance.
forward

Constructs the Megatron-BERT model and generates predictions for question answering tasks.
Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
2885
2886
2887
2888
2889
2890
2891
2892
2893
2894
2895
2896
2897
2898
2899
2900
2901
2902
2903
2904
2905
2906
2907
2908
2909
2910
2911
2912
2913
2914
2915
2916
2917
2918
2919
2920
2921
2922
2923
2924
2925
2926
2927
2928
2929
2930
2931
2932
2933
2934
2935
2936
2937
2938
2939
2940
2941
2942
2943
2944
2945
2946
2947
2948
2949
2950
2951
2952
2953
2954
2955
2956
2957
2958
2959
2960
2961
2962
2963
2964
2965
2966
2967
2968
2969
2970
2971
2972
2973
2974
2975
2976
2977
2978
2979
2980
2981
2982
2983
2984
2985
2986
2987
2988
2989
2990
2991
2992
2993
2994
2995
2996
2997
2998
2999
3000
class MegatronBertForQuestionAnswering(MegatronBertPreTrainedModel):

    """A class representing a Megatron-BERT model for question answering.

    This class inherits from the MegatronBertPreTrainedModel class and is specifically designed for question answering tasks.
    It includes methods for forwarding the model and generating predictions.

    Attributes:
        num_labels (int): The number of labels for token classification.
        bert (MegatronBertModel): The Megatron-BERT model.
        qa_outputs (nn.Linear): The dense layer for question answering outputs.

    Methods:
        __init__: Initializes the MegatronBertForQuestionAnswering instance.
        forward: Constructs the Megatron-BERT model and generates predictions for question answering tasks.

    """
    def __init__(self, config):
        """
        Initialize the MegatronBertForQuestionAnswering class.

        Args:
            self (object): The instance of the class.
            config (object):
                The configuration object containing the settings for the model.

                - num_labels (int): The number of labels for question answering.

        Returns:
            None.

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

        self.bert = MegatronBertModel(config, add_pooling_layer=False)
        self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)

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

    def forward(
        self,
        input_ids: Optional[mindspore.Tensor] = None,
        attention_mask: Optional[mindspore.Tensor] = None,
        token_type_ids: Optional[mindspore.Tensor] = None,
        position_ids: Optional[mindspore.Tensor] = None,
        head_mask: Optional[mindspore.Tensor] = None,
        inputs_embeds: Optional[mindspore.Tensor] = None,
        start_positions: Optional[mindspore.Tensor] = None,
        end_positions: Optional[mindspore.Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, QuestionAnsweringModelOutput]:
        r"""
        Args:
            start_positions (`mindspore.Tensor` of shape `(batch_size,)`, *optional*):
                Labels for position (index) of the start of the labelled span for computing the token classification loss.
                Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
                are not taken into account for computing the loss.
            end_positions (`mindspore.Tensor` of shape `(batch_size,)`, *optional*):
                Labels for position (index) of the end of the labelled span for computing the token classification loss.
                Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
                are not taken into account for computing the loss.
        """
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        outputs = self.bert(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        sequence_output = outputs[0]

        logits = self.qa_outputs(sequence_output)
        start_logits, end_logits = logits.split(1, axis=-1)
        start_logits = start_logits.squeeze(-1)
        end_logits = end_logits.squeeze(-1)

        total_loss = None
        if start_positions is not None and end_positions is not None:
            # If we are on multi-GPU, split add a dimension
            if len(start_positions.shape) > 1:
                start_positions = start_positions.squeeze(-1)
            if len(end_positions.shape) > 1:
                end_positions = end_positions.squeeze(-1)
            # sometimes the start/end positions are outside our model inputs, we ignore these terms
            ignored_index = start_logits.shape[1]
            start_positions = start_positions.clamp(0, ignored_index)
            end_positions = end_positions.clamp(0, ignored_index)

            start_loss = ops.cross_entropy(start_logits, start_positions, ignore_index=ignored_index)
            end_loss = ops.cross_entropy(end_logits, end_positions, ignore_index=ignored_index)
            total_loss = (start_loss + end_loss) / 2

        if not return_dict:
            output = (start_logits, end_logits) + outputs[2:]
            return ((total_loss,) + output) if total_loss is not None else output

        return QuestionAnsweringModelOutput(
            loss=total_loss,
            start_logits=start_logits,
            end_logits=end_logits,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertForQuestionAnswering.__init__(config)

Initialize the MegatronBertForQuestionAnswering class.

PARAMETER DESCRIPTION
self

The instance of the class.

TYPE: object

config

The configuration object containing the settings for the model.

  • num_labels (int): The number of labels for question answering.

TYPE: object

RETURNS DESCRIPTION

None.
Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
2902
2903
2904
2905
2906
2907
2908
2909
2910
2911
2912
2913
2914
2915
2916
2917
2918
2919
2920
2921
2922
2923
2924
2925
2926
def __init__(self, config):
    """
    Initialize the MegatronBertForQuestionAnswering class.

    Args:
        self (object): The instance of the class.
        config (object):
            The configuration object containing the settings for the model.

            - num_labels (int): The number of labels for question answering.

    Returns:
        None.

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

    self.bert = MegatronBertModel(config, add_pooling_layer=False)
    self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)

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

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertForQuestionAnswering.forward(input_ids=None, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None, inputs_embeds=None, start_positions=None, end_positions=None, output_attentions=None, output_hidden_states=None, return_dict=None)

PARAMETER DESCRIPTION
start_positions

Labels for position (index) of the start of the labelled span for computing the token classification loss. Positions are clamped to the length of the sequence (sequence_length). Position outside of the sequence are not taken into account for computing the loss.

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

end_positions

Labels for position (index) of the end of the labelled span for computing the token classification loss. Positions are clamped to the length of the sequence (sequence_length). Position outside of the sequence are not taken into account for computing the loss.

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

Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
2928
2929
2930
2931
2932
2933
2934
2935
2936
2937
2938
2939
2940
2941
2942
2943
2944
2945
2946
2947
2948
2949
2950
2951
2952
2953
2954
2955
2956
2957
2958
2959
2960
2961
2962
2963
2964
2965
2966
2967
2968
2969
2970
2971
2972
2973
2974
2975
2976
2977
2978
2979
2980
2981
2982
2983
2984
2985
2986
2987
2988
2989
2990
2991
2992
2993
2994
2995
2996
2997
2998
2999
3000
def forward(
    self,
    input_ids: Optional[mindspore.Tensor] = None,
    attention_mask: Optional[mindspore.Tensor] = None,
    token_type_ids: Optional[mindspore.Tensor] = None,
    position_ids: Optional[mindspore.Tensor] = None,
    head_mask: Optional[mindspore.Tensor] = None,
    inputs_embeds: Optional[mindspore.Tensor] = None,
    start_positions: Optional[mindspore.Tensor] = None,
    end_positions: Optional[mindspore.Tensor] = None,
    output_attentions: Optional[bool] = None,
    output_hidden_states: Optional[bool] = None,
    return_dict: Optional[bool] = None,
) -> Union[Tuple, QuestionAnsweringModelOutput]:
    r"""
    Args:
        start_positions (`mindspore.Tensor` of shape `(batch_size,)`, *optional*):
            Labels for position (index) of the start of the labelled span for computing the token classification loss.
            Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
            are not taken into account for computing the loss.
        end_positions (`mindspore.Tensor` of shape `(batch_size,)`, *optional*):
            Labels for position (index) of the end of the labelled span for computing the token classification loss.
            Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
            are not taken into account for computing the loss.
    """
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict

    outputs = self.bert(
        input_ids,
        attention_mask=attention_mask,
        token_type_ids=token_type_ids,
        position_ids=position_ids,
        head_mask=head_mask,
        inputs_embeds=inputs_embeds,
        output_attentions=output_attentions,
        output_hidden_states=output_hidden_states,
        return_dict=return_dict,
    )

    sequence_output = outputs[0]

    logits = self.qa_outputs(sequence_output)
    start_logits, end_logits = logits.split(1, axis=-1)
    start_logits = start_logits.squeeze(-1)
    end_logits = end_logits.squeeze(-1)

    total_loss = None
    if start_positions is not None and end_positions is not None:
        # If we are on multi-GPU, split add a dimension
        if len(start_positions.shape) > 1:
            start_positions = start_positions.squeeze(-1)
        if len(end_positions.shape) > 1:
            end_positions = end_positions.squeeze(-1)
        # sometimes the start/end positions are outside our model inputs, we ignore these terms
        ignored_index = start_logits.shape[1]
        start_positions = start_positions.clamp(0, ignored_index)
        end_positions = end_positions.clamp(0, ignored_index)

        start_loss = ops.cross_entropy(start_logits, start_positions, ignore_index=ignored_index)
        end_loss = ops.cross_entropy(end_logits, end_positions, ignore_index=ignored_index)
        total_loss = (start_loss + end_loss) / 2

    if not return_dict:
        output = (start_logits, end_logits) + outputs[2:]
        return ((total_loss,) + output) if total_loss is not None else output

    return QuestionAnsweringModelOutput(
        loss=total_loss,
        start_logits=start_logits,
        end_logits=end_logits,
        hidden_states=outputs.hidden_states,
        attentions=outputs.attentions,
    )

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertForSequenceClassification

Bases: MegatronBertPreTrainedModel

This class represents a MegatronBERT model for sequence classification tasks. It inherits from the MegatronBertPreTrainedModel class and includes methods for initializing the model and generating classification outputs.

The forward method takes various input tensors and computes the sequence classification/regression loss based on the configured problem type. It returns the classification logits and optionally the loss, hidden states, and attentions.

The __init__ method initializes the model with the provided configuration and sets up the BERT model, dropout layer, and classifier for sequence classification.

The class also provides detailed documentation for the forward method, including information about the input and output tensors, as well as the optional labels for computing the classification/regression loss.

For complete method signatures and code, please refer to the source code.

Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
2534
2535
2536
2537
2538
2539
2540
2541
2542
2543
2544
2545
2546
2547
2548
2549
2550
2551
2552
2553
2554
2555
2556
2557
2558
2559
2560
2561
2562
2563
2564
2565
2566
2567
2568
2569
2570
2571
2572
2573
2574
2575
2576
2577
2578
2579
2580
2581
2582
2583
2584
2585
2586
2587
2588
2589
2590
2591
2592
2593
2594
2595
2596
2597
2598
2599
2600
2601
2602
2603
2604
2605
2606
2607
2608
2609
2610
2611
2612
2613
2614
2615
2616
2617
2618
2619
2620
2621
2622
2623
2624
2625
2626
2627
2628
2629
2630
2631
2632
2633
2634
2635
2636
2637
2638
2639
2640
2641
2642
2643
2644
class MegatronBertForSequenceClassification(MegatronBertPreTrainedModel):

    """
    This class represents a MegatronBERT model for sequence classification tasks. It inherits from the
    MegatronBertPreTrainedModel class and includes methods for initializing the model and generating
    classification outputs.

    The `forward` method takes various input tensors and computes the sequence classification/regression loss based
    on the configured problem type. It returns the classification logits and optionally the loss, hidden states, and
    attentions.

    The `__init__` method initializes the model with the provided configuration and sets up the BERT model, dropout layer,
    and classifier for sequence classification.

    The class also provides detailed documentation for the `forward` method, including information about the input and
    output tensors, as well as the optional labels for computing the classification/regression loss.

    For complete method signatures and code, please refer to the source code.
    """
    def __init__(self, config):
        """
        Initializes an instance of the MegatronBertForSequenceClassification class.

        Args:
            self : The object instance.
            config : An object of type 'Config' containing the configuration settings for the model.

        Returns:
            None

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

        self.bert = MegatronBertModel(config)
        self.dropout = nn.Dropout(p=config.hidden_dropout_prob)
        self.classifier = nn.Linear(config.hidden_size, config.num_labels)

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

    def forward(
        self,
        input_ids: Optional[mindspore.Tensor] = None,
        attention_mask: Optional[mindspore.Tensor] = None,
        token_type_ids: Optional[mindspore.Tensor] = None,
        position_ids: Optional[mindspore.Tensor] = None,
        head_mask: Optional[mindspore.Tensor] = None,
        inputs_embeds: Optional[mindspore.Tensor] = None,
        labels: Optional[mindspore.Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, SequenceClassifierOutput]:
        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

        outputs = self.bert(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        pooled_output = outputs[1]

        pooled_output = self.dropout(pooled_output)
        logits = self.classifier(pooled_output)

        loss = None
        if labels is not None:
            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(logits.squeeze(), labels.squeeze())
                else:
                    loss = ops.mse_loss(logits, labels)
            elif self.config.problem_type == "single_label_classification":
                loss = ops.cross_entropy(logits.view(-1, self.num_labels), labels.view(-1))
            elif self.config.problem_type == "multi_label_classification":
                loss = ops.binary_cross_entropy_with_logits(logits, labels)
        if not return_dict:
            output = (logits,) + outputs[2:]
            return ((loss,) + output) if loss is not None else output

        return SequenceClassifierOutput(
            loss=loss,
            logits=logits,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertForSequenceClassification.__init__(config)

Initializes an instance of the MegatronBertForSequenceClassification class.

PARAMETER DESCRIPTION
self

The object instance.
config

An object of type 'Config' containing the configuration settings for the model.
RETURNS DESCRIPTION

None
Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
2553
2554
2555
2556
2557
2558
2559
2560
2561
2562
2563
2564
2565
2566
2567
2568
2569
2570
2571
2572
2573
2574
2575
def __init__(self, config):
    """
    Initializes an instance of the MegatronBertForSequenceClassification class.

    Args:
        self : The object instance.
        config : An object of type 'Config' containing the configuration settings for the model.

    Returns:
        None

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

    self.bert = MegatronBertModel(config)
    self.dropout = nn.Dropout(p=config.hidden_dropout_prob)
    self.classifier = nn.Linear(config.hidden_size, config.num_labels)

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

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertForSequenceClassification.forward(input_ids=None, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None, inputs_embeds=None, labels=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/megatron_bert/modeling_megatron_bert.py 
2577
2578
2579
2580
2581
2582
2583
2584
2585
2586
2587
2588
2589
2590
2591
2592
2593
2594
2595
2596
2597
2598
2599
2600
2601
2602
2603
2604
2605
2606
2607
2608
2609
2610
2611
2612
2613
2614
2615
2616
2617
2618
2619
2620
2621
2622
2623
2624
2625
2626
2627
2628
2629
2630
2631
2632
2633
2634
2635
2636
2637
2638
2639
2640
2641
2642
2643
2644
def forward(
    self,
    input_ids: Optional[mindspore.Tensor] = None,
    attention_mask: Optional[mindspore.Tensor] = None,
    token_type_ids: Optional[mindspore.Tensor] = None,
    position_ids: Optional[mindspore.Tensor] = None,
    head_mask: Optional[mindspore.Tensor] = None,
    inputs_embeds: Optional[mindspore.Tensor] = None,
    labels: Optional[mindspore.Tensor] = None,
    output_attentions: Optional[bool] = None,
    output_hidden_states: Optional[bool] = None,
    return_dict: Optional[bool] = None,
) -> Union[Tuple, SequenceClassifierOutput]:
    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

    outputs = self.bert(
        input_ids,
        attention_mask=attention_mask,
        token_type_ids=token_type_ids,
        position_ids=position_ids,
        head_mask=head_mask,
        inputs_embeds=inputs_embeds,
        output_attentions=output_attentions,
        output_hidden_states=output_hidden_states,
        return_dict=return_dict,
    )

    pooled_output = outputs[1]

    pooled_output = self.dropout(pooled_output)
    logits = self.classifier(pooled_output)

    loss = None
    if labels is not None:
        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(logits.squeeze(), labels.squeeze())
            else:
                loss = ops.mse_loss(logits, labels)
        elif self.config.problem_type == "single_label_classification":
            loss = ops.cross_entropy(logits.view(-1, self.num_labels), labels.view(-1))
        elif self.config.problem_type == "multi_label_classification":
            loss = ops.binary_cross_entropy_with_logits(logits, labels)
    if not return_dict:
        output = (logits,) + outputs[2:]
        return ((loss,) + output) if loss is not None else output

    return SequenceClassifierOutput(
        loss=loss,
        logits=logits,
        hidden_states=outputs.hidden_states,
        attentions=outputs.attentions,
    )

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertForTokenClassification

Bases: MegatronBertPreTrainedModel

This class represents a token classification model based on the Megatron BERT architecture. It inherits from the MegatronBertPreTrainedModel class and includes functionality for token classification tasks.

The init method initializes the MegatronBertForTokenClassification instance with the provided configuration. It sets the number of labels, initializes the BERT model without a pooling layer, sets the dropout probability, and initializes the classifier.

The forward method takes input tensors for token classification, such as input_ids, attention_mask, token_type_ids, position_ids, head_mask, and inputs_embeds. It also supports optional arguments for labels, output_attentions, output_hidden_states, and return_dict. The method returns TokenClassifierOutput containing the loss, logits, hidden states, and attentions. If labels are provided, it computes the token classification loss using cross-entropy.

The class provides detailed docstrings for each method, including parameter descriptions and return types for improved documentation and understanding.

Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
2782
2783
2784
2785
2786
2787
2788
2789
2790
2791
2792
2793
2794
2795
2796
2797
2798
2799
2800
2801
2802
2803
2804
2805
2806
2807
2808
2809
2810
2811
2812
2813
2814
2815
2816
2817
2818
2819
2820
2821
2822
2823
2824
2825
2826
2827
2828
2829
2830
2831
2832
2833
2834
2835
2836
2837
2838
2839
2840
2841
2842
2843
2844
2845
2846
2847
2848
2849
2850
2851
2852
2853
2854
2855
2856
2857
2858
2859
2860
2861
2862
2863
2864
2865
2866
2867
2868
2869
2870
2871
2872
2873
2874
2875
2876
2877
2878
2879
2880
2881
2882
class MegatronBertForTokenClassification(MegatronBertPreTrainedModel):

    """
    This class represents a token classification model based on the Megatron BERT architecture.
    It inherits from the MegatronBertPreTrainedModel class and includes functionality for token classification tasks.

    The __init__ method initializes the MegatronBertForTokenClassification instance with the provided configuration.
    It sets the number of labels, initializes the BERT model without a pooling layer, sets the dropout probability,
    and initializes the classifier.

    The forward method takes input tensors for token classification, such as input_ids, attention_mask, token_type_ids,
    position_ids, head_mask, and inputs_embeds. It also supports optional arguments for labels, output_attentions,
    output_hidden_states, and return_dict. The method returns TokenClassifierOutput containing the loss, logits,
    hidden states, and attentions. If labels are provided, it computes the token classification loss using cross-entropy.

    The class provides detailed docstrings for each method, including parameter descriptions and return types for
    improved documentation and understanding.
    """
    def __init__(self, config):
        """
        Initializes an instance of the MegatronBertForTokenClassification class.

        Args:
            self: The instance of the class.
            config: An object containing configuration parameters for the model.
                It should include the following attributes:

                - num_labels (int): The number of labels for token classification.
                - hidden_dropout_prob (float): The dropout probability for the hidden layers.

        Returns:
            None.

        Raises:
            TypeError: If the config parameter is not provided or is not of the correct type.
            ValueError: If the num_labels attribute in the config is not provided or is not a positive integer.
            ValueError: If the hidden_dropout_prob attribute in the config is not provided or is not a valid
                probability value (0 <= hidden_dropout_prob <= 1).
            RuntimeError: If an error occurs during the initialization process.
        """
        super().__init__(config)
        self.num_labels = config.num_labels

        self.bert = MegatronBertModel(config, add_pooling_layer=False)
        self.dropout = nn.Dropout(p=config.hidden_dropout_prob)
        self.classifier = nn.Linear(config.hidden_size, config.num_labels)

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

    def forward(
        self,
        input_ids: Optional[mindspore.Tensor] = None,
        attention_mask: Optional[mindspore.Tensor] = None,
        token_type_ids: Optional[mindspore.Tensor] = None,
        position_ids: Optional[mindspore.Tensor] = None,
        head_mask: Optional[mindspore.Tensor] = None,
        inputs_embeds: Optional[mindspore.Tensor] = None,
        labels: Optional[mindspore.Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, TokenClassifierOutput]:
        r"""
        Args:
            labels (`mindspore.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
                Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`.
        """
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        outputs = self.bert(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        sequence_output = outputs[0]

        sequence_output = self.dropout(sequence_output)
        logits = self.classifier(sequence_output)

        loss = None
        if labels is not None:
            loss = ops.cross_entropy(logits.view(-1, self.num_labels), labels.view(-1))

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

        return TokenClassifierOutput(
            loss=loss,
            logits=logits,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertForTokenClassification.__init__(config)

Initializes an instance of the MegatronBertForTokenClassification class.

PARAMETER DESCRIPTION
self

The instance of the class.
config

An object containing configuration parameters for the model. It should include the following attributes:

  • num_labels (int): The number of labels for token classification.
  • hidden_dropout_prob (float): The dropout probability for the hidden layers.
RETURNS DESCRIPTION

None.
RAISES DESCRIPTION
TypeError

If the config parameter is not provided or is not of the correct type.
ValueError

If the num_labels attribute in the config is not provided or is not a positive integer.
ValueError

If the hidden_dropout_prob attribute in the config is not provided or is not a valid probability value (0 <= hidden_dropout_prob <= 1).
RuntimeError

If an error occurs during the initialization process.
Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
2800
2801
2802
2803
2804
2805
2806
2807
2808
2809
2810
2811
2812
2813
2814
2815
2816
2817
2818
2819
2820
2821
2822
2823
2824
2825
2826
2827
2828
2829
2830
def __init__(self, config):
    """
    Initializes an instance of the MegatronBertForTokenClassification class.

    Args:
        self: The instance of the class.
        config: An object containing configuration parameters for the model.
            It should include the following attributes:

            - num_labels (int): The number of labels for token classification.
            - hidden_dropout_prob (float): The dropout probability for the hidden layers.

    Returns:
        None.

    Raises:
        TypeError: If the config parameter is not provided or is not of the correct type.
        ValueError: If the num_labels attribute in the config is not provided or is not a positive integer.
        ValueError: If the hidden_dropout_prob attribute in the config is not provided or is not a valid
            probability value (0 <= hidden_dropout_prob <= 1).
        RuntimeError: If an error occurs during the initialization process.
    """
    super().__init__(config)
    self.num_labels = config.num_labels

    self.bert = MegatronBertModel(config, add_pooling_layer=False)
    self.dropout = nn.Dropout(p=config.hidden_dropout_prob)
    self.classifier = nn.Linear(config.hidden_size, config.num_labels)

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

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertForTokenClassification.forward(input_ids=None, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None, inputs_embeds=None, labels=None, output_attentions=None, output_hidden_states=None, return_dict=None)

PARAMETER DESCRIPTION
labels

Labels for computing the token classification loss. Indices should be in [0, ..., config.num_labels - 1].

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

Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
2832
2833
2834
2835
2836
2837
2838
2839
2840
2841
2842
2843
2844
2845
2846
2847
2848
2849
2850
2851
2852
2853
2854
2855
2856
2857
2858
2859
2860
2861
2862
2863
2864
2865
2866
2867
2868
2869
2870
2871
2872
2873
2874
2875
2876
2877
2878
2879
2880
2881
2882
def forward(
    self,
    input_ids: Optional[mindspore.Tensor] = None,
    attention_mask: Optional[mindspore.Tensor] = None,
    token_type_ids: Optional[mindspore.Tensor] = None,
    position_ids: Optional[mindspore.Tensor] = None,
    head_mask: Optional[mindspore.Tensor] = None,
    inputs_embeds: Optional[mindspore.Tensor] = None,
    labels: Optional[mindspore.Tensor] = None,
    output_attentions: Optional[bool] = None,
    output_hidden_states: Optional[bool] = None,
    return_dict: Optional[bool] = None,
) -> Union[Tuple, TokenClassifierOutput]:
    r"""
    Args:
        labels (`mindspore.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
            Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`.
    """
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict

    outputs = self.bert(
        input_ids,
        attention_mask=attention_mask,
        token_type_ids=token_type_ids,
        position_ids=position_ids,
        head_mask=head_mask,
        inputs_embeds=inputs_embeds,
        output_attentions=output_attentions,
        output_hidden_states=output_hidden_states,
        return_dict=return_dict,
    )

    sequence_output = outputs[0]

    sequence_output = self.dropout(sequence_output)
    logits = self.classifier(sequence_output)

    loss = None
    if labels is not None:
        loss = ops.cross_entropy(logits.view(-1, self.num_labels), labels.view(-1))

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

    return TokenClassifierOutput(
        loss=loss,
        logits=logits,
        hidden_states=outputs.hidden_states,
        attentions=outputs.attentions,
    )

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertIntermediate

Bases: Module

Represents an intermediate layer of a Megatron-style BERT model for processing hidden states.

This class inherits from nn.Module and contains methods for initializing the intermediate layer and processing hidden states through dense and activation functions.

ATTRIBUTE DESCRIPTION
dense

The dense layer used for processing hidden states.

TYPE: Linear

intermediate_act_fn

The activation function applied to the hidden states.

TYPE: function

METHOD DESCRIPTION
__init__

Initializes the MegatronBertIntermediate instance with the provided configuration.
forward

Processes the input hidden states through the dense layer and activation function, returning the transformed hidden states.
Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
class MegatronBertIntermediate(nn.Module):

    """
    Represents an intermediate layer of a Megatron-style BERT model for processing hidden states.

    This class inherits from nn.Module and contains methods for initializing the intermediate layer and processing
    hidden states through dense and activation functions.

    Attributes:
        dense (nn.Linear): The dense layer used for processing hidden states.
        intermediate_act_fn (function): The activation function applied to the hidden states.

    Methods:
        __init__: Initializes the MegatronBertIntermediate instance with the provided configuration.
        forward: Processes the input hidden states through the dense layer and activation function, returning
            the transformed hidden states.
    """
    def __init__(self, config):
        """
        Initializes an instance of the MegatronBertIntermediate class.

        Args:
            self: The instance of the class.
            config (object): The configuration object containing the settings for the MegatronBertIntermediate.
                It should have the following attributes:

                - hidden_size (int): The size of the hidden layer.
                - intermediate_size (int): The size of the intermediate layer.
                - hidden_act (str or function): The activation function for the hidden layer.

                    - If it is a string, it should be one of the predefined activation functions available in the
                    ACT2FN dictionary.
                    - If it is a function, it should be a custom activation function.

        Returns:
            None.

        Raises:
            None.
        """
        super().__init__()
        self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
        if isinstance(config.hidden_act, str):
            self.intermediate_act_fn = ACT2FN[config.hidden_act]
        else:
            self.intermediate_act_fn = config.hidden_act

    def forward(self, hidden_states: mindspore.Tensor) -> mindspore.Tensor:
        """
        Constructs the intermediate layer of the Megatron BERT model.

        Args:
            self (MegatronBertIntermediate): An instance of the MegatronBertIntermediate class.
            hidden_states (mindspore.Tensor): The input hidden states tensor.

        Returns:
            mindspore.Tensor: The output hidden states tensor after applying the intermediate layer.

        Raises:
            None.
        """
        hidden_states = self.dense(hidden_states)
        hidden_states = self.intermediate_act_fn(hidden_states)
        return hidden_states

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertIntermediate.__init__(config)

Initializes an instance of the MegatronBertIntermediate class.

PARAMETER DESCRIPTION
self

The instance of the class.
config

The configuration object containing the settings for the MegatronBertIntermediate. It should have the following attributes:

  • hidden_size (int): The size of the hidden layer.
  • intermediate_size (int): The size of the intermediate layer.

  • hidden_act (str or function): The activation function for the hidden layer.

    • If it is a string, it should be one of the predefined activation functions available in the ACT2FN dictionary.
    • If it is a function, it should be a custom activation function.

TYPE: object

RETURNS DESCRIPTION

None.
Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
def __init__(self, config):
    """
    Initializes an instance of the MegatronBertIntermediate class.

    Args:
        self: The instance of the class.
        config (object): The configuration object containing the settings for the MegatronBertIntermediate.
            It should have the following attributes:

            - hidden_size (int): The size of the hidden layer.
            - intermediate_size (int): The size of the intermediate layer.
            - hidden_act (str or function): The activation function for the hidden layer.

                - If it is a string, it should be one of the predefined activation functions available in the
                ACT2FN dictionary.
                - If it is a function, it should be a custom activation function.

    Returns:
        None.

    Raises:
        None.
    """
    super().__init__()
    self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
    if isinstance(config.hidden_act, str):
        self.intermediate_act_fn = ACT2FN[config.hidden_act]
    else:
        self.intermediate_act_fn = config.hidden_act

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertIntermediate.forward(hidden_states)

Constructs the intermediate layer of the Megatron BERT model.

PARAMETER DESCRIPTION
self

An instance of the MegatronBertIntermediate class.

TYPE: MegatronBertIntermediate

hidden_states

The input hidden states tensor.

TYPE: Tensor

RETURNS DESCRIPTION
Tensor

mindspore.Tensor: The output hidden states tensor after applying the intermediate layer.
Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
def forward(self, hidden_states: mindspore.Tensor) -> mindspore.Tensor:
    """
    Constructs the intermediate layer of the Megatron BERT model.

    Args:
        self (MegatronBertIntermediate): An instance of the MegatronBertIntermediate class.
        hidden_states (mindspore.Tensor): The input hidden states tensor.

    Returns:
        mindspore.Tensor: The output hidden states tensor after applying the intermediate layer.

    Raises:
        None.
    """
    hidden_states = self.dense(hidden_states)
    hidden_states = self.intermediate_act_fn(hidden_states)
    return hidden_states

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertLMPredictionHead

Bases: Module

MegatronBertLMPredictionHead

This class represents the prediction head for the Megatron-BERT language model. It is responsible for transforming the hidden states and generating predictions for the next token in a sequence.

This class inherits from the nn.Module class.

ATTRIBUTE DESCRIPTION
transform

An instance of the MegatronBertPredictionHeadTransform class, used to transform the hidden states.

TYPE: MegatronBertPredictionHeadTransform

decoder

A fully connected layer that maps the transformed hidden states to the vocabulary size.

TYPE: Linear

bias

A learnable bias parameter used in the decoder layer.

TYPE: Parameter

METHOD DESCRIPTION
forward

Transforms the input hidden states and generates predictions for the next token in the sequence.
Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
class MegatronBertLMPredictionHead(nn.Module):

    """MegatronBertLMPredictionHead

    This class represents the prediction head for the Megatron-BERT language model. It is responsible for
    transforming the hidden states and generating predictions for the next token in a sequence.

    This class inherits from the nn.Module class.

    Attributes:
        transform (MegatronBertPredictionHeadTransform): An instance of the MegatronBertPredictionHeadTransform class,
            used to transform the hidden states.
        decoder (nn.Linear): A fully connected layer that maps the transformed hidden states to the vocabulary size.
        bias (Parameter): A learnable bias parameter used in the decoder layer.

    Methods:
        forward(hidden_states): Transforms the input hidden states and generates predictions for the next token
            in the sequence.

    """
    def __init__(self, config):
        """
        Initialize the MegatronBertLMPredictionHead object with the provided configuration.

        Args:
            self (object): The instance of the class.
            config (object): An object containing configuration parameters for the prediction head.
                It is expected to have attributes like 'hidden_size' and 'vocab_size' required for initialization.

        Returns:
            None.

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

        # The output weights are the same as the input embeddings, but there is
        # an output-only bias for each token.
        self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

        self.bias = Parameter(ops.zeros(config.vocab_size), 'bias')

        # Need a link between the two variables so that the bias is correctly resized with `resize_token_embeddings`
        self.decoder.bias = self.bias

    def forward(self, hidden_states):
        """
        Constructs the MegatronBertLMPredictionHead.

        Args:
            self (MegatronBertLMPredictionHead): The instance of the MegatronBertLMPredictionHead class.
            hidden_states (Tensor): The input hidden states to be processed. It should be a tensor of shape
                (batch_size, sequence_length, hidden_size).

        Returns:
            hidden_states (Tensor): The processed hidden states. It is a tensor of shape
                (batch_size, sequence_length, hidden_size) after applying the transformation and decoding.

        Raises:
            None.
        """
        hidden_states = self.transform(hidden_states)
        hidden_states = self.decoder(hidden_states)
        return hidden_states

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertLMPredictionHead.__init__(config)

Initialize the MegatronBertLMPredictionHead object with the provided configuration.

PARAMETER DESCRIPTION
self

The instance of the class.

TYPE: object

config

An object containing configuration parameters for the prediction head. It is expected to have attributes like 'hidden_size' and 'vocab_size' required for initialization.

TYPE: object

RETURNS DESCRIPTION

None.
Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
def __init__(self, config):
    """
    Initialize the MegatronBertLMPredictionHead object with the provided configuration.

    Args:
        self (object): The instance of the class.
        config (object): An object containing configuration parameters for the prediction head.
            It is expected to have attributes like 'hidden_size' and 'vocab_size' required for initialization.

    Returns:
        None.

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

    # The output weights are the same as the input embeddings, but there is
    # an output-only bias for each token.
    self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

    self.bias = Parameter(ops.zeros(config.vocab_size), 'bias')

    # Need a link between the two variables so that the bias is correctly resized with `resize_token_embeddings`
    self.decoder.bias = self.bias

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertLMPredictionHead.forward(hidden_states)

Constructs the MegatronBertLMPredictionHead.

PARAMETER DESCRIPTION
self

The instance of the MegatronBertLMPredictionHead class.

TYPE: MegatronBertLMPredictionHead

hidden_states

The input hidden states to be processed. It should be a tensor of shape (batch_size, sequence_length, hidden_size).

TYPE: Tensor

RETURNS DESCRIPTION
hidden_states

The processed hidden states. It is a tensor of shape (batch_size, sequence_length, hidden_size) after applying the transformation and decoding.

TYPE: Tensor

Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
def forward(self, hidden_states):
    """
    Constructs the MegatronBertLMPredictionHead.

    Args:
        self (MegatronBertLMPredictionHead): The instance of the MegatronBertLMPredictionHead class.
        hidden_states (Tensor): The input hidden states to be processed. It should be a tensor of shape
            (batch_size, sequence_length, hidden_size).

    Returns:
        hidden_states (Tensor): The processed hidden states. It is a tensor of shape
            (batch_size, sequence_length, hidden_size) after applying the transformation and decoding.

    Raises:
        None.
    """
    hidden_states = self.transform(hidden_states)
    hidden_states = self.decoder(hidden_states)
    return hidden_states

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertLayer

Bases: Module

This class represents a layer of the Megatron-Bert model. It is used to perform attention and feed-forward operations on input hidden states.

ATTRIBUTE DESCRIPTION
chunk_size_feed_forward

The chunk size used for chunking the feed-forward operation.

TYPE: int

seq_len_dim

The dimension of the sequence length.

TYPE: int

attention

The attention module used for self-attention.

TYPE: MegatronBertAttention

is_decoder

Indicates whether the layer is used as a decoder model.

TYPE: bool

add_cross_attention

Indicates whether cross-attention is added.

TYPE: bool

crossattention

The attention module used for cross-attention if add_cross_attention is True.

TYPE: MegatronBertAttention

ln

The layer normalization module.

TYPE: LayerNorm

intermediate

The intermediate module used for the feed-forward operation.

TYPE: MegatronBertIntermediate

output

The output module used for the feed-forward operation.

TYPE: MegatronBertOutput

METHOD DESCRIPTION
feed_forward_chunk

Applies the feed-forward operation to the attention output.

Args:

  • attention_output (mindspore.Tensor): The attention output.

Returns:

  • mindspore.Tensor: The output of the feed-forward operation.
Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
class MegatronBertLayer(nn.Module):

    """
    This class represents a layer of the Megatron-Bert model. It is used to perform attention and feed-forward
    operations on input hidden states.

    Attributes:
        chunk_size_feed_forward (int): The chunk size used for chunking the feed-forward operation.
        seq_len_dim (int): The dimension of the sequence length.
        attention (MegatronBertAttention): The attention module used for self-attention.
        is_decoder (bool): Indicates whether the layer is used as a decoder model.
        add_cross_attention (bool): Indicates whether cross-attention is added.
        crossattention (MegatronBertAttention): The attention module used for cross-attention
            if add_cross_attention is True.
        ln (nn.LayerNorm): The layer normalization module.
        intermediate (MegatronBertIntermediate): The intermediate module used for the feed-forward operation.
        output (MegatronBertOutput): The output module used for the feed-forward operation.

    Methods:
        forward(hidden_states, attention_mask=None, head_mask=None, encoder_hidden_states=None,
            encoder_attention_mask=None, past_key_value=None, output_attentions=False):
            Constructs the layer by performing attention and feed-forward operations on the input hidden states.

            Args:

            - hidden_states (mindspore.Tensor): The input hidden states.
            - attention_mask (mindspore.Tensor, optional): The attention mask tensor. Defaults to None.
            - head_mask (mindspore.Tensor, optional): The head mask tensor. Defaults to None.
            - encoder_hidden_states (mindspore.Tensor, optional): The hidden states of the encoder if the layer
            is used as a decoder model. Defaults to None.
            - encoder_attention_mask (mindspore.Tensor, optional): The attention mask of the encoder if the layer
            is used as a decoder model. Defaults to None.
            - past_key_value (Tuple[Tuple[mindspore.Tensor]], optional): The past key-value pairs for caching
            attention outputs. Defaults to None.
            - output_attentions (bool, optional): Whether to output attention scores. Defaults to False.

            Returns:

            - Tuple[mindspore.Tensor]: The outputs of the layer.

        feed_forward_chunk(attention_output):
            Applies the feed-forward operation to the attention output.

            Args:

            - attention_output (mindspore.Tensor): The attention output.

            Returns:

            - mindspore.Tensor: The output of the feed-forward operation.
    """
    def __init__(self, config):
        """Initializes an instance of the MegatronBertLayer class.

        Args:
            self: An instance of the MegatronBertLayer class.
            config:
                A configuration object containing the following attributes:

                - chunk_size_feed_forward: An integer indicating the chunk size for feedforward layers.
                - is_decoder: A boolean indicating whether the layer is a decoder.
                - add_cross_attention: A boolean indicating whether to add cross attention to the layer.
                - hidden_size: An integer indicating the size of the hidden layer.
                - layer_norm_eps: A float indicating the epsilon value for layer normalization.

        Returns:
            None.

        Raises:
            TypeError: If add_cross_attention is True and is_decoder is False.
        """
        super().__init__()
        self.chunk_size_feed_forward = config.chunk_size_feed_forward
        self.seq_len_dim = 1
        self.attention = MegatronBertAttention(config)
        self.is_decoder = config.is_decoder
        self.add_cross_attention = config.add_cross_attention
        if self.add_cross_attention:
            if not self.is_decoder:
                raise TypeError(f"{self} should be used as a decoder model if cross attention is added")
            self.crossattention = MegatronBertAttention(config)
        self.ln = nn.LayerNorm([config.hidden_size], eps=config.layer_norm_eps)
        self.intermediate = MegatronBertIntermediate(config)
        self.output = MegatronBertOutput(config)

    def forward(
        self,
        hidden_states: mindspore.Tensor,
        attention_mask: Optional[mindspore.Tensor] = None,
        head_mask: Optional[mindspore.Tensor] = None,
        encoder_hidden_states: Optional[mindspore.Tensor] = None,
        encoder_attention_mask: Optional[mindspore.Tensor] = None,
        past_key_value: Optional[Tuple[Tuple[mindspore.Tensor]]] = None,
        output_attentions: Optional[bool] = False,
    ) -> Tuple[mindspore.Tensor]:
        """
        Constructs a MegatronBertLayer.

        This method performs the forward pass of a MegatronBertLayer. It takes in various input tensors and returns
        the outputs after applying self-attention and cross-attention mechanisms, as well as feed-forward layers.

        Args:
            self (MegatronBertLayer): An instance of the MegatronBertLayer class.
            hidden_states (mindspore.Tensor): The input hidden states tensor of shape
                (batch_size, seq_length, hidden_size).
            attention_mask (Optional[mindspore.Tensor]): An optional attention mask tensor of shape
                (batch_size, seq_length) where 1s indicate tokens to attend to and 0s indicate tokens to mask.
            head_mask (Optional[mindspore.Tensor]): An optional head mask tensor of shape (num_heads,) or
                (num_layers, num_heads) where 1s indicate heads to keep and 0s indicate heads to mask.
            encoder_hidden_states (Optional[mindspore.Tensor]): An optional tensor of shape
                (batch_size, seq_length, hidden_size) representing the hidden states of the encoder.
            encoder_attention_mask (Optional[mindspore.Tensor]): An optional attention mask tensor of shape
                (batch_size, seq_length) for the encoder.
            past_key_value (Optional[Tuple[Tuple[mindspore.Tensor]]]): An optional tuple of past key-value tensors
                for self-attention and cross-attention.
            output_attentions (Optional[bool]): An optional flag indicating whether to output attentions.

        Returns:
            Tuple[mindspore.Tensor]: A tuple containing the outputs of the MegatronBertLayer.
                The first element is the layer output tensor of shape (batch_size, seq_length, hidden_size).
                If the layer is a decoder, the tuple also contains the present key-value tensor of shape
                (2, batch_size, num_heads, seq_length, hidden_size).

        Raises:
            AttributeError: If `encoder_hidden_states` are passed and cross-attention layers are not instantiated
                by setting `config.add_cross_attention=True`.
        """
        # decoder uni-directional self-attention cached key/values tuple is at positions 1,2
        self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None
        self_attention_outputs = self.attention(
            hidden_states,
            attention_mask,
            head_mask,
            output_attentions=output_attentions,
            past_key_value=self_attn_past_key_value,
        )
        attention_output = self_attention_outputs[0]

        # if decoder, the last output is tuple of self-attn cache
        if self.is_decoder:
            outputs = self_attention_outputs[1:-1]
            present_key_value = self_attention_outputs[-1]
        else:
            outputs = self_attention_outputs[1:]  # add self attentions if we output attention weights

        cross_attn_present_key_value = None
        if self.is_decoder and encoder_hidden_states is not None:
            if not hasattr(self, "crossattention"):
                raise AttributeError(
                    f"If `encoder_hidden_states` are passed, {self} has to be instantiated with cross-attention layers"
                    " by setting `config.add_cross_attention=True`"
                )

            # cross_attn cached key/values tuple is at positions 3,4 of past_key_value tuple
            cross_attn_past_key_value = past_key_value[-2:] if past_key_value is not None else None
            cross_attention_outputs = self.crossattention(
                attention_output,
                attention_mask,
                head_mask,
                encoder_hidden_states,
                encoder_attention_mask,
                cross_attn_past_key_value,
                output_attentions,
            )
            attention_output = cross_attention_outputs[0]
            outputs = outputs + cross_attention_outputs[1:-1]  # add cross attentions if we output attention weights

            # add cross-attn cache to positions 3,4 of present_key_value tuple
            cross_attn_present_key_value = cross_attention_outputs[-1]
            present_key_value = present_key_value + cross_attn_present_key_value

        layer_output = apply_chunking_to_forward(
            self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output
        )
        outputs = (layer_output,) + outputs

        # if decoder, return the attn key/values as the last output
        if self.is_decoder:
            outputs = outputs + (present_key_value,)

        return outputs

    def feed_forward_chunk(self, attention_output):
        """
        Feed forward chunk of the MegatronBertLayer class.

        This method applies feed forward operations to the attention_output tensor.

        Args:
            self (MegatronBertLayer): An instance of the MegatronBertLayer class.
            attention_output (Tensor): The input tensor to be processed. It represents the attention output.

        Returns:
            None.

        Raises:
            None.

        """
        ln_output = self.ln(attention_output)
        intermediate_output = self.intermediate(ln_output)
        layer_output = self.output(intermediate_output, attention_output)
        return layer_output

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertLayer.__init__(config)

Initializes an instance of the MegatronBertLayer class.

PARAMETER DESCRIPTION
self

An instance of the MegatronBertLayer class.
config

A configuration object containing the following attributes:

  • chunk_size_feed_forward: An integer indicating the chunk size for feedforward layers.
  • is_decoder: A boolean indicating whether the layer is a decoder.
  • add_cross_attention: A boolean indicating whether to add cross attention to the layer.
  • hidden_size: An integer indicating the size of the hidden layer.
  • layer_norm_eps: A float indicating the epsilon value for layer normalization.
RETURNS DESCRIPTION

None.
RAISES DESCRIPTION
TypeError

If add_cross_attention is True and is_decoder is False.
Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
def __init__(self, config):
    """Initializes an instance of the MegatronBertLayer class.

    Args:
        self: An instance of the MegatronBertLayer class.
        config:
            A configuration object containing the following attributes:

            - chunk_size_feed_forward: An integer indicating the chunk size for feedforward layers.
            - is_decoder: A boolean indicating whether the layer is a decoder.
            - add_cross_attention: A boolean indicating whether to add cross attention to the layer.
            - hidden_size: An integer indicating the size of the hidden layer.
            - layer_norm_eps: A float indicating the epsilon value for layer normalization.

    Returns:
        None.

    Raises:
        TypeError: If add_cross_attention is True and is_decoder is False.
    """
    super().__init__()
    self.chunk_size_feed_forward = config.chunk_size_feed_forward
    self.seq_len_dim = 1
    self.attention = MegatronBertAttention(config)
    self.is_decoder = config.is_decoder
    self.add_cross_attention = config.add_cross_attention
    if self.add_cross_attention:
        if not self.is_decoder:
            raise TypeError(f"{self} should be used as a decoder model if cross attention is added")
        self.crossattention = MegatronBertAttention(config)
    self.ln = nn.LayerNorm([config.hidden_size], eps=config.layer_norm_eps)
    self.intermediate = MegatronBertIntermediate(config)
    self.output = MegatronBertOutput(config)

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertLayer.feed_forward_chunk(attention_output)

Feed forward chunk of the MegatronBertLayer class.

This method applies feed forward operations to the attention_output tensor.

PARAMETER DESCRIPTION
self

An instance of the MegatronBertLayer class.

TYPE: MegatronBertLayer

attention_output

The input tensor to be processed. It represents the attention output.

TYPE: Tensor

RETURNS DESCRIPTION

None.
Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
def feed_forward_chunk(self, attention_output):
    """
    Feed forward chunk of the MegatronBertLayer class.

    This method applies feed forward operations to the attention_output tensor.

    Args:
        self (MegatronBertLayer): An instance of the MegatronBertLayer class.
        attention_output (Tensor): The input tensor to be processed. It represents the attention output.

    Returns:
        None.

    Raises:
        None.

    """
    ln_output = self.ln(attention_output)
    intermediate_output = self.intermediate(ln_output)
    layer_output = self.output(intermediate_output, attention_output)
    return layer_output

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertLayer.forward(hidden_states, attention_mask=None, head_mask=None, encoder_hidden_states=None, encoder_attention_mask=None, past_key_value=None, output_attentions=False)

Constructs a MegatronBertLayer.

This method performs the forward pass of a MegatronBertLayer. It takes in various input tensors and returns the outputs after applying self-attention and cross-attention mechanisms, as well as feed-forward layers.

PARAMETER DESCRIPTION
self

An instance of the MegatronBertLayer class.

TYPE: MegatronBertLayer

hidden_states

The input hidden states tensor of shape (batch_size, seq_length, hidden_size).

TYPE: Tensor

attention_mask

An optional attention mask tensor of shape (batch_size, seq_length) where 1s indicate tokens to attend to and 0s indicate tokens to mask.

TYPE: Optional[Tensor] DEFAULT: None

head_mask

An optional head mask tensor of shape (num_heads,) or (num_layers, num_heads) where 1s indicate heads to keep and 0s indicate heads to mask.

TYPE: Optional[Tensor] DEFAULT: None

encoder_hidden_states

An optional tensor of shape (batch_size, seq_length, hidden_size) representing the hidden states of the encoder.

TYPE: Optional[Tensor] DEFAULT: None

encoder_attention_mask

An optional attention mask tensor of shape (batch_size, seq_length) for the encoder.

TYPE: Optional[Tensor] DEFAULT: None

past_key_value

An optional tuple of past key-value tensors for self-attention and cross-attention.

TYPE: Optional[Tuple[Tuple[Tensor]]] DEFAULT: None

output_attentions

An optional flag indicating whether to output attentions.

TYPE: Optional[bool] DEFAULT: False

RETURNS DESCRIPTION
Tuple[Tensor]

Tuple[mindspore.Tensor]: A tuple containing the outputs of the MegatronBertLayer. The first element is the layer output tensor of shape (batch_size, seq_length, hidden_size). If the layer is a decoder, the tuple also contains the present key-value tensor of shape (2, batch_size, num_heads, seq_length, hidden_size).
RAISES DESCRIPTION
AttributeError

If encoder_hidden_states are passed and cross-attention layers are not instantiated by setting config.add_cross_attention=True.
Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
def forward(
    self,
    hidden_states: mindspore.Tensor,
    attention_mask: Optional[mindspore.Tensor] = None,
    head_mask: Optional[mindspore.Tensor] = None,
    encoder_hidden_states: Optional[mindspore.Tensor] = None,
    encoder_attention_mask: Optional[mindspore.Tensor] = None,
    past_key_value: Optional[Tuple[Tuple[mindspore.Tensor]]] = None,
    output_attentions: Optional[bool] = False,
) -> Tuple[mindspore.Tensor]:
    """
    Constructs a MegatronBertLayer.

    This method performs the forward pass of a MegatronBertLayer. It takes in various input tensors and returns
    the outputs after applying self-attention and cross-attention mechanisms, as well as feed-forward layers.

    Args:
        self (MegatronBertLayer): An instance of the MegatronBertLayer class.
        hidden_states (mindspore.Tensor): The input hidden states tensor of shape
            (batch_size, seq_length, hidden_size).
        attention_mask (Optional[mindspore.Tensor]): An optional attention mask tensor of shape
            (batch_size, seq_length) where 1s indicate tokens to attend to and 0s indicate tokens to mask.
        head_mask (Optional[mindspore.Tensor]): An optional head mask tensor of shape (num_heads,) or
            (num_layers, num_heads) where 1s indicate heads to keep and 0s indicate heads to mask.
        encoder_hidden_states (Optional[mindspore.Tensor]): An optional tensor of shape
            (batch_size, seq_length, hidden_size) representing the hidden states of the encoder.
        encoder_attention_mask (Optional[mindspore.Tensor]): An optional attention mask tensor of shape
            (batch_size, seq_length) for the encoder.
        past_key_value (Optional[Tuple[Tuple[mindspore.Tensor]]]): An optional tuple of past key-value tensors
            for self-attention and cross-attention.
        output_attentions (Optional[bool]): An optional flag indicating whether to output attentions.

    Returns:
        Tuple[mindspore.Tensor]: A tuple containing the outputs of the MegatronBertLayer.
            The first element is the layer output tensor of shape (batch_size, seq_length, hidden_size).
            If the layer is a decoder, the tuple also contains the present key-value tensor of shape
            (2, batch_size, num_heads, seq_length, hidden_size).

    Raises:
        AttributeError: If `encoder_hidden_states` are passed and cross-attention layers are not instantiated
            by setting `config.add_cross_attention=True`.
    """
    # decoder uni-directional self-attention cached key/values tuple is at positions 1,2
    self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None
    self_attention_outputs = self.attention(
        hidden_states,
        attention_mask,
        head_mask,
        output_attentions=output_attentions,
        past_key_value=self_attn_past_key_value,
    )
    attention_output = self_attention_outputs[0]

    # if decoder, the last output is tuple of self-attn cache
    if self.is_decoder:
        outputs = self_attention_outputs[1:-1]
        present_key_value = self_attention_outputs[-1]
    else:
        outputs = self_attention_outputs[1:]  # add self attentions if we output attention weights

    cross_attn_present_key_value = None
    if self.is_decoder and encoder_hidden_states is not None:
        if not hasattr(self, "crossattention"):
            raise AttributeError(
                f"If `encoder_hidden_states` are passed, {self} has to be instantiated with cross-attention layers"
                " by setting `config.add_cross_attention=True`"
            )

        # cross_attn cached key/values tuple is at positions 3,4 of past_key_value tuple
        cross_attn_past_key_value = past_key_value[-2:] if past_key_value is not None else None
        cross_attention_outputs = self.crossattention(
            attention_output,
            attention_mask,
            head_mask,
            encoder_hidden_states,
            encoder_attention_mask,
            cross_attn_past_key_value,
            output_attentions,
        )
        attention_output = cross_attention_outputs[0]
        outputs = outputs + cross_attention_outputs[1:-1]  # add cross attentions if we output attention weights

        # add cross-attn cache to positions 3,4 of present_key_value tuple
        cross_attn_present_key_value = cross_attention_outputs[-1]
        present_key_value = present_key_value + cross_attn_present_key_value

    layer_output = apply_chunking_to_forward(
        self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output
    )
    outputs = (layer_output,) + outputs

    # if decoder, return the attn key/values as the last output
    if self.is_decoder:
        outputs = outputs + (present_key_value,)

    return outputs

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertModel

Bases: MegatronBertPreTrainedModel

The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of cross-attention is added between the self-attention layers, following the architecture described in Attention is all you need by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin.

To behave as an decoder the model needs to be initialized with the is_decoder argument of the configuration set to True. To be used in a Seq2Seq model, the model needs to initialized with both is_decoder argument and add_cross_attention set to True; an encoder_hidden_states is then expected as an input to the forward pass.

Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
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
1706
1707
1708
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
class MegatronBertModel(MegatronBertPreTrainedModel):
    """

    The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of
    cross-attention is added between the self-attention layers, following the architecture described in [Attention is
    all you need](https://arxiv.org/abs/1706.03762) by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit,
    Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin.

    To behave as an decoder the model needs to be initialized with the `is_decoder` argument of the configuration set
    to `True`. To be used in a Seq2Seq model, the model needs to initialized with both `is_decoder` argument and
    `add_cross_attention` set to `True`; an `encoder_hidden_states` is then expected as an input to the forward pass.
    """
    def __init__(self, config, add_pooling_layer=True):
        """
        __init__ method in the MegatronBertModel class.

        Args:
            self: The instance of the class.
            config: A dictionary containing configuration parameters for the MegatronBertModel.
                It is used to initialize the model's embeddings, encoder, and pooler.
            add_pooling_layer: A boolean flag indicating whether to add a pooling layer to the model.
                Default is True.

        Returns:
            None.

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

        self.embeddings = MegatronBertEmbeddings(config)
        self.encoder = MegatronBertEncoder(config)

        self.pooler = MegatronBertPooler(config) if add_pooling_layer else None

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

    def get_input_embeddings(self):
        """
        Method: get_input_embeddings

        Description:
        This method returns the word embeddings used for input in a MegatronBertModel instance.

        Args:
            self (MegatronBertModel): The instance of the MegatronBertModel class.

        Returns:
            None.

        Raises:
            None.

        """
        return self.embeddings.word_embeddings

    def set_input_embeddings(self, value):
        """
        Sets the input embeddings for the MegatronBertModel instance.

        Args:
            self (MegatronBertModel): The instance of the MegatronBertModel class.
            value: The new input embeddings to be set for the model. Should be of type torch.Tensor.

        Returns:
            None.

        Raises:
            None.
        """
        self.embeddings.word_embeddings = value

    def _prune_heads(self, heads_to_prune):
        """
        Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
        class PreTrainedModel
        """
        for layer, heads in heads_to_prune.items():
            self.encoder.layer[layer].attention.prune_heads(heads)

    def forward(
        self,
        input_ids: Optional[mindspore.Tensor] = None,
        attention_mask: Optional[mindspore.Tensor] = None,
        token_type_ids: Optional[mindspore.Tensor] = None,
        position_ids: Optional[mindspore.Tensor] = None,
        head_mask: Optional[mindspore.Tensor] = None,
        inputs_embeds: Optional[mindspore.Tensor] = None,
        encoder_hidden_states: Optional[mindspore.Tensor] = None,
        encoder_attention_mask: Optional[mindspore.Tensor] = None,
        past_key_values: Optional[Tuple[Tuple[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, BaseModelOutputWithPoolingAndCrossAttentions]:
        r"""
        Args:
            encoder_hidden_states  (`mindspore.Tensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
                Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
                the model is configured as a decoder.
            encoder_attention_mask (`mindspore.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
                Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
                the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`:

                - 1 for tokens that are **not masked**,
                - 0 for tokens that are **masked**.
            past_key_values (`tuple(tuple(mindspore.Tensor))` of length `config.n_layers` with each tuple having
                4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
                Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.

                If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
                don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
                `decoder_input_ids` of shape `(batch_size, sequence_length)`.
            use_cache (`bool`, *optional*):
                If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
                (see `past_key_values`).
        """
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        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

        if self.config.is_decoder:
            use_cache = use_cache if use_cache is not None else self.config.use_cache
        else:
            use_cache = False

        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")
        if input_ids is not None:
            self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
            input_shape = input_ids.shape
        elif inputs_embeds is not None:
            input_shape = inputs_embeds.shape[:-1]
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds")

        batch_size, seq_length = input_shape

        # past_key_values_length
        past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0

        if attention_mask is None:
            attention_mask = ops.ones(((batch_size, seq_length + past_key_values_length)))
        if token_type_ids is None:
            token_type_ids = ops.zeros(input_shape, dtype=mindspore.int64)

        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
        # ourselves in which case we just need to make it broadcastable to all heads.
        extended_attention_mask: mindspore.Tensor = self.get_extended_attention_mask(attention_mask, input_shape)

        # If a 2D or 3D attention mask is provided for the cross-attention
        # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
        if self.config.is_decoder and encoder_hidden_states is not None:
            encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.shape
            encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
            if encoder_attention_mask is None:
                encoder_attention_mask = ops.ones(encoder_hidden_shape)
            encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
        else:
            encoder_extended_attention_mask = None

        # Prepare head mask if needed
        # 1.0 in head_mask indicate we keep the head
        # attention_probs has shape bsz x n_heads x N x N
        # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
        # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
        head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)

        embedding_output = self.embeddings(
            input_ids=input_ids,
            position_ids=position_ids,
            token_type_ids=token_type_ids,
            inputs_embeds=inputs_embeds,
            past_key_values_length=past_key_values_length,
        )
        encoder_outputs = self.encoder(
            embedding_output,
            attention_mask=extended_attention_mask,
            head_mask=head_mask,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_extended_attention_mask,
            past_key_values=past_key_values,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        sequence_output = encoder_outputs[0]
        pooled_output = self.pooler(sequence_output) if self.pooler is not None else None

        if not return_dict:
            return (sequence_output, pooled_output) + encoder_outputs[1:]

        return BaseModelOutputWithPoolingAndCrossAttentions(
            last_hidden_state=sequence_output,
            pooler_output=pooled_output,
            past_key_values=encoder_outputs.past_key_values,
            hidden_states=encoder_outputs.hidden_states,
            attentions=encoder_outputs.attentions,
            cross_attentions=encoder_outputs.cross_attentions,
        )

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertModel.__init__(config, add_pooling_layer=True)

init method in the MegatronBertModel class.

PARAMETER DESCRIPTION
self

The instance of the class.
config

A dictionary containing configuration parameters for the MegatronBertModel. It is used to initialize the model's embeddings, encoder, and pooler.
add_pooling_layer

A boolean flag indicating whether to add a pooling layer to the model. Default is True.

DEFAULT: True

RETURNS DESCRIPTION

None.
Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
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
def __init__(self, config, add_pooling_layer=True):
    """
    __init__ method in the MegatronBertModel class.

    Args:
        self: The instance of the class.
        config: A dictionary containing configuration parameters for the MegatronBertModel.
            It is used to initialize the model's embeddings, encoder, and pooler.
        add_pooling_layer: A boolean flag indicating whether to add a pooling layer to the model.
            Default is True.

    Returns:
        None.

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

    self.embeddings = MegatronBertEmbeddings(config)
    self.encoder = MegatronBertEncoder(config)

    self.pooler = MegatronBertPooler(config) if add_pooling_layer else None

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

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertModel.forward(input_ids=None, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None, inputs_embeds=None, encoder_hidden_states=None, encoder_attention_mask=None, past_key_values=None, use_cache=None, output_attentions=None, output_hidden_states=None, return_dict=None)

PARAMETER DESCRIPTION
encoder_hidden_states

Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if the model is configured as a decoder.

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

encoder_attention_mask

Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in the cross-attention if the model is configured as a decoder. Mask values selected in [0, 1]:

  • 1 for tokens that are not masked,
  • 0 for tokens that are masked.

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

use_cache

If set to True, past_key_values key value states are returned and can be used to speed up decoding (see past_key_values).

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

Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
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
1706
1707
1708
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
def forward(
    self,
    input_ids: Optional[mindspore.Tensor] = None,
    attention_mask: Optional[mindspore.Tensor] = None,
    token_type_ids: Optional[mindspore.Tensor] = None,
    position_ids: Optional[mindspore.Tensor] = None,
    head_mask: Optional[mindspore.Tensor] = None,
    inputs_embeds: Optional[mindspore.Tensor] = None,
    encoder_hidden_states: Optional[mindspore.Tensor] = None,
    encoder_attention_mask: Optional[mindspore.Tensor] = None,
    past_key_values: Optional[Tuple[Tuple[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, BaseModelOutputWithPoolingAndCrossAttentions]:
    r"""
    Args:
        encoder_hidden_states  (`mindspore.Tensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
            Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
            the model is configured as a decoder.
        encoder_attention_mask (`mindspore.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
            Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
            the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`:

            - 1 for tokens that are **not masked**,
            - 0 for tokens that are **masked**.
        past_key_values (`tuple(tuple(mindspore.Tensor))` of length `config.n_layers` with each tuple having
            4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
            Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.

            If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
            don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
            `decoder_input_ids` of shape `(batch_size, sequence_length)`.
        use_cache (`bool`, *optional*):
            If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
            (see `past_key_values`).
    """
    output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
    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

    if self.config.is_decoder:
        use_cache = use_cache if use_cache is not None else self.config.use_cache
    else:
        use_cache = False

    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")
    if input_ids is not None:
        self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
        input_shape = input_ids.shape
    elif inputs_embeds is not None:
        input_shape = inputs_embeds.shape[:-1]
    else:
        raise ValueError("You have to specify either input_ids or inputs_embeds")

    batch_size, seq_length = input_shape

    # past_key_values_length
    past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0

    if attention_mask is None:
        attention_mask = ops.ones(((batch_size, seq_length + past_key_values_length)))
    if token_type_ids is None:
        token_type_ids = ops.zeros(input_shape, dtype=mindspore.int64)

    # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
    # ourselves in which case we just need to make it broadcastable to all heads.
    extended_attention_mask: mindspore.Tensor = self.get_extended_attention_mask(attention_mask, input_shape)

    # If a 2D or 3D attention mask is provided for the cross-attention
    # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
    if self.config.is_decoder and encoder_hidden_states is not None:
        encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.shape
        encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
        if encoder_attention_mask is None:
            encoder_attention_mask = ops.ones(encoder_hidden_shape)
        encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
    else:
        encoder_extended_attention_mask = None

    # Prepare head mask if needed
    # 1.0 in head_mask indicate we keep the head
    # attention_probs has shape bsz x n_heads x N x N
    # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
    # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
    head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)

    embedding_output = self.embeddings(
        input_ids=input_ids,
        position_ids=position_ids,
        token_type_ids=token_type_ids,
        inputs_embeds=inputs_embeds,
        past_key_values_length=past_key_values_length,
    )
    encoder_outputs = self.encoder(
        embedding_output,
        attention_mask=extended_attention_mask,
        head_mask=head_mask,
        encoder_hidden_states=encoder_hidden_states,
        encoder_attention_mask=encoder_extended_attention_mask,
        past_key_values=past_key_values,
        use_cache=use_cache,
        output_attentions=output_attentions,
        output_hidden_states=output_hidden_states,
        return_dict=return_dict,
    )
    sequence_output = encoder_outputs[0]
    pooled_output = self.pooler(sequence_output) if self.pooler is not None else None

    if not return_dict:
        return (sequence_output, pooled_output) + encoder_outputs[1:]

    return BaseModelOutputWithPoolingAndCrossAttentions(
        last_hidden_state=sequence_output,
        pooler_output=pooled_output,
        past_key_values=encoder_outputs.past_key_values,
        hidden_states=encoder_outputs.hidden_states,
        attentions=encoder_outputs.attentions,
        cross_attentions=encoder_outputs.cross_attentions,
    )

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertModel.get_input_embeddings()

Description: This method returns the word embeddings used for input in a MegatronBertModel instance.

PARAMETER DESCRIPTION
self

The instance of the MegatronBertModel class.

TYPE: MegatronBertModel

RETURNS DESCRIPTION

None.
Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
1634
1635
1636
1637
1638
1639
1640
1641
def get_input_embeddings(self):
    """
    Method: get_input_embeddings

    Description:
    This method returns the word embeddings used for input in a MegatronBertModel instance.

    Args:
        self (MegatronBertModel): The instance of the MegatronBertModel class.

    Returns:
        None.

    Raises:
        None.

    """
    return self.embeddings.word_embeddings

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertModel.set_input_embeddings(value)

Sets the input embeddings for the MegatronBertModel instance.

PARAMETER DESCRIPTION
self

The instance of the MegatronBertModel class.

TYPE: MegatronBertModel

value

The new input embeddings to be set for the model. Should be of type torch.Tensor.
RETURNS DESCRIPTION

None.
Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655
1656
1657
def set_input_embeddings(self, value):
    """
    Sets the input embeddings for the MegatronBertModel instance.

    Args:
        self (MegatronBertModel): The instance of the MegatronBertModel class.
        value: The new input embeddings to be set for the model. Should be of type torch.Tensor.

    Returns:
        None.

    Raises:
        None.
    """
    self.embeddings.word_embeddings = value

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertOnlyMLMHead

Bases: Module

Represents a Megatron-style MLM head for BERT models, which includes only the MLM prediction head without the rest of the model.

This class inherits from nn.Module and is designed to be used in conjunction with a BERT model for masked language modeling tasks. It contains methods for initializing the prediction head and generating prediction scores based on the input sequence output.

The class includes an init method to initialize the prediction head with the provided configuration, and a forward method to generate prediction scores using the sequence output tensor. The prediction scores are obtained by passing the sequence output through the prediction head.

Note

This class assumes that the MegatronBertLMPredictionHead class is available for use in creating the MLM prediction head.

Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
class MegatronBertOnlyMLMHead(nn.Module):

    """
    Represents a Megatron-style MLM head for BERT models, which includes only the MLM prediction head without the rest
    of the model.

    This class inherits from nn.Module and is designed to be used in conjunction with a BERT model for masked language
    modeling tasks. It contains methods for initializing the prediction head and generating prediction scores based on
    the input sequence output.

    The class includes an __init__ method to initialize the prediction head with the provided configuration, and a
    forward method to generate prediction scores using the sequence output tensor. The prediction scores are obtained
    by passing the sequence output through the prediction head.

    Note:
        This class assumes that the MegatronBertLMPredictionHead class is available for use in creating the MLM
        prediction head.

    """
    def __init__(self, config):
        """
        Initialize the MegatronBertOnlyMLMHead class.

        Args:
            self (object): The instance of the class.
            config (object): An object containing configuration settings for the MegatronBertOnlyMLMHead class.

        Returns:
            None.

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

    def forward(self, sequence_output: mindspore.Tensor) -> mindspore.Tensor:
        """
        This method forwards predictions for masked language modeling using the MegatronBertOnlyMLMHead class.

        Args:
            self (MegatronBertOnlyMLMHead): The instance of the MegatronBertOnlyMLMHead class.
            sequence_output (mindspore.Tensor): The output tensor from the previous layer representing the
                input sequence for prediction. This tensor should be compatible with the model architecture and contain
                the necessary information for prediction.

        Returns:
            mindspore.Tensor: A tensor containing the prediction scores generated by the model for masked language modeling.
                The prediction scores represent the likelihood of each token being the correct masked token.

        Raises:
            ValueError: If the input sequence_output is not a valid mindspore.Tensor object.
            RuntimeError: If there are issues during the prediction process within the self.predictions() method.
        """
        prediction_scores = self.predictions(sequence_output)
        return prediction_scores

mindnlp.transformers.models.megatron_bert.modeling_megatron_bert.MegatronBertOnlyMLMHead.__init__(config)

Initialize the MegatronBertOnlyMLMHead class.

PARAMETER DESCRIPTION
self

The instance of the class.

TYPE: object

config

An object containing configuration settings for the MegatronBertOnlyMLMHead class.

TYPE: object

RETURNS DESCRIPTION

None.
Source code in mindnlp/transformers/models/megatron_bert/modeling_megatron_bert.py 
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
def __init__(self, config):
    """
    Initialize the MegatronBertOnlyMLMHead class.

    Args:
        self (object): The instance of the class.
        config (object): An object conta