text: import torch.nn as nn
  from torch.nn import CrossEntropyLoss, MSELoss
  from transformers.models.distilbert.modeling_distilbert import DistilBertModel, DistilBertPreTrainedModel
  
  
  class DistilBertForSequenceClassification(DistilBertPreTrainedModel):
      r"""
          **labels**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:
              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).
      Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
          **loss**: (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
              Classification (or regression if config.num_labels==1) loss.
          **logits**: ``torch.FloatTensor`` of shape ``(batch_size, config.num_labels)``
              Classification (or regression if config.num_labels==1) scores (before SoftMax).
          **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
              list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)
              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**: (`optional`, returned when ``config.output_attentions=True``)
              list of ``torch.FloatTensor`` (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.
      Examples::
          tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')
          model = DistilBertForSequenceClassification.from_pretrained('distilbert-base-uncased')
          input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute")).unsqueeze(0)  # Batch size 1
          labels = torch.tensor([1]).unsqueeze(0)  # Batch size 1
          outputs = model(input_ids, labels=labels)
          loss, logits = outputs[:2]
      """  # noqa: ignore flake8"
  
      def __init__(self, config, weight=None):
          super(DistilBertForSequenceClassification, self).__init__(config)
          self.num_labels = config.num_labels
          self.weight = weight
  
          self.distilbert = DistilBertModel(config)
          self.pre_classifier = nn.Linear(config.dim, config.dim)
          self.classifier = nn.Linear(config.dim, config.num_labels)
          self.dropout = nn.Dropout(config.seq_classif_dropout)
  
          self.init_weights()
  
      def forward(
          self,
          input_ids=None,
          attention_mask=None,
          head_mask=None,
          inputs_embeds=None,
          labels=None,
          class_weights=None,
      ):
          distilbert_output = self.distilbert(input_ids=input_ids, attention_mask=attention_mask, head_mask=head_mask)
          hidden_state = distilbert_output[0]  # (bs, seq_len, dim)
          pooled_output = hidden_state[:, 0]  # (bs, dim)
          pooled_output = self.pre_classifier(pooled_output)  # (bs, dim)
          pooled_output = nn.ReLU()(pooled_output)  # (bs, dim)
          pooled_output = self.dropout(pooled_output)  # (bs, dim)
          logits = self.classifier(pooled_output)  # (bs, dim)
  
          outputs = (logits,) + distilbert_output[1:]
          if labels is not None:
              if self.num_labels == 1:
                  loss_fct = nn.MSELoss()
                  loss = loss_fct(logits.view(-1), labels.view(-1))
              else:
                  if self.weight is not None:
                      weight = self.weight.to(labels.device)
                  else:
                      weight = None
                  loss_fct = CrossEntropyLoss(weight=weight)
                  loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
              outputs = (loss,) + outputs
  
          return outputs  # (loss), logits, (hidden_states), (attentions)
  