mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-28 09:08:54 +00:00
Introduce code from Rust monorepo. This is a mess
This commit is contained in:
Brandon Thomas
@@ -0,0 +1,20 @@
|
||||
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from nemo.collections.common.losses.aggregator import AggregatorLoss
|
||||
from nemo.collections.common.losses.cross_entropy import CrossEntropyLoss, NLLLoss
|
||||
from nemo.collections.common.losses.mse_loss import MSELoss
|
||||
from nemo.collections.common.losses.multi_similarity_loss import MultiSimilarityLoss
|
||||
from nemo.collections.common.losses.smoothed_cross_entropy import SmoothedCrossEntropyLoss
|
||||
from nemo.collections.common.losses.spanning_loss import SpanningLoss
|
||||
@@ -0,0 +1,67 @@
|
||||
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from typing import List
|
||||
|
||||
import torch
|
||||
|
||||
from nemo.core.classes import Loss, typecheck
|
||||
from nemo.core.neural_types import LossType, NeuralType
|
||||
|
||||
__all__ = ['AggregatorLoss']
|
||||
|
||||
|
||||
class AggregatorLoss(Loss):
|
||||
"""
|
||||
Sums several losses into one.
|
||||
|
||||
Args:
|
||||
num_inputs: number of input losses
|
||||
weights: a list of coefficient for merging losses
|
||||
"""
|
||||
|
||||
@property
|
||||
def input_types(self):
|
||||
"""Returns definitions of module input ports.
|
||||
"""
|
||||
input_types = {}
|
||||
for i in range(self._num_losses):
|
||||
input_types["loss_" + str(i + 1)] = NeuralType(elements_type=LossType())
|
||||
|
||||
return input_types
|
||||
|
||||
@property
|
||||
def output_types(self):
|
||||
"""Returns definitions of module output ports.
|
||||
"""
|
||||
return {"loss": NeuralType(elements_type=LossType())}
|
||||
|
||||
def __init__(self, num_inputs: int = 2, weights: List[float] = None):
|
||||
super().__init__()
|
||||
self._num_losses = num_inputs
|
||||
if weights is not None and len(weights) != num_inputs:
|
||||
raise ValueError("Length of weights should be equal to the number of inputs (num_inputs)")
|
||||
|
||||
self._weights = weights
|
||||
|
||||
@typecheck()
|
||||
def forward(self, **kwargs):
|
||||
values = [kwargs[x] for x in sorted(kwargs.keys())]
|
||||
loss = torch.zeros_like(values[0])
|
||||
for loss_idx, loss_value in enumerate(values):
|
||||
if self._weights is not None:
|
||||
loss = loss.add(loss_value, alpha=self._weights[loss_idx])
|
||||
else:
|
||||
loss = loss.add(loss_value)
|
||||
return loss
|
||||
@@ -0,0 +1,138 @@
|
||||
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from nemo.core.classes import Serialization, Typing, typecheck
|
||||
from nemo.core.neural_types import LabelsType, LogitsType, LogprobsType, LossType, MaskType, NeuralType
|
||||
|
||||
__all__ = ['CrossEntropyLoss', 'NLLLoss']
|
||||
|
||||
|
||||
class CrossEntropyLoss(nn.CrossEntropyLoss, Serialization, Typing):
|
||||
"""
|
||||
CrossEntropyLoss
|
||||
"""
|
||||
|
||||
@property
|
||||
def input_types(self):
|
||||
"""Returns definitions of module input ports.
|
||||
"""
|
||||
return {
|
||||
"logits": NeuralType(['B'] + ['ANY'] * (self._logits_dim - 1), LogitsType()),
|
||||
"labels": NeuralType(['B'] + ['ANY'] * (self._logits_dim - 2), LabelsType()),
|
||||
"loss_mask": NeuralType(['B'] + ['ANY'] * (self._logits_dim - 2), MaskType(), optional=True),
|
||||
}
|
||||
|
||||
@property
|
||||
def output_types(self):
|
||||
"""Returns definitions of module output ports.
|
||||
"""
|
||||
return {"loss": NeuralType(elements_type=LossType())}
|
||||
|
||||
def __init__(self, logits_ndim=2, weight=None, reduction='mean', ignore_index=-100):
|
||||
"""
|
||||
Args:
|
||||
logits_ndim (int): number of dimensions (or rank) of the logits tensor
|
||||
weight (list): list of rescaling weight given to each class
|
||||
reduction (str): type of the reduction over the batch
|
||||
"""
|
||||
if weight is not None and not torch.is_tensor(weight):
|
||||
weight = torch.FloatTensor(weight)
|
||||
super().__init__(weight=weight, reduction=reduction, ignore_index=ignore_index)
|
||||
self._logits_dim = logits_ndim
|
||||
|
||||
@typecheck()
|
||||
def forward(self, logits, labels, loss_mask=None):
|
||||
"""
|
||||
Args:
|
||||
logits (float): output of the classifier
|
||||
labels (long): ground truth labels
|
||||
loss_mask (bool/float/int): tensor to specify the masking
|
||||
"""
|
||||
logits_flatten = torch.flatten(logits, start_dim=0, end_dim=-2)
|
||||
labels_flatten = torch.flatten(labels, start_dim=0, end_dim=-1)
|
||||
|
||||
if loss_mask is not None:
|
||||
if loss_mask.dtype is not torch.bool:
|
||||
loss_mask = loss_mask > 0.5
|
||||
loss_mask_flatten = torch.flatten(loss_mask, start_dim=0, end_dim=-1)
|
||||
logits_flatten = logits_flatten[loss_mask_flatten]
|
||||
labels_flatten = labels_flatten[loss_mask_flatten]
|
||||
|
||||
if len(labels_flatten) == 0:
|
||||
return super().forward(logits, torch.argmax(logits, dim=-1))
|
||||
|
||||
loss = super().forward(logits_flatten, labels_flatten)
|
||||
return loss
|
||||
|
||||
|
||||
class NLLLoss(nn.NLLLoss, Serialization, Typing):
|
||||
"""
|
||||
NLLLoss
|
||||
"""
|
||||
|
||||
@property
|
||||
def input_types(self):
|
||||
"""Returns definitions of module input ports.
|
||||
"""
|
||||
return {
|
||||
"log_probs": NeuralType(("B", "T", "D"), LogprobsType()),
|
||||
"labels": NeuralType(("B", "T"), LabelsType()),
|
||||
"output_mask": NeuralType(("B", "T"), MaskType(), optional=True),
|
||||
}
|
||||
|
||||
@property
|
||||
def output_types(self):
|
||||
"""Returns definitions of module output ports.
|
||||
"""
|
||||
return {"loss": NeuralType(elements_type=LossType())}
|
||||
|
||||
def __init__(self, log_probs_ndim=2, weight=None, reduction='mean', ignore_index=-100):
|
||||
"""
|
||||
Args:
|
||||
log_probs_ndim (int): number of dimensions (or rank) of the logprobs tensor
|
||||
weight (list): list of rescaling weight given to each class
|
||||
reduction (str): type of the reduction over the batch
|
||||
ignore_index (int): mask out loss computation where labels = ignore_index
|
||||
"""
|
||||
if weight is not None and not torch.is_tensor(weight):
|
||||
weight = torch.FloatTensor(weight)
|
||||
super().__init__(weight=weight, reduction=reduction, ignore_index=ignore_index)
|
||||
self._log_probs_dim = log_probs_ndim
|
||||
|
||||
@typecheck()
|
||||
def forward(self, log_probs, labels, loss_mask=None):
|
||||
"""
|
||||
Args:
|
||||
log_probs (float): output log probability tensor
|
||||
labels (long): ground truth labels
|
||||
loss_mask (bool/float/int): tensor to specify the masking
|
||||
"""
|
||||
log_probs_flatten = torch.flatten(log_probs, start_dim=0, end_dim=-2)
|
||||
labels_flatten = torch.flatten(labels, start_dim=0, end_dim=-1)
|
||||
|
||||
if loss_mask is not None:
|
||||
if loss_mask.dtype is not torch.bool:
|
||||
loss_mask = loss_mask > 0.5
|
||||
loss_mask_flatten = torch.flatten(loss_mask, start_dim=0, end_dim=-1)
|
||||
log_probs_flatten = log_probs_flatten[loss_mask_flatten]
|
||||
labels_flatten = labels_flatten[loss_mask_flatten]
|
||||
|
||||
if len(labels_flatten) == 0:
|
||||
return super().forward(log_probs, torch.argmax(log_probs, dim=-1))
|
||||
|
||||
loss = super().forward(log_probs_flatten, labels_flatten)
|
||||
return loss
|
||||
@@ -0,0 +1,57 @@
|
||||
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from torch import Tensor, nn
|
||||
|
||||
from nemo.core.classes import Serialization, Typing, typecheck
|
||||
from nemo.core.neural_types import LabelsType, LossType, NeuralType, RegressionValuesType
|
||||
|
||||
__all__ = ['MSELoss']
|
||||
|
||||
|
||||
class MSELoss(nn.MSELoss, Serialization, Typing):
|
||||
"""
|
||||
MSELoss
|
||||
"""
|
||||
|
||||
@property
|
||||
def input_types(self):
|
||||
"""Returns definitions of module input ports.
|
||||
"""
|
||||
return {
|
||||
"preds": NeuralType(tuple('B'), RegressionValuesType()),
|
||||
"labels": NeuralType(tuple('B'), LabelsType()),
|
||||
}
|
||||
|
||||
@property
|
||||
def output_types(self):
|
||||
"""Returns definitions of module output ports.
|
||||
"""
|
||||
return {"loss": NeuralType(elements_type=LossType())}
|
||||
|
||||
def __init__(self, reduction: str = 'mean'):
|
||||
"""
|
||||
Args:
|
||||
reduction: type of the reduction over the batch
|
||||
"""
|
||||
super().__init__(reduction=reduction)
|
||||
|
||||
@typecheck()
|
||||
def forward(self, preds: Tensor, labels: Tensor) -> Tensor:
|
||||
"""
|
||||
Args:
|
||||
preds: output of the classifier
|
||||
labels: ground truth labels
|
||||
"""
|
||||
return super().forward(preds, labels)
|
||||
@@ -0,0 +1,95 @@
|
||||
# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
|
||||
from nemo.core.classes import Loss
|
||||
from nemo.core.classes.common import typecheck
|
||||
from nemo.core.neural_types import LabelsType, LogitsType, LossType, NeuralType
|
||||
from nemo.utils import logging
|
||||
|
||||
__all__ = ['MultiSimilarityLoss']
|
||||
|
||||
|
||||
class MultiSimilarityLoss(Loss):
|
||||
@property
|
||||
def input_types(self):
|
||||
"""Returns definitions of module input ports.
|
||||
"""
|
||||
return {"logits": NeuralType(('B', 'D'), LogitsType()), "labels": NeuralType(('B'), LabelsType())}
|
||||
|
||||
@property
|
||||
def output_types(self):
|
||||
"""Returns definitions of module output ports.
|
||||
"""
|
||||
return {"loss": NeuralType(elements_type=LossType())}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
scale_pos: Optional[float] = 2.0, # Params found to work best in our experiments
|
||||
scale_neg: Optional[float] = 40.0,
|
||||
offset: Optional[float] = 0.5,
|
||||
margin: Optional[float] = 0.1,
|
||||
):
|
||||
super().__init__()
|
||||
self._scale_pos = scale_pos
|
||||
self._scale_neg = scale_neg
|
||||
self._offset = offset
|
||||
self._margin = margin
|
||||
self._epsilon = 1e-5
|
||||
|
||||
@typecheck()
|
||||
def forward(self, logits, labels):
|
||||
cos_sim = torch.matmul(logits, torch.t(logits))
|
||||
losses = []
|
||||
|
||||
for i in range(logits.size(0)):
|
||||
# mine hard pairs relative to anchor i
|
||||
positive_sims = cos_sim[i][labels.eq(labels[i])]
|
||||
positive_sims = positive_sims[positive_sims.lt(1 - self._epsilon)] # omit identical pairs
|
||||
negative_sims = cos_sim[i][labels.ne(labels[i])]
|
||||
|
||||
if len(negative_sims) == 0 or len(positive_sims) == 0:
|
||||
continue
|
||||
|
||||
# negatives that are more similar than the least-similar positive
|
||||
hard_negatives = negative_sims[negative_sims.gt(min(positive_sims) - self._margin)]
|
||||
|
||||
# positives that are less similar than the most-similar negative
|
||||
hard_positives = positive_sims[positive_sims.lt(max(negative_sims) + self._margin)]
|
||||
|
||||
if len(hard_negatives) == 0 or len(hard_positives) == 0:
|
||||
continue
|
||||
|
||||
pos_term = (
|
||||
1.0
|
||||
/ self._scale_pos
|
||||
* torch.log(1 + torch.sum(torch.exp(-self._scale_pos * (hard_positives - self._offset))))
|
||||
)
|
||||
neg_term = (
|
||||
1.0
|
||||
/ self._scale_neg
|
||||
* torch.log(1 + torch.sum(torch.exp(self._scale_neg * (hard_negatives - self._offset))))
|
||||
)
|
||||
losses.append(pos_term + neg_term)
|
||||
|
||||
if len(losses) == 0:
|
||||
loss = torch.zeros([], requires_grad=True).cuda()
|
||||
logging.info(f'Encountered zero loss in multisimloss, loss = {loss}. No hard examples found in the batch')
|
||||
else:
|
||||
loss = torch.sum(torch.stack(losses)) / logits.size(0)
|
||||
|
||||
return loss
|
||||
@@ -0,0 +1,103 @@
|
||||
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
|
||||
from nemo.core.classes import Loss, typecheck
|
||||
from nemo.core.neural_types import LabelsType, LogprobsType, LossType, MaskType, NeuralType
|
||||
|
||||
__all__ = ["SmoothedCrossEntropyLoss"]
|
||||
|
||||
|
||||
class SmoothedCrossEntropyLoss(Loss):
|
||||
"""
|
||||
Calculates Cross-entropy loss with label smoothing for a batch of sequences.
|
||||
|
||||
SmoothedCrossEntropyLoss:
|
||||
1) excludes padding tokens from loss calculation
|
||||
2) allows to use label smoothing regularization
|
||||
3) allows to calculate loss for the desired number of last tokens
|
||||
|
||||
Args:
|
||||
label_smoothing (float): label smoothing regularization coefficient
|
||||
predict_last_k (int): parameter which sets the number of last tokens to calculate the loss for, for example
|
||||
0: (default) calculate loss on the entire sequence (e.g., NMT)
|
||||
1: calculate loss on the last token only (e.g., LM evaluation)
|
||||
Intermediate values allow to control the trade-off between eval
|
||||
time (proportional to the number of batches) and eval performance
|
||||
(proportional to the number of context tokens)
|
||||
pad_id (int): padding id
|
||||
eps (float): the small eps number to avoid division buy zero
|
||||
"""
|
||||
|
||||
@property
|
||||
def input_types(self):
|
||||
"""Returns definitions of module input ports.
|
||||
"""
|
||||
return {
|
||||
"log_probs": NeuralType(("B", "T", "D"), LogprobsType()),
|
||||
"labels": NeuralType(("B", "T"), LabelsType()),
|
||||
"output_mask": NeuralType(("B", "T"), MaskType(), optional=True),
|
||||
}
|
||||
|
||||
@property
|
||||
def output_types(self):
|
||||
"""Returns definitions of module output ports.
|
||||
"""
|
||||
return {"loss": NeuralType(elements_type=LossType())}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
pad_id: Optional[int] = None,
|
||||
label_smoothing: Optional[float] = 0.0,
|
||||
predict_last_k: Optional[int] = 0,
|
||||
eps: float = 1e-6,
|
||||
):
|
||||
super().__init__()
|
||||
self._pad_id = pad_id
|
||||
self._eps = eps
|
||||
self._predict_last_k = predict_last_k
|
||||
self._label_smoothing = label_smoothing
|
||||
|
||||
@typecheck()
|
||||
def forward(self, log_probs, labels, output_mask=None):
|
||||
"""
|
||||
Args:
|
||||
log_probs: float tensor of shape batch_size x seq_len x vocab_size, values should be log probabilities
|
||||
labels: int tensor of shape batch_size x seq_len
|
||||
output_mask: binary tensor of shape batch_size x seq_len
|
||||
eps: epsilon param to avoid divide by zero in loss calculation
|
||||
"""
|
||||
if output_mask is None and self._pad_id is None:
|
||||
raise ValueError("Both output_mask and pad_id are None")
|
||||
if output_mask is None and self._pad_id is not None:
|
||||
output_mask = (labels != self._pad_id).to(log_probs.dtype)
|
||||
|
||||
if output_mask.dtype is not log_probs.dtype:
|
||||
output_mask = output_mask.to(log_probs.dtype)
|
||||
|
||||
batch_size, seq_len, vocab_size = log_probs.size()
|
||||
smoothing = vocab_size * self._label_smoothing / (vocab_size - 1)
|
||||
target_log_probs = log_probs.gather(2, labels.unsqueeze(2)).squeeze(2)
|
||||
|
||||
smoothing_log_probs = log_probs.mean(dim=-1)
|
||||
neg_log_likelihood = (1.0 - smoothing) * target_log_probs + smoothing * smoothing_log_probs
|
||||
neg_log_likelihood = neg_log_likelihood[:, -self._predict_last_k :]
|
||||
output_mask = output_mask[:, -self._predict_last_k :]
|
||||
neg_log_likelihood = -torch.sum(neg_log_likelihood * output_mask)
|
||||
neg_log_likelihood = neg_log_likelihood / (output_mask.sum() + self._eps)
|
||||
|
||||
return neg_log_likelihood
|
||||
@@ -0,0 +1,79 @@
|
||||
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from torch import nn
|
||||
|
||||
from nemo.core.classes import Loss, typecheck
|
||||
from nemo.core.neural_types import ChannelType, LogitsType, LossType, NeuralType
|
||||
|
||||
__all__ = ['SpanningLoss']
|
||||
|
||||
|
||||
class SpanningLoss(Loss):
|
||||
"""
|
||||
implements start and end loss of a span e.g. for Question Answering.
|
||||
"""
|
||||
|
||||
@property
|
||||
def input_types(self):
|
||||
"""Returns definitions of module input ports.
|
||||
"""
|
||||
return {
|
||||
"logits": NeuralType(('B', 'T', 'D'), LogitsType()),
|
||||
"start_positions": NeuralType(tuple('B'), ChannelType()),
|
||||
"end_positions": NeuralType(tuple('B'), ChannelType()),
|
||||
}
|
||||
|
||||
@property
|
||||
def output_types(self):
|
||||
"""Returns definitions of module output ports.
|
||||
"""
|
||||
return {
|
||||
"loss": NeuralType(elements_type=LossType()),
|
||||
"start_logits": NeuralType(('B', 'T'), LogitsType()),
|
||||
"end_logits": NeuralType(('B', 'T'), LogitsType()),
|
||||
}
|
||||
|
||||
def __init__(self,):
|
||||
super().__init__()
|
||||
|
||||
@typecheck()
|
||||
def forward(self, logits, start_positions, end_positions):
|
||||
"""
|
||||
Args:
|
||||
logits: Output of question answering head, which is a token classfier.
|
||||
start_positions: Ground truth start positions of the answer w.r.t.
|
||||
input sequence. If question is unanswerable, this will be
|
||||
pointing to start token, e.g. [CLS], of the input sequence.
|
||||
end_positions: Ground truth end positions of the answer w.r.t.
|
||||
input sequence. If question is unanswerable, this will be
|
||||
pointing to start token, e.g. [CLS], of the input sequence.
|
||||
"""
|
||||
start_logits, end_logits = logits.split(1, dim=-1)
|
||||
start_logits = start_logits.squeeze(-1)
|
||||
end_logits = end_logits.squeeze(-1)
|
||||
# If we are on multi-GPU, split add a dimension
|
||||
if len(start_positions.size()) > 1:
|
||||
start_positions = start_positions.squeeze(-1)
|
||||
if len(end_positions.size()) > 1:
|
||||
end_positions = end_positions.squeeze(-1)
|
||||
ignored_index = start_logits.size(1)
|
||||
start_positions.clamp_(0, ignored_index)
|
||||
end_positions.clamp_(0, ignored_index)
|
||||
|
||||
loss_fct = nn.CrossEntropyLoss(ignore_index=ignored_index)
|
||||
start_loss = loss_fct(start_logits, start_positions)
|
||||
end_loss = loss_fct(end_logits, end_positions)
|
||||
total_loss = (start_loss + end_loss) / 2
|
||||
return total_loss, start_logits, end_logits
|
||||
Reference in New Issue
Block a user