Added pred func

bf1f674 almost 3 years ago

15.1 kB

"""
Modified HuggingFace transformer model classes
"""
from typing import Tuple

import numpy as np
import torch
from torch import nn
from torch.nn import BCELoss, BCEWithLogitsLoss, MSELoss, PoissonNLLLoss, KLDivLoss

from transformers import BertConfig, BertModel, RobertaConfig, RobertaModel
from transformers import BertPreTrainedModel
from transformers.modeling_outputs import SequenceClassifierOutput
from transformers import RobertaPreTrainedModel


class RobertaMeanPoolConfig(RobertaConfig):
    model_type = "roberta"

    def __init__(
        self,
        output_mode="regression",
        freeze_base=True,
        start_token_idx=0,
        end_token_idx=1,
        threshold=1,
        alpha=0.5,
        log_offset=1,
        batch_norm=False,
        **kwargs,
    ):
        """Constructs RobertaConfig."""
        super().__init__(**kwargs)
        self.output_mode = output_mode
        self.freeze_base = freeze_base
        self.start_token_idx = start_token_idx
        self.end_token_idx = end_token_idx
        self.threshold = threshold
        self.alpha = alpha
        self.log_offset = log_offset
        self.batch_norm = batch_norm


class ClassificationHeadMeanPool(nn.Module):
    """Head for sentence-level classification tasks.
    Modifications:
        1. Using mean-pooling over tokens instead of CLS token
        2. Multi-output regression
    """

    def __init__(self, config: RobertaMeanPoolConfig):
        super().__init__()
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        self.dense2 = nn.Linear(config.hidden_size, config.hidden_size)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        self.out_proj = nn.Linear(config.hidden_size, config.num_labels)
        self.start_token_idx = config.start_token_idx
        self.end_token_idx = config.end_token_idx
        self.batch_norm = (
            nn.BatchNorm1d(config.hidden_size) if config.batch_norm else None
        )
        if self.batch_norm is not None:
            print("Using batch_norm")

    def forward(self, features, attention_mask=None, input_ids=None, **kwargs):
        x = self.embed(features, attention_mask, input_ids, **kwargs)
        x = self.out_proj(x)
        return x

    def embed(self, features, attention_mask=None, input_ids=None, **kwargs):
        attention_mask[input_ids == self.start_token_idx] = 0
        attention_mask[input_ids == self.end_token_idx] = 0
        x = torch.sum(features * attention_mask.unsqueeze(2), dim=1) / torch.sum(
            attention_mask, dim=1, keepdim=True
        )  # Mean pooling over non-padding tokens

        x = self.dropout(x)
        x = self.dense(x)
        x = torch.tanh(x)
        x = self.dropout(x)

        # Batchnorm
        x = self.normalize(x)

        # Second linear layer
        x = self.dense2(x)
        x = torch.tanh(x)
        return x

    def normalize(self, x: torch.Tensor) -> torch.Tensor:
        if self.batch_norm is not None:
            return self.batch_norm(x)
        return x


class ClassificationHeadMeanPoolSparse(nn.Module):
    """Classification head that predicts binary outcome (expressed/not)
    and real-valued gene expression values.
    """

    def __init__(self, config):
        super().__init__()
        self.classification_head = ClassificationHeadMeanPool(config)
        self.regression_head = ClassificationHeadMeanPool(config)

    def forward(
        self, features, attention_mask=None, input_ids=None, **kwargs
    ) -> Tuple[torch.Tensor]:
        """Compute binarized logits and real-valued gene expressions for each tissue.
        Args:
            features (torch.Tensor): outputs of RoBERTa
            attention_mask (Optional[torch.Tensor]): attention mask for sentence
            input_ids (Optional[torch.Tensor]): original sequence inputs
        Returns:
            (torch.Tensor): classification logits (whether gene is expressed/not for tissue)
            (torch.Tensor): gene expression value predictions (real-valued)
        """
        # Consider using .clone().detach()
        attention_mask_copy = attention_mask.clone()
        return (
            self.classification_head(
                features, attention_mask=attention_mask, input_ids=input_ids, **kwargs
            ),
            self.regression_head(
                features,
                attention_mask=attention_mask_copy,
                input_ids=input_ids,
                **kwargs,
            ),
        )


