text: from transformers import BertForSequenceClassification, BertModel, BertConfig, XLNetForSequenceClassification, RobertaModel, RobertaConfig, BertPreTrainedModel, RobertaForSequenceClassification, DistilBertForSequenceClassification
  
  import torch
  from torch import Tensor
  from torch.nn import BCEWithLogitsLoss
  
  
  class DistilBertForMultiLabelSequenceClassification(DistilBertForSequenceClassification):
      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]
      """
  
      def forward(self, input_ids,  attention_mask=None, labels=None, head_mask=None):
          distilbert_output = self.distilbert(input_ids=input_ids,
                                              attention_mask=attention_mask,
                                              head_mask=head_mask)
          # (bs, seq_len, dim)
          hidden_state = distilbert_output[0]
          pooled_output = hidden_state[:, 0]                    # (bs, dim)
          pooled_output = self.pre_classifier(pooled_output)   # (bs, dim)
          pooled_output = torch.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:
              loss_fct = BCEWithLogitsLoss()
  
              loss = loss_fct(logits.view(-1, self.num_labels),
                              labels.view(-1, self.num_labels))
              outputs = (loss,) + outputs
  
          return outputs  # (loss), logits, (hidden_states), (attentions)
  
  
  class RobertaForMultiLabelSequenceClassification(RobertaForSequenceClassification):
      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]``.
              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 = RoertaTokenizer.from_pretrained('roberta-base')
          model = RobertaForSequenceClassification.from_pretrained('roberta-base')
          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]
      """
      config_class = RobertaConfig
  #    pretrained_model_archive_map = ROBERTA_PRETRAINED_MODEL_ARCHIVE_MAP
      base_model_prefix = "roberta"
  
      def forward(self, input_ids, token_type_ids=None, attention_mask=None, labels=None,
                  position_ids=None, head_mask=None):
          outputs = self.roberta(input_ids, position_ids=position_ids, token_type_ids=token_type_ids,
                                 attention_mask=attention_mask, head_mask=head_mask)
          sequence_output = outputs[0]
          logits = self.classifier(sequence_output)
  
          outputs = (logits,) + outputs[2:]
  
          if labels is not None:
              loss_fct = BCEWithLogitsLoss()
  
              loss = loss_fct(logits.view(-1, self.num_labels),
                              labels.view(-1, self.num_labels))
              outputs = (loss,) + outputs
  
          return outputs  # (loss), logits, (hidden_states), (attentions)
  
  
  class BertForMultiLabelSequenceClassification(BertForSequenceClassification):
      """BERT model for classification.
      This module is composed of the BERT model with a linear layer on top of
      the pooled output.
      Params:
          `config`: a BertConfig class instance with the configuration to build a new model.
          `num_labels`: the number of classes for the classifier. Default = 2.
      Inputs:
          `input_ids`: a torch.LongTensor of shape [batch_size, sequence_length]
              with the word token indices in the vocabulary(see the tokens preprocessing logic in the scripts
              `extract_features.py`, `run_classifier.py` and `run_squad.py`)
          `token_type_ids`: an optional torch.LongTensor of shape [batch_size, sequence_length] with the token
              types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to
              a `sentence B` token (see BERT paper for more details).
          `attention_mask`: an optional torch.LongTensor of shape [batch_size, sequence_length] with indices
              selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max
              input sequence length in the current batch. It's the mask that we typically use for attention when
              a batch has varying length sentences.
          `labels`: labels for the classification output: torch.LongTensor of shape [batch_size]
              with indices selected in [0, ..., num_labels].
      Outputs:
          if `labels` is not `None`:
              Outputs the CrossEntropy classification loss of the output with the labels.
          if `labels` is `None`:
              Outputs the classification logits of shape [batch_size, num_labels].
      Example usage:
      ```python
      # Already been converted into WordPiece token ids
      input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]])
      input_mask = torch.LongTensor([[1, 1, 1], [1, 1, 0]])
      token_type_ids = torch.LongTensor([[0, 0, 1], [0, 1, 0]])
      config = BertConfig(vocab_size_or_config_json_file=32000, hidden_size=768,
          num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072)
      num_labels = 2
      model = BertForSequenceClassification(config, num_labels)
      logits = model(input_ids, token_type_ids, input_mask)
      ```
      """
  
      def forward(self, input_ids, token_type_ids=None, attention_mask=None, labels=None,
                  position_ids=None, head_mask=None):
  
          import pdb
          # pdb.set_trace()
          outputs = self.bert(input_ids, position_ids=position_ids, token_type_ids=token_type_ids,
                              attention_mask=attention_mask, head_mask=head_mask)
          pooled_output = outputs[1]
  
          pooled_output = self.dropout(pooled_output)
          logits = self.classifier(pooled_output)
  
          # add hidden states and attention if they are here
          outputs = (logits,) + outputs[2:]
  
          if labels is not None:
              loss_fct = BCEWithLogitsLoss()
  
              loss = loss_fct(logits.view(-1, self.num_labels),
                              labels.view(-1, self.num_labels))
              outputs = (loss,) + outputs
  
          return outputs  # (loss), logits, (hidden_states), (attentions)
  
  
  class XLNetForMultiLabelSequenceClassification(XLNetForSequenceClassification):
      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]``.
              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).
          **mems**:
              list of ``torch.FloatTensor`` (one for each layer):
              that contains pre-computed hidden-states (key and values in the attention blocks) as computed by the model
              (see `mems` input above). Can be used to speed up sequential decoding and attend to longer context.
          **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::
          >>> config = XLNetConfig.from_pretrained('xlnet-large-cased')
          >>> tokenizer = XLNetTokenizer.from_pretrained('xlnet-large-cased')
          >>> 
          >>> model = XLNetForSequenceClassification(config)
          >>> 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]
      """
  
      def forward(self, input_ids, token_type_ids=None, input_mask=None, attention_mask=None,
                  mems=None, perm_mask=None, target_mapping=None,
                  labels=None, head_mask=None):
          transformer_outputs = self.transformer(input_ids, token_type_ids=token_type_ids,
                                                 input_mask=input_mask, attention_mask=attention_mask,
                                                 mems=mems, perm_mask=perm_mask, target_mapping=target_mapping,
                                                 head_mask=head_mask)
          output = transformer_outputs[0]
  
          output = self.sequence_summary(output)
          logits = self.logits_proj(output)
  
          # Keep mems, hidden states, attentions if there are in it
          outputs = (logits,) + transformer_outputs[1:]
  
          if labels is not None:
              loss_fct = BCEWithLogitsLoss()
  
              loss = loss_fct(logits.view(-1, self.num_labels),
                              labels.view(-1, self.num_labels))
              outputs = (loss,) + outputs
  
          return outputs  # (loss), logits, (hidden_states), (attentions)
  