|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
import importlib.util |
|
import logging |
|
import re |
|
from collections import OrderedDict |
|
from collections.abc import Sequence |
|
from functools import partial |
|
from typing import List, Optional, Tuple, Union |
|
|
|
import numpy as np |
|
import torch |
|
import torch.nn as nn |
|
import torch.nn.functional as F |
|
import torch.utils.checkpoint |
|
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss |
|
from transformers import AutoTokenizer, PretrainedConfig |
|
from transformers.modeling_outputs import (MaskedLMOutput, |
|
SequenceClassifierOutput) |
|
from transformers.modeling_utils import PreTrainedModel |
|
from transformers.models.bert.modeling_bert import ( |
|
BaseModelOutputWithPoolingAndCrossAttentions, BertForPreTrainingOutput) |
|
from transformers.models.xlm_roberta.modeling_xlm_roberta import \ |
|
XLMRobertaLMHead |
|
|
|
from .rotary import RotaryEmbedding |
|
from .block import Block |
|
from .configuration_xlm_roberta import XLMRobertaFlashConfig |
|
from .embedding import XLMRobertaEmbeddings |
|
from .mha import MHA |
|
from .mlp import FusedMLP, Mlp |
|
from .xlm_padding import index_first_axis_residual, pad_input, unpad_input |
|
|
|
try: |
|
from flash_attn.ops.fused_dense import FusedDense |
|
except ImportError: |
|
FusedDense = None |
|
|
|
try: |
|
from flash_attn.ops.triton.layer_norm import layer_norm_fn |
|
except ImportError: |
|
layer_norm_fn = None |
|
|
|
|
|
try: |
|
from flash_attn.losses.cross_entropy import CrossEntropyLoss |
|
except ImportError: |
|
CrossEntropyLoss = torch.nn.CrossEntropyLoss |
|
|
|
try: |
|
from tqdm.autonotebook import trange |
|
except ImportError: |
|
trange = None |
|
|
|
|
|
logger = logging.getLogger(__name__) |
|
|
|
|
|
def get_use_flash_attn(config: XLMRobertaFlashConfig): |
|
if not getattr(config, "use_flash_attn", False) or not torch.cuda.is_available(): |
|
return False |
|
if importlib.util.find_spec("flash_attn") is None: |
|
logger.warning( |
|
"flash_attn is not installed. Using PyTorch native attention implementation." |
|
) |
|
return False |
|
return True |
|
|
|
|
|
def create_mixer_cls(config, cross_attn=False, return_residual=False): |
|
use_flash_attn = get_use_flash_attn(config) |
|
fused_bias_fc = getattr(config, "fused_bias_fc", False) |
|
rotary_kwargs = {} |
|
if config.position_embedding_type == "rotary": |
|
rotary_kwargs["rotary_emb_dim"] = getattr( |
|
config, "rotary_emb_dim", config.hidden_size / config.num_attention_heads |
|
) |
|
rotary_kwargs["rotary_emb_base"] = config.rotary_emb_base |
|
rotary_kwargs["rotary_emb_scale_base"] = getattr( |
|
config, "rotary_emb_scale_base", None |
|
) |
|
rotary_kwargs["rotary_emb_interleaved"] = getattr( |
|
config, "rotary_emb_interleaved", False |
|
) |
|
mixer_cls = partial( |
|
MHA, |
|
num_heads=config.num_attention_heads, |
|
cross_attn=cross_attn, |
|
dropout=config.attention_probs_dropout_prob, |
|
causal=False, |
|
fused_bias_fc=fused_bias_fc, |
|
use_flash_attn=use_flash_attn, |
|
return_residual=return_residual, |
|
use_alibi=config.position_embedding_type == "alibi", |
|
**rotary_kwargs, |
|
) |
|
return mixer_cls |
|
|
|
|
|
def create_mlp_cls(config, layer_idx=None, return_residual=False): |
|
inner_dim = config.intermediate_size |
|
fused_mlp = getattr(config, "fused_mlp", False) |
|
if fused_mlp: |
|
assert config.hidden_act in ["gelu_new", "gelu_fast", "gelu_pytorch_tanh"], ( |
|
"fused_mlp only " "supports approximate gelu" |
|
) |
|
if not fused_mlp: |
|
approximate = ( |
|
"tanh" |
|
if config.hidden_act in ["gelu_new", "gelu_fast", "gelu_pytorch_tanh"] |
|
else "none" |
|
) |
|
mlp_cls = partial( |
|
Mlp, |
|
hidden_features=inner_dim, |
|
activation=partial(F.gelu, approximate=approximate), |
|
return_residual=return_residual, |
|
) |
|
else: |
|
if FusedMLP is None: |
|
raise ImportError("fused_dense is not installed") |
|
mlp_checkpoint_lvl = getattr(config, "mlp_checkpoint_lvl", 0) |
|
|
|
if isinstance(mlp_checkpoint_lvl, Sequence): |
|
assert layer_idx is not None |
|
mlp_checkpoint_lvl = mlp_checkpoint_lvl[layer_idx] |
|
mlp_cls = partial( |
|
FusedMLP, |
|
hidden_features=inner_dim, |
|
checkpoint_lvl=mlp_checkpoint_lvl, |
|
return_residual=return_residual, |
|
) |
|
return mlp_cls |
|
|
|
|
|
def create_block(config, layer_idx=None): |
|
last_layer_subset = getattr(config, "last_layer_subset", False) |
|
cross_attn = last_layer_subset and layer_idx == config.num_hidden_layers - 1 |
|
|
|
|
|
|
|
return_residual = not cross_attn |
|
mixer_cls = create_mixer_cls(config, cross_attn, return_residual=return_residual) |
|
mlp_cls = create_mlp_cls(config, layer_idx, return_residual=return_residual) |
|
norm_cls = partial(nn.LayerNorm, eps=config.layer_norm_eps) |
|
block = Block( |
|
config.hidden_size, |
|
mixer_cls, |
|
mlp_cls, |
|
norm_cls=norm_cls, |
|
prenorm=False, |
|
resid_dropout1=config.hidden_dropout_prob, |
|
resid_dropout2=config.hidden_dropout_prob, |
|
fused_dropout_add_ln=getattr(config, "fused_dropout_add_ln", False), |
|
return_residual=return_residual, |
|
) |
|
return block |
|
|
|
|
|
|
|
def _init_weights(module, initializer_range=0.02): |
|
if isinstance(module, nn.Linear): |
|
nn.init.normal_(module.weight, std=initializer_range) |
|
if module.bias is not None: |
|
nn.init.zeros_(module.bias) |
|
elif isinstance(module, nn.Embedding): |
|
nn.init.normal_(module.weight, std=initializer_range) |
|
if module.padding_idx is not None: |
|
nn.init.zeros_(module.weight[module.padding_idx]) |
|
|
|
|
|
class XLMRobertaEncoder(nn.Module): |
|
def __init__(self, config: XLMRobertaFlashConfig): |
|
super().__init__() |
|
self.use_flash_attn = get_use_flash_attn(config) |
|
self.use_reentrant = config.use_reentrant |
|
self.layers = nn.ModuleList( |
|
[create_block(config, layer_idx=i) for i in range(config.num_hidden_layers)] |
|
) |
|
self._grad_checkpointing = False |
|
|
|
@property |
|
def gradient_checkpointing(self): |
|
return self._grad_checkpointing |
|
|
|
@gradient_checkpointing.setter |
|
def gradient_checkpointing(self, value): |
|
self._grad_checkpointing = value |
|
|
|
def forward( |
|
self, hidden_states, key_padding_mask=None, subset_mask=None, adapter_mask=None |
|
): |
|
"""If subset_mask is not None, we only want output for the subset of the sequence. |
|
This means that we only compute the last layer output for these tokens. |
|
subset_mask: (batch, seqlen), dtype=torch.bool |
|
""" |
|
if key_padding_mask is None or not self.use_flash_attn: |
|
mixer_kwargs = {"adapter_mask": adapter_mask} |
|
if key_padding_mask is not None: |
|
mixer_kwargs["key_padding_mask"] = key_padding_mask.bool() |
|
for layer in self.layers: |
|
if self._grad_checkpointing: |
|
hidden_states = torch.utils.checkpoint.checkpoint( |
|
layer, |
|
hidden_states, |
|
use_reentrant=self.use_reentrant, |
|
mixer_kwargs=mixer_kwargs, |
|
) |
|
else: |
|
hidden_states = layer(hidden_states, mixer_kwargs=mixer_kwargs) |
|
if subset_mask is not None: |
|
hidden_states = hidden_states[subset_mask] |
|
else: |
|
batch, seqlen = hidden_states.shape[:2] |
|
hidden_states, indices, cu_seqlens, max_seqlen_in_batch, cu_adapter_mask = ( |
|
unpad_input(hidden_states, key_padding_mask, adapter_mask) |
|
) |
|
mixer_kwargs = { |
|
"cu_seqlens": cu_seqlens, |
|
"max_seqlen": max_seqlen_in_batch, |
|
"adapter_mask": cu_adapter_mask, |
|
} |
|
|
|
if subset_mask is None: |
|
for layer in self.layers: |
|
if self._grad_checkpointing: |
|
hidden_states = torch.utils.checkpoint.checkpoint( |
|
layer, |
|
hidden_states, |
|
use_reentrant=self.use_reentrant, |
|
mixer_kwargs=mixer_kwargs, |
|
) |
|
else: |
|
hidden_states = layer(hidden_states, mixer_kwargs=mixer_kwargs) |
|
hidden_states = pad_input(hidden_states, indices, batch, seqlen) |
|
else: |
|
for layer in self.layers[:-1]: |
|
if self._grad_checkpointing: |
|
hidden_states = torch.utils.checkpoint.checkpoint( |
|
layer, |
|
hidden_states, |
|
use_reentrant=self.use_reentrant, |
|
mixer_kwargs=mixer_kwargs, |
|
) |
|
else: |
|
hidden_states = layer(hidden_states, mixer_kwargs=mixer_kwargs) |
|
if key_padding_mask is not None: |
|
subset_idx = torch.nonzero( |
|
subset_mask[key_padding_mask], as_tuple=False |
|
).flatten() |
|
subset_seqlens = (subset_mask & key_padding_mask).sum( |
|
dim=-1, dtype=torch.int32 |
|
) |
|
subset_cu_seqlens = F.pad( |
|
torch.cumsum(subset_seqlens, dim=0, dtype=torch.torch.int32), |
|
(1, 0), |
|
) |
|
else: |
|
subset_idx = torch.nonzero(subset_mask, as_tuple=False).flatten() |
|
subset_seqlens = subset_mask.sum(dim=-1, dtype=torch.int32) |
|
subset_cu_seqlens = F.pad( |
|
torch.cumsum(subset_seqlens, dim=0, dtype=torch.torch.int32), |
|
(1, 0), |
|
) |
|
hidden_states_subset, hidden_states = index_first_axis_residual( |
|
hidden_states, subset_idx |
|
) |
|
|
|
mixer_kwargs = { |
|
"x_kv": hidden_states, |
|
"cu_seqlens": subset_cu_seqlens, |
|
"max_seqlen": max_seqlen_in_batch, |
|
"cu_seqlens_k": cu_seqlens, |
|
"max_seqlen_k": max_seqlen_in_batch, |
|
} |
|
if self._grad_checkpointing: |
|
torch.utils.checkpoint.checkpoint( |
|
self.layers[-1], |
|
hidden_states_subset, |
|
use_reentrant=self.use_reentrant, |
|
mixer_kwargs=mixer_kwargs, |
|
) |
|
else: |
|
hidden_states = self.layers[-1]( |
|
hidden_states_subset, mixer_kwargs=mixer_kwargs |
|
) |
|
return hidden_states |
|
|
|
|
|
class XLMRobertaPooler(nn.Module): |
|
def __init__(self, config): |
|
super().__init__() |
|
fused_bias_fc = getattr(config, "fused_bias_fc", False) |
|
if fused_bias_fc and FusedDense is None: |
|
raise ImportError("fused_dense is not installed") |
|
linear_cls = nn.Linear if not fused_bias_fc else FusedDense |
|
self.dense = linear_cls(config.hidden_size, config.hidden_size) |
|
self.activation = nn.Tanh() |
|
|
|
def forward(self, hidden_states, pool=True, adapter_mask=None): |
|
|
|
|
|
first_token_tensor = hidden_states[:, 0] if pool else hidden_states |
|
if adapter_mask is not None: |
|
unique_tasks = torch.unique(adapter_mask) |
|
pool_dtype = next(self.dense.parameters()).dtype |
|
pooled_output = torch.empty( |
|
first_token_tensor.shape[0], |
|
self.dense.out_features, |
|
dtype=pool_dtype, |
|
device=first_token_tensor.device, |
|
) |
|
for task_id in unique_tasks: |
|
task_indices = (adapter_mask == task_id).nonzero(as_tuple=True)[0] |
|
task_first_token_tensor = first_token_tensor[task_indices] |
|
task_pooled_output = self.dense( |
|
task_first_token_tensor, task_id=task_id |
|
) |
|
pooled_output[task_indices] = task_pooled_output |
|
else: |
|
pooled_output = self.dense(first_token_tensor) |
|
pooled_output = self.activation(pooled_output) |
|
return pooled_output |
|
|
|
|
|
class XLMRobertaPredictionHeadTransform(nn.Module): |
|
def __init__(self, config): |
|
super().__init__() |
|
fused_bias_fc = getattr(config, "fused_bias_fc", False) |
|
if fused_bias_fc and FusedDense is None: |
|
raise ImportError("fused_dense is not installed") |
|
self.fused_dropout_add_ln = getattr(config, "fused_dropout_add_ln", False) |
|
if self.fused_dropout_add_ln and layer_norm_fn is None: |
|
raise ImportError("Triton is not installed") |
|
linear_cls = nn.Linear if not fused_bias_fc else FusedDense |
|
self.dense = linear_cls(config.hidden_size, config.hidden_size) |
|
approximate = ( |
|
"tanh" |
|
if config.hidden_act in ["gelu_new", "gelu_fast", "gelu_pytorch_tanh"] |
|
else "none" |
|
) |
|
self.transform_act_fn = nn.GELU(approximate=approximate) |
|
self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) |
|
|
|
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: |
|
hidden_states = self.dense(hidden_states) |
|
hidden_states = self.transform_act_fn(hidden_states) |
|
if not self.fused_dropout_add_ln: |
|
hidden_states = self.layer_norm(hidden_states) |
|
else: |
|
hidden_states = layer_norm_fn( |
|
hidden_states, |
|
self.layer_norm.weight, |
|
self.layer_norm.bias, |
|
eps=self.layer_norm.eps, |
|
) |
|
return hidden_states |
|
|
|
|
|
class XLMRobertaLMPredictionHead(nn.Module): |
|
def __init__(self, config): |
|
super().__init__() |
|
fused_bias_fc = getattr(config, "fused_bias_fc", False) |
|
if fused_bias_fc and FusedDense is None: |
|
raise ImportError("fused_dense is not installed") |
|
linear_cls = nn.Linear if not fused_bias_fc else FusedDense |
|
|
|
self.transform = XLMRobertaPredictionHeadTransform(config) |
|
|
|
|
|
|
|
self.decoder = linear_cls(config.hidden_size, config.vocab_size, bias=True) |
|
|
|
def forward(self, hidden_states): |
|
hidden_states = self.transform(hidden_states) |
|
hidden_states = self.decoder(hidden_states) |
|
return hidden_states |
|
|
|
|
|
class XLMRobertaPreTrainingHeads(nn.Module): |
|
def __init__(self, config): |
|
super().__init__() |
|
self.predictions = XLMRobertaLMPredictionHead(config) |
|
self.seq_relationship = nn.Linear(config.hidden_size, 2) |
|
|
|
def forward(self, sequence_output, pooled_output): |
|
prediction_scores = self.predictions(sequence_output) |
|
seq_relationship_score = self.seq_relationship(pooled_output) |
|
return prediction_scores, seq_relationship_score |
|
|
|
|
|
class XLMRobertaPreTrainedModel(PreTrainedModel): |
|
"""An abstract class to handle weights initialization and |
|
a simple interface for dowloading and loading pretrained models. |
|
""" |
|
|
|
config_class = XLMRobertaFlashConfig |
|
base_model_prefix = "roberta" |
|
supports_gradient_checkpointing = True |
|
_supports_param_buffer_assignment = False |
|
|
|
def _set_gradient_checkpointing(self, module, value=False): |
|
if isinstance(module, XLMRobertaEncoder): |
|
module.gradient_checkpointing = value |
|
|
|
@classmethod |
|
def from_pretrained( |
|
cls, |
|
*args, |
|
**kwargs, |
|
): |
|
if not "torch_dtype" in kwargs: |
|
kwargs["torch_dtype"] = "auto" |
|
return super().from_pretrained(*args, **kwargs) |
|
|
|
|
|
class XLMRobertaModel(XLMRobertaPreTrainedModel): |
|
def __init__(self, config: XLMRobertaFlashConfig, add_pooling_layer=True): |
|
super().__init__(config) |
|
self.pad_vocab_size_multiple = getattr(config, "pad_vocab_size_multiple", 1) |
|
if config.vocab_size % self.pad_vocab_size_multiple != 0: |
|
config.vocab_size += self.pad_vocab_size_multiple - ( |
|
config.vocab_size % self.pad_vocab_size_multiple |
|
) |
|
self.fused_dropout_add_ln = getattr(config, "fused_dropout_add_ln", False) |
|
if self.fused_dropout_add_ln and layer_norm_fn is None: |
|
raise ImportError("Triton is not installed") |
|
assert config.hidden_act in [ |
|
"gelu", |
|
"gelu_new", |
|
"gelu_fast", |
|
"gelu_pytorch_tanh", |
|
] |
|
self.embeddings = XLMRobertaEmbeddings( |
|
config.hidden_size, |
|
config.vocab_size, |
|
( |
|
config.max_position_embeddings |
|
if config.position_embedding_type == "absolute" |
|
else -1 |
|
), |
|
config.type_vocab_size, |
|
padding_idx=config.pad_token_id, |
|
) |
|
self.emb_drop = nn.Dropout(config.hidden_dropout_prob) |
|
self.emb_ln = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) |
|
self.encoder = XLMRobertaEncoder(config) |
|
self.pooler = XLMRobertaPooler(config) if add_pooling_layer else None |
|
|
|
self.apply(partial(_init_weights, initializer_range=config.initializer_range)) |
|
self.tokenizer = AutoTokenizer.from_pretrained( |
|
self.name_or_path, trust_remote_code=True |
|
) |
|
self._rotary_emb_base = config.rotary_emb_base |
|
|
|
@torch.inference_mode() |
|
def encode( |
|
self: "XLMRobertaModel", |
|
sentences: Union[str, List[str]], |
|
batch_size: int = 32, |
|
show_progress_bar: Optional[bool] = None, |
|
output_value: str = "sentence_embedding", |
|
convert_to_numpy: bool = True, |
|
convert_to_tensor: bool = False, |
|
device: Optional[torch.device] = None, |
|
normalize_embeddings: bool = True, |
|
truncate_dim: Optional[int] = None, |
|
adapter_mask: Optional[torch.Tensor] = None, |
|
task: Optional[str] = None, |
|
**tokenizer_kwargs, |
|
) -> Union[List[torch.Tensor], np.ndarray, torch.Tensor]: |
|
""" |
|
Computes sentence embeddings |
|
Args: |
|
sentences(`str` or `List[str]`): |
|
Sentence or sentences to be encoded |
|
batch_size(`int`, *optional*, defaults to 32): |
|
Batch size for the computation |
|
show_progress_bar(`bool`, *optional*, defaults to None): |
|
Show a progress bar when encoding sentences. |
|
If set to None, progress bar is only shown when |
|
`logger.level == logging.INFO` or `logger.level == logging.DEBUG`. |
|
output_value(`str`, *optional*, defaults to 'sentence_embedding'): |
|
Default sentence_embedding, to get sentence embeddings. |
|
Can be set to token_embeddings to get wordpiece token embeddings. |
|
Set to None, to get all output values |
|
convert_to_numpy(`bool`, *optional*, defaults to True): |
|
If true, the output is a list of numpy vectors. |
|
Else, it is a list of pytorch tensors. |
|
convert_to_tensor(`bool`, *optional*, defaults to False): |
|
If true, you get one large tensor as return. |
|
Overwrites any setting from convert_to_numpy |
|
device(`torch.device`, *optional*, defaults to None): |
|
Which torch.device to use for the computation |
|
normalize_embeddings(`bool`, *optional*, defaults to True): |
|
If set to true, returned vectors will have length 1. In that case, the |
|
faster dot-product (util.dot_score) instead of cosine similarity can |
|
be used. |
|
truncate_dim(`int`, *optional*, defaults to None): |
|
The dimension to truncate sentence embeddings to. `None` does no truncation. |
|
tokenizer_kwargs(`Dict[str, Any]`, *optional*, defaults to {}): |
|
Keyword arguments for the tokenizer |
|
Returns: |
|
By default, a list of tensors is returned. |
|
If convert_to_tensor, a stacked tensor is returned. |
|
If convert_to_numpy, a numpy matrix is returned. |
|
""" |
|
is_training = self.training |
|
self.eval() |
|
|
|
if show_progress_bar is None: |
|
show_progress_bar = ( |
|
logger.getEffectiveLevel() == logging.INFO |
|
or logger.getEffectiveLevel() == logging.DEBUG |
|
) |
|
|
|
if convert_to_tensor: |
|
convert_to_numpy = False |
|
|
|
if output_value != "sentence_embedding": |
|
convert_to_tensor = False |
|
convert_to_numpy = False |
|
|
|
input_was_string = False |
|
if isinstance(sentences, str) or not hasattr(sentences, "__len__"): |
|
sentences = [sentences] |
|
input_was_string = True |
|
|
|
if device is not None: |
|
self.to(device) |
|
|
|
permutation = np.argsort([-len(i) for i in sentences]) |
|
inverse_permutation = np.argsort(permutation) |
|
sentences = [sentences[idx] for idx in permutation] |
|
|
|
tokenizer_kwargs["padding"] = tokenizer_kwargs.get("padding", True) |
|
tokenizer_kwargs["max_length"] = tokenizer_kwargs.get( |
|
"max_length", self.tokenizer.init_kwargs.get("model_max_length", 8192) |
|
) |
|
tokenizer_kwargs["truncation"] = tokenizer_kwargs.get("truncation", True) |
|
|
|
all_embeddings = [] |
|
|
|
if trange is not None: |
|
range_iter = trange( |
|
0, |
|
len(sentences), |
|
batch_size, |
|
desc="Encoding", |
|
disable=not show_progress_bar, |
|
) |
|
else: |
|
range_iter = range(0, len(sentences), batch_size) |
|
|
|
for i in range_iter: |
|
encoded_input = self.tokenizer( |
|
sentences[i : i + batch_size], |
|
return_tensors="pt", |
|
**tokenizer_kwargs, |
|
).to(self.device) |
|
lora_arguments = ( |
|
{"adapter_mask": adapter_mask[i : i + batch_size]} |
|
if adapter_mask is not None |
|
else {} |
|
) |
|
token_embs = self.forward(**encoded_input, **lora_arguments)[0] |
|
|
|
|
|
token_embs = token_embs.float() |
|
|
|
if output_value == "token_embeddings": |
|
raise NotImplementedError |
|
elif output_value is None: |
|
raise NotImplementedError |
|
else: |
|
if self.config.emb_pooler == "cls": |
|
embeddings = self.cls_pooling( |
|
token_embs, encoded_input["attention_mask"] |
|
) |
|
else: |
|
embeddings = self.mean_pooling( |
|
token_embs, encoded_input["attention_mask"] |
|
) |
|
|
|
all_embeddings.extend(embeddings) |
|
|
|
all_embeddings = [all_embeddings[idx] for idx in inverse_permutation] |
|
|
|
truncate_dim = truncate_dim or self.config.truncate_dim |
|
if truncate_dim: |
|
all_embeddings = self.truncate_embeddings(all_embeddings, truncate_dim) |
|
|
|
if normalize_embeddings: |
|
all_embeddings = [torch.nn.functional.normalize(embedding, p=2, dim=0) for embedding in all_embeddings] |
|
|
|
if convert_to_tensor: |
|
all_embeddings = torch.stack(all_embeddings) |
|
elif convert_to_numpy: |
|
all_embeddings = np.asarray([emb.cpu().numpy() for emb in all_embeddings]) |
|
|
|
if input_was_string: |
|
all_embeddings = all_embeddings[0] |
|
|
|
self.train(is_training) |
|
return all_embeddings |
|
|
|
def truncate_embeddings(self, embeddings, truncate_dim): |
|
if not self.config.matryoshka_dimensions: |
|
logger.warning( |
|
"Matryoshka embeddings are not supported, so dimension truncation will not be performed." |
|
) |
|
return embeddings |
|
elif truncate_dim in self.config.matryoshka_dimensions: |
|
return [tensor[:truncate_dim] for tensor in embeddings] |
|
else: |
|
raise ValueError( |
|
f"The provided `truncate_dim` value of {truncate_dim} is not supported. " |
|
f"Supported dimensions are {self.config.matryoshka_dimensions}." |
|
) |
|
|
|
def mean_pooling( |
|
self, token_embeddings: torch.Tensor, attention_mask: torch.Tensor |
|
): |
|
input_mask_expanded = ( |
|
attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float() |
|
) |
|
return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp( |
|
input_mask_expanded.sum(1), min=1e-9 |
|
) |
|
|
|
def cls_pooling(self, token_embeddings: torch.Tensor, attention_mask: torch.Tensor): |
|
return token_embeddings[:, 0] |
|
|
|
@property |
|
def rotary_emb_base(self): |
|
return self._rotary_emb_base |
|
|
|
@rotary_emb_base.setter |
|
def rotary_emb_base(self, base): |
|
if not isinstance(base, (int, float)): |
|
raise TypeError("Base must be an integer or float") |
|
logger.info(f"Changing RoPE base value to {base}") |
|
for layer in self.encoder.layers: |
|
layer.mixer.rotary_emb.base = base |
|
self._rotary_emb_base = base |
|
|
|
def forward( |
|
self, |
|
input_ids, |
|
position_ids=None, |
|
token_type_ids=None, |
|
attention_mask=None, |
|
masked_tokens_mask=None, |
|
return_dict=None, |
|
**kwargs, |
|
): |
|
"""If masked_tokens_mask is not None (i.e. last_layer_subset == True in XLMForPreTraining), |
|
we only want the output for the masked tokens. This means that we only compute the last |
|
layer output for these tokens. |
|
masked_tokens_mask: (batch, seqlen), dtype=torch.bool |
|
""" |
|
adapter_mask = kwargs.pop("adapter_mask", None) |
|
if kwargs: |
|
for key, value in kwargs.items(): |
|
if value is not None: |
|
logger.warning( |
|
"Flash attention implementation does not support kwargs: %s", |
|
key, |
|
) |
|
|
|
return_dict = ( |
|
return_dict if return_dict is not None else self.config.use_return_dict |
|
) |
|
|
|
hidden_states = self.embeddings( |
|
input_ids, |
|
position_ids=position_ids, |
|
token_type_ids=token_type_ids, |
|
adapter_mask=adapter_mask, |
|
) |
|
|
|
|
|
if not self.fused_dropout_add_ln: |
|
hidden_states = self.emb_ln(hidden_states) |
|
else: |
|
hidden_states = layer_norm_fn( |
|
hidden_states, self.emb_ln.weight, self.emb_ln.bias, eps=self.emb_ln.eps |
|
) |
|
hidden_states = self.emb_drop(hidden_states) |
|
|
|
if masked_tokens_mask is not None: |
|
batch_size, seqlen = input_ids.shape[:2] |
|
|
|
first_col_mask = torch.zeros( |
|
batch_size, seqlen, dtype=torch.bool, device=input_ids.device |
|
) |
|
first_col_mask[:, 0] = True |
|
subset_mask = masked_tokens_mask | first_col_mask |
|
else: |
|
subset_mask = None |
|
|
|
sequence_output = self.encoder( |
|
hidden_states, |
|
key_padding_mask=attention_mask, |
|
subset_mask=subset_mask, |
|
adapter_mask=adapter_mask, |
|
) |
|
|
|
if masked_tokens_mask is None: |
|
pooled_output = ( |
|
self.pooler(sequence_output, adapter_mask=adapter_mask) |
|
if self.pooler is not None |
|
else None |
|
) |
|
else: |
|
|
|
if attention_mask is not None: |
|
subset_idx = subset_mask[attention_mask] |
|
pool_input = sequence_output[first_col_mask[attention_mask][subset_idx]] |
|
sequence_output = sequence_output[ |
|
masked_tokens_mask[attention_mask][subset_idx] |
|
] |
|
else: |
|
pool_input = sequence_output[first_col_mask[subset_mask]] |
|
sequence_output = sequence_output[masked_tokens_mask[subset_mask]] |
|
pooled_output = ( |
|
self.pooler(pool_input, pool=False, adapter_mask=adapter_mask) |
|
if self.pooler is not None |
|
else None |
|
) |
|
|
|
if not return_dict: |
|
return sequence_output, pooled_output |
|
|
|
return BaseModelOutputWithPoolingAndCrossAttentions( |
|
last_hidden_state=sequence_output, |
|
pooler_output=pooled_output, |
|
) |
|
|
|
|
|
class XLMRobertaForMaskedLM(XLMRobertaPreTrainedModel): |
|
_tied_weights_keys = ["lm_head.decoder.weight", "lm_head.decoder.bias"] |
|
|
|
def __init__(self, config): |
|
super().__init__(config) |
|
|
|
if config.is_decoder: |
|
logger.warning( |
|
"If you want to use `XLMRobertaForMaskedLM` make sure `config.is_decoder=False` for " |
|
"bi-directional self-attention." |
|
) |
|
|
|
self.roberta = XLMRobertaModel(config, add_pooling_layer=False) |
|
self.lm_head = XLMRobertaLMHead(config) |
|
|
|
|
|
self.post_init() |
|
|
|
def get_input_embeddings(self): |
|
return self.roberta.embeddings.word_embeddings |
|
|
|
def get_output_embeddings(self): |
|
return self.lm_head.decoder |
|
|
|
def set_output_embeddings(self, new_embeddings): |
|
self.lm_head.decoder = new_embeddings |
|
|
|
def forward( |
|
self, |
|
input_ids: Optional[torch.LongTensor] = None, |
|
attention_mask: Optional[torch.FloatTensor] = None, |
|
token_type_ids: Optional[torch.LongTensor] = None, |
|
position_ids: Optional[torch.LongTensor] = None, |
|
head_mask: Optional[torch.FloatTensor] = None, |
|
inputs_embeds: Optional[torch.FloatTensor] = None, |
|
encoder_hidden_states: Optional[torch.FloatTensor] = None, |
|
encoder_attention_mask: Optional[torch.FloatTensor] = None, |
|
labels: Optional[torch.LongTensor] = None, |
|
output_attentions: Optional[bool] = None, |
|
output_hidden_states: Optional[bool] = None, |
|
return_dict: Optional[bool] = None, |
|
) -> Union[Tuple[torch.Tensor], MaskedLMOutput]: |
|
r""" |
|
labels (`torch.LongTensor` 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]` |
|
kwargs (`Dict[str, any]`, optional, defaults to *{}*): |
|
Used to hide legacy arguments that have been deprecated. |
|
""" |
|
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, |
|
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.lm_head(sequence_output) |
|
|
|
masked_lm_loss = None |
|
if labels is not None: |
|
|
|
labels = labels.to(prediction_scores.device) |
|
loss_fct = CrossEntropyLoss() |
|
masked_lm_loss = loss_fct( |
|
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 remap_state_dict(state_dict, config: PretrainedConfig): |
|
""" |
|
Map the state_dict of a Huggingface BERT model to be flash_attn compatible. |
|
""" |
|
|
|
|
|
def key_mapping_ln_gamma_beta(key): |
|
key = re.sub(r"LayerNorm.gamma$", "LayerNorm.weight", key) |
|
key = re.sub(r"LayerNorm.beta$", "LayerNorm.bias", key) |
|
return key |
|
|
|
state_dict = OrderedDict( |
|
(key_mapping_ln_gamma_beta(k), v) for k, v in state_dict.items() |
|
) |
|
|
|
|
|
def key_mapping_layers(key): |
|
return re.sub(r"^bert.encoder.layer.", "bert.encoder.layers.", key) |
|
|
|
state_dict = OrderedDict((key_mapping_layers(k), v) for k, v in state_dict.items()) |
|
|
|
|
|
def key_mapping_ln(key): |
|
key = re.sub(r"^bert.embeddings.LayerNorm.", "bert.emb_ln.", key) |
|
key = re.sub( |
|
r"^bert.encoder.layers.(\d+).attention.output.LayerNorm.(weight|bias)", |
|
r"bert.encoder.layers.\1.norm1.\2", |
|
key, |
|
) |
|
key = re.sub( |
|
r"^bert.encoder.layers.(\d+).output.LayerNorm.(weight|bias)", |
|
r"bert.encoder.layers.\1.norm2.\2", |
|
key, |
|
) |
|
key = re.sub( |
|
r"^cls.predictions.transform.LayerNorm.(weight|bias)", |
|
r"cls.predictions.transform.layer_norm.\1", |
|
key, |
|
) |
|
return key |
|
|
|
state_dict = OrderedDict((key_mapping_ln(k), v) for k, v in state_dict.items()) |
|
|
|
|
|
def key_mapping_mlp(key): |
|
key = re.sub( |
|
r"^bert.encoder.layers.(\d+).intermediate.dense.(weight|bias)", |
|
r"bert.encoder.layers.\1.mlp.fc1.\2", |
|
key, |
|
) |
|
key = re.sub( |
|
r"^bert.encoder.layers.(\d+).output.dense.(weight|bias)", |
|
r"bert.encoder.layers.\1.mlp.fc2.\2", |
|
key, |
|
) |
|
return key |
|
|
|
state_dict = OrderedDict((key_mapping_mlp(k), v) for k, v in state_dict.items()) |
|
|
|
|
|
last_layer_subset = getattr(config, "last_layer_subset", False) |
|
for d in range(config.num_hidden_layers): |
|
Wq = state_dict.pop(f"bert.encoder.layers.{d}.attention.self.query.weight") |
|
Wk = state_dict.pop(f"bert.encoder.layers.{d}.attention.self.key.weight") |
|
Wv = state_dict.pop(f"bert.encoder.layers.{d}.attention.self.value.weight") |
|
bq = state_dict.pop(f"bert.encoder.layers.{d}.attention.self.query.bias") |
|
bk = state_dict.pop(f"bert.encoder.layers.{d}.attention.self.key.bias") |
|
bv = state_dict.pop(f"bert.encoder.layers.{d}.attention.self.value.bias") |
|
if not (last_layer_subset and d == config.num_hidden_layers - 1): |
|
state_dict[f"bert.encoder.layers.{d}.mixer.Wqkv.weight"] = torch.cat( |
|
[Wq, Wk, Wv], dim=0 |
|
) |
|
state_dict[f"bert.encoder.layers.{d}.mixer.Wqkv.bias"] = torch.cat( |
|
[bq, bk, bv], dim=0 |
|
) |
|
else: |
|
state_dict[f"bert.encoder.layers.{d}.mixer.Wq.weight"] = Wq |
|
state_dict[f"bert.encoder.layers.{d}.mixer.Wkv.weight"] = torch.cat( |
|
[Wk, Wv], dim=0 |
|
) |
|
state_dict[f"bert.encoder.layers.{d}.mixer.Wq.bias"] = bq |
|
state_dict[f"bert.encoder.layers.{d}.mixer.Wkv.bias"] = torch.cat( |
|
[bk, bv], dim=0 |
|
) |
|
|
|
def key_mapping_attn(key): |
|
return re.sub( |
|
r"^bert.encoder.layers.(\d+).attention.output.dense.(weight|bias)", |
|
r"bert.encoder.layers.\1.mixer.out_proj.\2", |
|
key, |
|
) |
|
|
|
state_dict = OrderedDict((key_mapping_attn(k), v) for k, v in state_dict.items()) |
|
|
|
def key_mapping_decoder_bias(key): |
|
return re.sub(r"^cls.predictions.bias", "cls.predictions.decoder.bias", key) |
|
|
|
state_dict = OrderedDict( |
|
(key_mapping_decoder_bias(k), v) for k, v in state_dict.items() |
|
) |
|
|
|
|
|
pad_vocab_size_multiple = getattr(config, "pad_vocab_size_multiple", 1) |
|
if pad_vocab_size_multiple > 1: |
|
word_embeddings = state_dict["bert.embeddings.word_embeddings.weight"] |
|
state_dict["bert.embeddings.word_embeddings.weight"] = F.pad( |
|
word_embeddings, (0, 0, 0, config.vocab_size - word_embeddings.shape[0]) |
|
) |
|
decoder_weight = state_dict["cls.predictions.decoder.weight"] |
|
state_dict["cls.predictions.decoder.weight"] = F.pad( |
|
decoder_weight, (0, 0, 0, config.vocab_size - decoder_weight.shape[0]) |
|
) |
|
|
|
|
|
|
|
decoder_bias = state_dict["cls.predictions.decoder.bias"] |
|
state_dict["cls.predictions.decoder.bias"] = F.pad( |
|
decoder_bias, (0, config.vocab_size - decoder_bias.shape[0]), value=-100.0 |
|
) |
|
|
|
return state_dict |
|
|
|
|
|
def inv_remap_state_dict(state_dict, config: PretrainedConfig): |
|
""" |
|
Map the state_dict of a flash_attn model to be Huggingface BERT compatible. |
|
|
|
This function is meant to be the inverse of remap_state_dict. |
|
""" |
|
|
|
pad_vocab_size_multiple = getattr(config, "pad_vocab_size_multiple", 1) |
|
if pad_vocab_size_multiple > 1: |
|
word_embeddings = state_dict["bert.embeddings.word_embeddings.weight"] |
|
decoder_weight = state_dict["cls.predictions.decoder.weight"] |
|
decoder_bias = state_dict["cls.predictions.decoder.bias"] |
|
|
|
state_dict["bert.embeddings.word_embeddings.weight"] = word_embeddings[ |
|
: config.orig_vocab_size, : |
|
] |
|
state_dict["cls.predictions.decoder.weight"] = decoder_weight[ |
|
: config.orig_vocab_size, : |
|
] |
|
state_dict["cls.predictions.decoder.bias"] = decoder_bias[ |
|
: config.orig_vocab_size |
|
] |
|
|
|
for d in range(config.num_hidden_layers): |
|
last_layer_subset = getattr(config, "last_layer_subset", False) |
|
if not last_layer_subset or d != (config.num_hidden_layers - 1): |
|
Wqkv_weights = state_dict.pop(f"bert.encoder.layers.{d}.mixer.Wqkv.weight") |
|
Wqkv_biases = state_dict.pop(f"bert.encoder.layers.{d}.mixer.Wqkv.bias") |
|
state_dict[f"bert.encoder.layers.{d}.attention.self.query.weight"] = ( |
|
Wqkv_weights[: Wqkv_weights.shape[0] // 3, :] |
|
) |
|
state_dict[f"bert.encoder.layers.{d}.attention.self.key.weight"] = ( |
|
Wqkv_weights[ |
|
Wqkv_weights.shape[0] // 3 : 2 * Wqkv_weights.shape[0] // 3, : |
|
] |
|
) |
|
state_dict[f"bert.encoder.layers.{d}.attention.self.value.weight"] = ( |
|
Wqkv_weights[2 * Wqkv_weights.shape[0] // 3 :, :] |
|
) |
|
state_dict[f"bert.encoder.layers.{d}.attention.self.query.bias"] = ( |
|
Wqkv_biases[: Wqkv_biases.shape[0] // 3] |
|
) |
|
state_dict[f"bert.encoder.layers.{d}.attention.self.key.bias"] = ( |
|
Wqkv_biases[Wqkv_biases.shape[0] // 3 : 2 * Wqkv_biases.shape[0] // 3] |
|
) |
|
state_dict[f"bert.encoder.layers.{d}.attention.self.value.bias"] = ( |
|
Wqkv_biases[2 * Wqkv_biases.shape[0] // 3 :] |
|
) |
|
else: |
|
Wq_weight = state_dict.pop(f"bert.encoder.layers.{d}.mixer.Wq.weight") |
|
Wkv_weights = state_dict.pop(f"bert.encoder.layers.{d}.mixer.Wkv.weight") |
|
Wq_bias = state_dict.pop(f"bert.encoder.layers.{d}.mixer.Wq.bias") |
|
Wkv_biases = state_dict.pop(f"bert.encoder.layers.{d}.mixer.Wkv.bias") |
|
state_dict[f"bert.encoder.layers.{d}.attention.self.query.weight"] = ( |
|
Wq_weight |
|
) |
|
state_dict[f"bert.encoder.layers.{d}.attention.self.key.weight"] = ( |
|
Wkv_weights[: Wkv_weights.shape[0] // 2, :] |
|
) |
|
state_dict[f"bert.encoder.layers.{d}.attention.self.value.weight"] = ( |
|
Wkv_weights[Wkv_weights.shape[0] // 2 :, :] |
|
) |
|
state_dict[f"bert.encoder.layers.{d}.attention.self.query.bias"] = Wq_bias |
|
state_dict[f"bert.encoder.layers.{d}.attention.self.key.bias"] = Wkv_biases[ |
|
: Wkv_biases.shape[0] // 2 |
|
] |
|
state_dict[f"bert.encoder.layers.{d}.attention.self.value.bias"] = ( |
|
Wkv_biases[Wkv_biases.shape[0] // 2 :] |
|
) |
|
|
|
def inv_key_mapping_ln(key): |
|
key = re.sub(r"bert.emb_ln.", "bert.embeddings.LayerNorm.", key) |
|
key = re.sub( |
|
r"bert.encoder.layers.(\d+).norm1.(weight|bias)", |
|
r"bert.encoder.layers.\1.attention.output.LayerNorm.\2", |
|
key, |
|
) |
|
key = re.sub( |
|
r"bert.encoder.layers.(\d+).norm2.(weight|bias)", |
|
r"bert.encoder.layers.\1.output.LayerNorm.\2", |
|
key, |
|
) |
|
key = re.sub( |
|
r"cls.predictions.transform.layer_norm.(weight|bias)", |
|
r"cls.predictions.transform.LayerNorm.\1", |
|
key, |
|
) |
|
return key |
|
|
|
def inv_key_mapping_ln_gamma_beta(key): |
|
key = re.sub(r"LayerNorm.weight$", "LayerNorm.gamma", key) |
|
key = re.sub(r"LayerNorm.bias$", "LayerNorm.beta", key) |
|
return key |
|
|
|
def inv_key_mapping_layers(key): |
|
return re.sub(r"bert.encoder.layers.", "bert.encoder.layer.", key) |
|
|
|
def inv_key_mapping_mlp(key): |
|
key = re.sub( |
|
r"bert.encoder.layer.(\d+).mlp.fc1.(weight|bias)", |
|
r"bert.encoder.layer.\1.intermediate.dense.\2", |
|
key, |
|
) |
|
key = re.sub( |
|
r"bert.encoder.layer.(\d+).mlp.fc2.(weight|bias)", |
|
r"bert.encoder.layer.\1.output.dense.\2", |
|
key, |
|
) |
|
return key |
|
|
|
def inv_key_mapping_attn(key): |
|
return re.sub( |
|
r"bert.encoder.layer.(\d+).mixer.out_proj.(weight|bias)", |
|
r"bert.encoder.layer.\1.attention.output.dense.\2", |
|
key, |
|
) |
|
|
|
def inv_key_mapping_decoder_bias(key): |
|
return re.sub(r"cls.predictions.decoder.bias", "cls.predictions.bias", key) |
|
|
|
state_dict = OrderedDict( |
|
(inv_key_mapping_ln(key), value) for key, value in state_dict.items() |
|
) |
|
state_dict = OrderedDict( |
|
(inv_key_mapping_ln_gamma_beta(key), value) for key, value in state_dict.items() |
|
) |
|
state_dict = OrderedDict( |
|
(inv_key_mapping_layers(key), value) for key, value in state_dict.items() |
|
) |
|
state_dict = OrderedDict( |
|
(inv_key_mapping_mlp(key), value) for key, value in state_dict.items() |
|
) |
|
state_dict = OrderedDict( |
|
(inv_key_mapping_attn(key), value) for key, value in state_dict.items() |
|
) |
|
state_dict = OrderedDict( |
|
(inv_key_mapping_decoder_bias(key), value) for key, value in state_dict.items() |
|
) |
|
|
|
return state_dict |
|
|
|
|
|
|
|
class XLMRobertaClassificationHead(nn.Module): |
|
"""Head for sentence-level classification tasks.""" |
|
|
|
def __init__(self, config): |
|
super().__init__() |
|
fused_bias_fc = getattr(config, "fused_bias_fc", False) |
|
if fused_bias_fc and FusedDense is None: |
|
raise ImportError("fused_dense is not installed") |
|
linear_cls = nn.Linear if not fused_bias_fc else FusedDense |
|
self.dense = linear_cls(config.hidden_size, config.hidden_size) |
|
classifier_dropout = ( |
|
config.classifier_dropout |
|
if config.classifier_dropout is not None |
|
else config.hidden_dropout_prob |
|
) |
|
self.dropout = nn.Dropout(classifier_dropout) |
|
self.out_proj = linear_cls(config.hidden_size, config.num_labels) |
|
|
|
def forward(self, features, **kwargs): |
|
x = features[:, 0, :] |
|
x = self.dropout(x) |
|
x = self.dense(x) |
|
x = torch.tanh(x) |
|
x = self.dropout(x) |
|
x = self.out_proj(x) |
|
return x |
|
|
|
|
|
|
|
class XLMRobertaForSequenceClassification(XLMRobertaPreTrainedModel): |
|
def __init__(self, config): |
|
super().__init__(config) |
|
self.num_labels = config.num_labels |
|
self.config = config |
|
|
|
self.roberta = XLMRobertaModel(config, add_pooling_layer=False) |
|
self.classifier = XLMRobertaClassificationHead(config) |
|
|
|
|
|
self.post_init() |
|
|
|
def forward( |
|
self, |
|
input_ids: Optional[torch.LongTensor] = None, |
|
attention_mask: Optional[torch.FloatTensor] = None, |
|
token_type_ids: Optional[torch.LongTensor] = None, |
|
position_ids: Optional[torch.LongTensor] = None, |
|
head_mask: Optional[torch.FloatTensor] = None, |
|
inputs_embeds: Optional[torch.FloatTensor] = None, |
|
labels: Optional[torch.LongTensor] = None, |
|
output_attentions: Optional[bool] = None, |
|
output_hidden_states: Optional[bool] = None, |
|
return_dict: Optional[bool] = None, |
|
) -> Union[Tuple[torch.Tensor], SequenceClassifierOutput]: |
|
r""" |
|
labels (`torch.LongTensor` 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.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) |
|
|
|
loss = None |
|
if labels is not None: |
|
|
|
labels = labels.to(logits.device) |
|
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 == torch.long or labels.dtype == torch.int |
|
): |
|
self.config.problem_type = "single_label_classification" |
|
else: |
|
self.config.problem_type = "multi_label_classification" |
|
|
|
if self.config.problem_type == "regression": |
|
loss_fct = MSELoss() |
|
if self.num_labels == 1: |
|
loss = loss_fct(logits.squeeze(), labels.squeeze()) |
|
else: |
|
loss = loss_fct(logits, labels) |
|
elif self.config.problem_type == "single_label_classification": |
|
loss_fct = CrossEntropyLoss() |
|
loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) |
|
elif self.config.problem_type == "multi_label_classification": |
|
loss_fct = BCEWithLogitsLoss() |
|
loss = loss_fct(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, |
|
) |
|
|