class SparseMSELoss(nn.Module):
    """Custom loss function that takes in two inputs:
    1. Predicted logits for whether gene is expressed (1) or not (0)
    2. Real-valued log-TPM values for gene expression predictions.
    """

    def __init__(self, threshold: float = 1, alpha: float = 0.5):
        """
        Args:
            threshold (float): any value below this threshold (in natural
                scale, NOT log-scale) is considered "not expressed"
            alpha (float): parameter controlling importance of classification
                in overall accuracy. alpha == 1 means this is identical to
                classification. alpha == 0 means this is identical to regression.
        """
        super().__init__()
        self.threshold = np.log(threshold)
        self.alpha = alpha
        self.mse = MSELoss()
        self.bce = BCEWithLogitsLoss()

    def forward(self, logits: Tuple[torch.Tensor], labels: torch.Tensor):
        classification_outputs, regression_outputs = logits
        binarized_labels = (labels >= self.threshold).float()

        mse_loss = self.mse(regression_outputs, labels)
        bce_loss = self.bce(classification_outputs, binarized_labels)

        # Weight the losses by the logits
        # the mse loss should be weighted by the probability of being expressed
        # the bce loss should be weighted by the probability of not being expressed

        loss = self.alpha * bce_loss + (1 - self.alpha) * mse_loss
        return loss


class ZeroInflatedNegativeBinomialNLL(nn.Module):
    """Custom loss function that calculates the negative log-likelihood
    according to a zero-inflated negative binomial model.
    """

    pass


# -------------------------------------- #
#                                        #
# ---------- Modified RoBERTa ---------- #
#                                        #
# -------------------------------------- #


