def forward(
            self,
            input_ids=None,
            attention_mask=None,
            token_type_ids=None,
            position_ids=None,
            head_mask=None,
            inputs_embeds=None,
            labels=None,
            pooling_ids=None,
            noised_input_ids=None,
            noised_attention_mask=None,
            noised_token_type_ids=None,
            noised_labels=None,
            noised_pooling_ids=None,
            noised_r1_mask=None,
            original_r1_mask=None,
            src_pooling_ids=None,
            tgt_pooling_ids=None,
            is_augmented=None,
            first_stage_model_logits=None,
    ):
        r"""
        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
            Labels for computing the token classification loss.
            Indices should be in ``[0, ..., config.num_labels - 1]``.

    Returns:
        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.RobertaConfig`) and inputs:
        loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when ``labels`` is provided) :
            Classification loss.
        scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, config.num_labels)`)
            Classification scores (before SoftMax).
        hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_hidden_states=True``):
            Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(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.

    Examples::

        from transformers import RobertaTokenizer, RobertaForTokenClassification
        import torch

        tokenizer = RobertaTokenizer.from_pretrained('roberta-base')
        model = RobertaForTokenClassification.from_pretrained('roberta-base')
        input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze(0)  # Batch size 1
        labels = torch.tensor([1] * input_ids.size(1)).unsqueeze(0)  # Batch size 1
        outputs = model(input_ids, labels=labels)
        loss, scores = outputs[:2]

        """

        word_embeds = self.roberta.embeddings.word_embeddings(input_ids)

        if is_augmented is not None:
            # ground truth data indices
            # optional when data augmentation is considered noisy
            # if self.augment_method != "mt" and first_stage_model_logits is not None:
            #     gt_indices = (~is_augmented.bool()).view(-1).nonzero(as_tuple=False).view(-1).tolist()
            # else:
            #     gt_indices = list(range(0, input_ids.size(0)))

            gt_indices = list(range(0, input_ids.size(0)))
            # optional to conduct r2 on original corpus
            augmented_indices = is_augmented.view(-1).nonzero(as_tuple=False).view(-1).tolist()
        else:
            gt_indices = list(range(0, input_ids.size(0)))
            augmented_indices = None

        if is_augmented is not None and self.noise_sampler is not None:
            noise = self.noise_sampler.sample(sample_shape=word_embeds.shape).to(word_embeds)
            if self.detach_embeds:
                noised_word_embeds = word_embeds.detach().clone() + noise
            else:
                noised_word_embeds = word_embeds + noise
            if len(augmented_indices) > 0:
                # print(word_embeds[indices].size(), indices, noised_word_embeds[indices].size())
                word_embeds[augmented_indices] = noised_word_embeds[augmented_indices]

        outputs = self.roberta(
            input_ids=None,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=word_embeds.detach().clone() if self.detach_embeds else word_embeds,
        )

        # (batch_size, seq_len, hidden_size)
        sequence_output = outputs[0]

        if self.use_pooling_strategy:
            sum_sequence_output = torch.zeros_like(sequence_output)
            expanded_pooling_ids = pooling_ids.unsqueeze(-1).expand(-1, -1, sequence_output.size(-1))
            sum_sequence_output.scatter_add_(dim=1, index=expanded_pooling_ids, src=sequence_output)

            pooled_sequence_output = self.dropout(sum_sequence_output)
            logits = self.classifier(pooled_sequence_output)
        else:
            sequence_output = self.dropout(sequence_output)
            logits = self.classifier(sequence_output)

        outputs = (logits,) + outputs[2:]  # add hidden states and attention if they are here

        if labels is not None:
            loss_fct = CrossEntropyLoss()
            # Only keep active parts of the loss
            if attention_mask is not None:
                active_loss = attention_mask[gt_indices].view(-1) == 1
                active_logits = logits[gt_indices].view(-1, self.num_labels)
                active_labels = torch.where(
                    active_loss, labels[gt_indices].view(-1), torch.tensor(loss_fct.ignore_index).type_as(labels)
                )
                loss = loss_fct(active_logits, active_labels)
            else:
                loss = loss_fct(logits[gt_indices].view(-1, self.num_labels), labels[gt_indices].view(-1))
            outputs = (loss,) + outputs

            if self.training:
                if first_stage_model_logits is not None and is_augmented is not None and len(augmented_indices) > 0:
                    sum_token_count = torch.zeros_like(pooling_ids)
                    sum_token_count.scatter_add_(dim=1, index=pooling_ids, src=torch.ones_like(pooling_ids))
                    sum_token_count[:, 0] = 0
                    # optional to use hard labels and augmented indices when data augmenation is considered noisy
                    # hard labels for NER task
                    if self.use_hard_labels:
                        hard_labels = first_stage_model_logits.max(dim=-1)[1]
                        active_loss = sum_token_count[augmented_indices].gt(0).view(-1) == 1
                        active_logits = logits[augmented_indices].view(-1, self.num_labels)
                        active_labels = torch.where(active_loss, hard_labels[augmented_indices].view(-1),
                                                    torch.tensor(loss_fct.ignore_index).type_as(labels))
                        r2_loss = loss_fct(active_logits, active_labels) * self.r2_lambda
                    else:
                        sum_mask = sum_token_count[augmented_indices].view(-1).gt(0).unsqueeze(-1).expand(-1,
                                                                                                          self.num_labels)
                        token_teacher_logits = first_stage_model_logits[augmented_indices].view(-1,
                                                                                                self.num_labels).masked_select(
                            sum_mask).view(-1, self.num_labels)
                        token_student_logits = logits[augmented_indices].view(-1, self.num_labels).masked_select(
                            sum_mask).view(-1, self.num_labels)
                        r2_loss = KL(token_student_logits, token_teacher_logits.detach()) * self.r2_lambda
                else:
                    r2_loss = loss.data.new([0.0])

                if self.enable_r1_loss or self.noised_loss:
                    if noised_input_ids is not None:
                        noised_word_embeds = self.roberta.embeddings.word_embeddings(noised_input_ids)
                        assert noised_attention_mask is not None
                    elif self.noise_sampler is not None:
                        noised_word_embeds = self.roberta.embeddings.word_embeddings(input_ids)
                        noised_attention_mask = attention_mask
                        noised_token_type_ids = token_type_ids
                    else:
                        assert False

                    if self.noise_sampler is not None:
                        noise = self.noise_sampler.sample(sample_shape=noised_word_embeds.shape).to(noised_word_embeds)
                        if self.detach_embeds:
                            noised_word_embeds = noised_word_embeds.detach().clone() + noise
                        else:
                            noised_word_embeds = noised_word_embeds + noise
                    else:
                        noised_word_embeds = noised_word_embeds.detach().clone() if self.detach_embeds else noised_word_embeds

                    noised_outputs = self.roberta(
                        input_ids=None,
                        attention_mask=noised_attention_mask,
                        token_type_ids=noised_token_type_ids,
                        position_ids=position_ids,
                        head_mask=head_mask,
                        inputs_embeds=noised_word_embeds,
                    )

                    noised_sequence_output = noised_outputs[0]
                    if self.use_pooling_strategy:
                        noised_sum_token_count = torch.zeros_like(noised_pooling_ids)
                        noised_sum_token_count.scatter_add_(dim=1, index=noised_pooling_ids,
                                                            src=torch.ones_like(noised_pooling_ids))
                        noised_sum_sequence_output = torch.zeros_like(noised_sequence_output)
                        noised_expanded_pooling_ids = noised_pooling_ids.unsqueeze(-1).expand(-1, -1,
                                                                                              noised_sequence_output.size(
                                                                                                  -1))
                        noised_sum_sequence_output.scatter_add_(dim=1, index=noised_expanded_pooling_ids,
                                                                src=noised_sequence_output)

                        pooled_noised_sequence_output = self.dropout(noised_sum_sequence_output)
                        noised_logits = self.classifier(pooled_noised_sequence_output)
                    else:
                        noised_sequence_output = self.dropout(noised_sequence_output)
                        noised_logits = self.classifier(noised_sequence_output)

                if self.original_loss:
                    original_loss = loss
                else:
                    original_loss = loss.data.new([0.0])

                if self.noised_loss:
                    if noised_attention_mask is not None:
                        noised_active_loss = noised_attention_mask.view(-1) == 1
                        noised_active_logits = noised_logits.view(-1, self.num_labels)
                        noised_active_labels = torch.where(
                            noised_active_loss, noised_labels.view(-1),
                            torch.tensor(loss_fct.ignore_index).type_as(noised_labels)
                        )
                        noised_loss = loss_fct(noised_active_logits, noised_active_labels)
                    else:
                        noised_loss = loss_fct(noised_logits.view(-1, self.num_labels), noised_labels.view(-1))
                else:
                    noised_loss = loss.data.new([0.0])

                if self.enable_r1_loss:

                    if self.use_align_label_probs:
                        src_pooling_count, src_probs, src_mask = get_align_probs(logits, src_pooling_ids)
                        tgt_pooling_count, tgt_probs, tgt_mask = get_align_probs(noised_logits, tgt_pooling_ids)

                        assert src_mask.eq(tgt_mask).sum() == src_mask.size(0) * src_mask.size(1)

                        indices = src_mask.view(-1).nonzero(as_tuple=False).view(-1).tolist()

                        src_probs = src_probs.view(-1, src_probs.size(-1))[indices]
                        tgt_probs = tgt_probs.view(-1, src_probs.size(-1))[indices]

                        align_r1_loss_f = KL_probs(src_probs, tgt_probs.detach()) / src_probs.size(0)
                        align_r1_loss_b = KL_probs(tgt_probs, src_probs.detach()) / src_probs.size(0)
                        align_r1_loss = align_r1_loss_b + align_r1_loss_f
                    else:
                        align_r1_loss = loss.data.new([0.0])

                    if self.use_token_label_probs:
                        original_indices = original_r1_mask.view(-1).eq(1).nonzero(as_tuple=False).view(-1).tolist()
                        noised_indices = noised_r1_mask.view(-1).eq(1).nonzero(as_tuple=False).view(-1).tolist()
                        token_original_logits = logits.view(-1, self.num_labels)[original_indices]
                        token_noised_logits = noised_logits.view(-1, self.num_labels)[noised_indices]

                        token_r1_loss_f = KL(token_noised_logits, token_original_logits.detach())
                        token_r1_loss_b = KL(token_original_logits, token_noised_logits.detach())
                        if not self.disable_backward_kl:
                            token_r1_loss = token_r1_loss_f + token_r1_loss_b
                        else:
                            token_r1_loss = token_r1_loss_f
                    else:
                        token_r1_loss = loss.data.new([0.0])

                    if self.use_sentence_label_probs:
                        original_probs = get_label_probs(logits, labels.ne(loss_fct.ignore_index))
                        noised_probs = get_label_probs(noised_logits, noised_labels.ne(loss_fct.ignore_index))

                        sentence_r1_loss_f = KL_probs(noised_probs, original_probs.detach()) / input_ids.size(0)
                        sentence_r1_loss_b = KL_probs(original_probs, noised_probs.detach()) / input_ids.size(0)
                        sentence_r1_loss = sentence_r1_loss_f + sentence_r1_loss_b
                    else:
                        sentence_r1_loss = loss.data.new([0.0])

                    r1_loss = (token_r1_loss + sentence_r1_loss + align_r1_loss) * self.r1_lambda
                else:
                    r1_loss = loss.data.new([0.0])

                loss = original_loss + noised_loss + r1_loss + r2_loss

                # print(loss, original_loss, r1_loss, loss.eq(original_loss), loss - r1_loss)
                outputs = (loss, original_loss, noised_loss, r1_loss, r2_loss) + outputs[1:]

        return outputs  # (loss), scores, (hidden_states), (attentions)

        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
            Labels for computing the token classification loss.
            Indices should be in ``[0, ..., config.num_labels - 1]``.

    Returns:
        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.RobertaConfig`) and inputs:
        loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when ``labels`` is provided) :
            Classification loss.
        scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, config.num_labels)`)
            Classification scores (before SoftMax).
        hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_hidden_states=True``):
            Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(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.

    Examples::

        from transformers import RobertaTokenizer, RobertaForTokenClassification
        import torch

        tokenizer = RobertaTokenizer.from_pretrained('roberta-base')
        model = RobertaForTokenClassification.from_pretrained('roberta-base')
        input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze(0)  # Batch size 1
        labels = torch.tensor([1] * input_ids.size(1)).unsqueeze(0)  # Batch size 1
        outputs = model(input_ids, labels=labels)
        loss, scores = outputs[:2]

        
xtune/src/transformers/modeling_xlm_roberta.py