class RobertaForSequenceClassificationMeanPool(RobertaPreTrainedModel):
    """RobertaForSequenceClassification using modified classification head
    Args:
        RobertaPreTrainedModel ([type]): [description]
    Returns:
        [type]: [description]
    """

    _keys_to_ignore_on_load_missing = [r"position_ids"]

    def __init__(self, config: RobertaMeanPoolConfig):
        super().__init__(config)
        self.num_labels = config.num_labels
        self.output_mode = config.output_mode or "regression"
        self.roberta = RobertaModel(config, add_pooling_layer=False)
        self.threshold = config.threshold
        self.alpha = config.alpha
        self.log_offset = config.log_offset

        if self.output_mode == "sparse":
            self.classifier = ClassificationHeadMeanPoolSparse(config)
        else:
            self.classifier = ClassificationHeadMeanPool(config)

        self.init_weights()

    def forward(
        self,
        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,
    ):
        r"""
        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`):
            Labels for computing the sequence classification/regression loss. Indices should be in :obj:`[0, ...,
            config.num_labels - 1]`. If :obj:`config.num_labels == 1` a regression loss is computed (Mean-Square loss),
            If :obj:`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.roberta(
            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.classifier(
            sequence_output, attention_mask=attention_mask, input_ids=input_ids
        )

        loss = None
        if labels is not None:
            if self.output_mode == "regression":
                loss_fct = MSELoss()
            elif self.output_mode == "sparse":
                loss_fct = SparseMSELoss(threshold=self.threshold, alpha=self.alpha)
            elif self.output_mode == "classification":
                loss_fct = BCEWithLogitsLoss()
            elif self.output_mode == "poisson":
                loss_fct = PoissonNLLLoss()

            loss = loss_fct(
                logits.view(-1, self.num_labels), labels.view(-1, self.num_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,
        )

    def embed(
        self,
        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,
    ):
        """Embed sequences by running the `forward` method up to the dense layer of the classifier"""
        outputs = self.roberta(
            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]
        embeddings = self.classifier.embed(
            sequence_output, attention_mask=attention_mask, input_ids=input_ids
        )
        return embeddings

    def get_tissue_embeddings(self):
        return self.classifier.out_proj.weight.detach()

    def predict(
        self,
        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,
    ):
        logits = self.forward(
            input_ids=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,
        )[0]
        if self.output_mode == "sparse":
            binary_logits, pred_values = logits
            # Convert logits to binary predictions
            binary_preds = binary_logits < 0
            # return binary_preds * pred_values
            pred_values[binary_preds] = np.log(self.log_offset)
            return pred_values
        return logits


# -------------------------------------- #
#                                        #
# ----------  Modified BERT  ----------- #
#                                        #
# -------------------------------------- #


class BertMeanPoolConfig(BertConfig):
    model_type = "bert"

    def __init__(
        self, output_mode="regression", start_token_idx=2, end_token_idx=3, **kwargs
    ):
        """Constructs BertConfig."""
        super().__init__(**kwargs)
        self.output_mode = output_mode
        self.start_token_idx = start_token_idx
        self.end_token_idx = end_token_idx


class BertForSequenceClassificationMeanPool(BertPreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_labels
        self.output_mode = config.output_mode or "regression"
        self.bert = BertModel(config)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

        self.classifier = ClassificationHeadMeanPool(config)

        self.init_weights()

    def forward(
        self,
        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,
    ):
        r"""
        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`):
            Labels for computing the sequence classification/regression loss. Indices should be in :obj:`[0, ...,
            config.num_labels - 1]`. If :obj:`config.num_labels == 1` a regression loss is computed (Mean-Square loss),
            If :obj:`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[0]

        pooled_output = self.dropout(pooled_output)
        logits = self.classifier(
            pooled_output, attention_mask=attention_mask, input_ids=input_ids
        )

        loss = None
        if labels is not None:
            if self.output_mode == "regression":
                #  We are doing regression
                loss_fct = MSELoss()
                loss = loss_fct(logits.view(-1), labels.view(-1))
            else:
                loss_fct = BCELoss()
                loss = loss_fct(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 SequenceClassifierOutput(
            loss=loss,
            logits=logits,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

	"""
	Modified HuggingFace transformer model classes
	"""
	from typing import Tuple

	import numpy as np
	import torch
	from torch import nn
	from torch.nn import BCELoss, BCEWithLogitsLoss, MSELoss, PoissonNLLLoss, KLDivLoss

	from transformers import BertConfig, BertModel, RobertaConfig, RobertaModel
	from transformers import BertPreTrainedModel
	from transformers.modeling_outputs import SequenceClassifierOutput
	from transformers import RobertaPreTrainedModel


	class RobertaMeanPoolConfig(RobertaConfig):
	model_type = "roberta"

	def __init__(
	self,
	output_mode="regression",
	freeze_base=True,
	start_token_idx=0,
	end_token_idx=1,
	threshold=1,
	alpha=0.5,
	log_offset=1,
	batch_norm=False,
	**kwargs,
	):
	"""Constructs RobertaConfig."""
	super().__init__(**kwargs)
	self.output_mode = output_mode
	self.freeze_base = freeze_base
	self.start_token_idx = start_token_idx
	self.end_token_idx = end_token_idx
	self.threshold = threshold
	self.alpha = alpha
	self.log_offset = log_offset
	self.batch_norm = batch_norm


	class ClassificationHeadMeanPool(nn.Module):
	"""Head for sentence-level classification tasks.

	Modifications:
	1. Using mean-pooling over tokens instead of CLS token
	2. Multi-output regression
	"""

	def __init__(self, config: RobertaMeanPoolConfig):
	super().__init__()
	self.dense = nn.Linear(config.hidden_size, config.hidden_size)
	self.dense2 = nn.Linear(config.hidden_size, config.hidden_size)
	self.dropout = nn.Dropout(config.hidden_dropout_prob)
	self.out_proj = nn.Linear(config.hidden_size, config.num_labels)
	self.start_token_idx = config.start_token_idx
	self.end_token_idx = config.end_token_idx
	self.batch_norm = (
	nn.BatchNorm1d(config.hidden_size) if config.batch_norm else None
	)
	if self.batch_norm is not None:
	print("Using batch_norm")

	def forward(self, features, attention_mask=None, input_ids=None, **kwargs):
	x = self.embed(features, attention_mask, input_ids, **kwargs)
	x = self.out_proj(x)
	return x

	def embed(self, features, attention_mask=None, input_ids=None, **kwargs):
	attention_mask[input_ids == self.start_token_idx] = 0
	attention_mask[input_ids == self.end_token_idx] = 0
	x = torch.sum(features * attention_mask.unsqueeze(2), dim=1) / torch.sum(
	attention_mask, dim=1, keepdim=True
	) # Mean pooling over non-padding tokens

	x = self.dropout(x)
	x = self.dense(x)
	x = torch.tanh(x)
	x = self.dropout(x)

	# Batchnorm
	x = self.normalize(x)

	# Second linear layer
	x = self.dense2(x)
	x = torch.tanh(x)
	return x

	def normalize(self, x: torch.Tensor) -> torch.Tensor:
	if self.batch_norm is not None:
	return self.batch_norm(x)
	return x


	class ClassificationHeadMeanPoolSparse(nn.Module):
	"""Classification head that predicts binary outcome (expressed/not)
	and real-valued gene expression values.
	"""

	def __init__(self, config):
	super().__init__()
	self.classification_head = ClassificationHeadMeanPool(config)
	self.regression_head = ClassificationHeadMeanPool(config)

	def forward(
	self, features, attention_mask=None, input_ids=None, **kwargs
	) -> Tuple[torch.Tensor]:
	"""Compute binarized logits and real-valued gene expressions for each tissue.

	Args:
	features (torch.Tensor): outputs of RoBERTa
	attention_mask (Optional[torch.Tensor]): attention mask for sentence
	input_ids (Optional[torch.Tensor]): original sequence inputs

	Returns:
	(torch.Tensor): classification logits (whether gene is expressed/not for tissue)
	(torch.Tensor): gene expression value predictions (real-valued)
	"""
	# Consider using .clone().detach()
	attention_mask_copy = attention_mask.clone()
	return (
	self.classification_head(
	features, attention_mask=attention_mask, input_ids=input_ids, **kwargs
	),
	self.regression_head(
	features,
	attention_mask=attention_mask_copy,
	input_ids=input_ids,
	**kwargs,
	),
	)


	class SparseMSELoss(nn.Module):
	"""Custom loss function that takes in two inputs:
	1. Predicted logits for whether gene is expressed (1) or not (0)
	2. Real-valued log-TPM values for gene expression predictions.
	"""

	def __init__(self, threshold: float = 1, alpha: float = 0.5):
	"""
	Args:
	threshold (float): any value below this threshold (in natural
	scale, NOT log-scale) is considered "not expressed"
	alpha (float): parameter controlling importance of classification
	in overall accuracy. alpha == 1 means this is identical to
	classification. alpha == 0 means this is identical to regression.
	"""
	super().__init__()
	self.threshold = np.log(threshold)
	self.alpha = alpha
	self.mse = MSELoss()
	self.bce = BCEWithLogitsLoss()

	def forward(self, logits: Tuple[torch.Tensor], labels: torch.Tensor):
	classification_outputs, regression_outputs = logits
	binarized_labels = (labels >= self.threshold).float()

	mse_loss = self.mse(regression_outputs, labels)
	bce_loss = self.bce(classification_outputs, binarized_labels)

	# Weight the losses by the logits
	# the mse loss should be weighted by the probability of being expressed
	# the bce loss should be weighted by the probability of not being expressed

	loss = self.alpha * bce_loss + (1 - self.alpha) * mse_loss
	return loss


	class ZeroInflatedNegativeBinomialNLL(nn.Module):
	"""Custom loss function that calculates the negative log-likelihood
	according to a zero-inflated negative binomial model.
	"""

	pass


	# -------------------------------------- #
	# #
	# ---------- Modified RoBERTa ---------- #
	# #
	# -------------------------------------- #


	class RobertaForSequenceClassificationMeanPool(RobertaPreTrainedModel):
	"""RobertaForSequenceClassification using modified classification head

	Args:
	RobertaPreTrainedModel ([type]): [description]

	Returns:
	[type]: [description]
	"""

	_keys_to_ignore_on_load_missing = [r"position_ids"]

	def __init__(self, config: RobertaMeanPoolConfig):
	super().__init__(config)
	self.num_labels = config.num_labels
	self.output_mode = config.output_mode or "regression"
	self.roberta = RobertaModel(config, add_pooling_layer=False)
	self.threshold = config.threshold
	self.alpha = config.alpha
	self.log_offset = config.log_offset

	if self.output_mode == "sparse":
	self.classifier = ClassificationHeadMeanPoolSparse(config)
	else:
	self.classifier = ClassificationHeadMeanPool(config)

	self.init_weights()

	def forward(
	self,
	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,
	):
	r"""
	labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`):
	Labels for computing the sequence classification/regression loss. Indices should be in :obj:`[0, ...,
	config.num_labels - 1]`. If :obj:`config.num_labels == 1` a regression loss is computed (Mean-Square loss),
	If :obj:`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.roberta(
	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.classifier(
	sequence_output, attention_mask=attention_mask, input_ids=input_ids
	)

	loss = None
	if labels is not None:
	if self.output_mode == "regression":
	loss_fct = MSELoss()
	elif self.output_mode == "sparse":
	loss_fct = SparseMSELoss(threshold=self.threshold, alpha=self.alpha)
	elif self.output_mode == "classification":
	loss_fct = BCEWithLogitsLoss()
	elif self.output_mode == "poisson":
	loss_fct = PoissonNLLLoss()

	loss = loss_fct(
	logits.view(-1, self.num_labels), labels.view(-1, self.num_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,
	)

	def embed(
	self,
	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,
	):
	"""Embed sequences by running the `forward` method up to the dense layer of the classifier"""
	outputs = self.roberta(
	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]
	embeddings = self.classifier.embed(
	sequence_output, attention_mask=attention_mask, input_ids=input_ids
	)
	return embeddings

	def get_tissue_embeddings(self):
	return self.classifier.out_proj.weight.detach()

	def predict(
	self,
	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,
	):
	logits = self.forward(
	input_ids=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,
	)[0]
	if self.output_mode == "sparse":
	binary_logits, pred_values = logits
	# Convert logits to binary predictions
	binary_preds = binary_logits < 0
	# return binary_preds * pred_values
	pred_values[binary_preds] = np.log(self.log_offset)
	return pred_values
	return logits


	# -------------------------------------- #
	# #
	# ---------- Modified BERT ----------- #
	# #
	# -------------------------------------- #


	class BertMeanPoolConfig(BertConfig):
	model_type = "bert"

	def __init__(
	self, output_mode="regression", start_token_idx=2, end_token_idx=3, **kwargs
	):
	"""Constructs BertConfig."""
	super().__init__(**kwargs)
	self.output_mode = output_mode
	self.start_token_idx = start_token_idx
	self.end_token_idx = end_token_idx


	class BertForSequenceClassificationMeanPool(BertPreTrainedModel):
	def __init__(self, config):
	super().__init__(config)
	self.num_labels = config.num_labels
	self.output_mode = config.output_mode or "regression"
	self.bert = BertModel(config)
	self.dropout = nn.Dropout(config.hidden_dropout_prob)

	self.classifier = ClassificationHeadMeanPool(config)

	self.init_weights()

	def forward(
	self,
	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,
	):
	r"""
	labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`):
	Labels for computing the sequence classification/regression loss. Indices should be in :obj:`[0, ...,
	config.num_labels - 1]`. If :obj:`config.num_labels == 1` a regression loss is computed (Mean-Square loss),
	If :obj:`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[0]

	pooled_output = self.dropout(pooled_output)
	logits = self.classifier(
	pooled_output, attention_mask=attention_mask, input_ids=input_ids
	)

	loss = None
	if labels is not None:
	if self.output_mode == "regression":
	# We are doing regression
	loss_fct = MSELoss()
	loss = loss_fct(logits.view(-1), labels.view(-1))
	else:
	loss_fct = BCELoss()
	loss = loss_fct(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 SequenceClassifierOutput(
	loss=loss,
	logits=logits,
	hidden_states=outputs.hidden_states,
	attentions=outputs.attentions,
	)