Spaces:
Runtime error
Runtime error
| import torch | |
| import torch.nn as nn | |
| from tencentpretrain.utils.constants import * | |
| class LmTarget(nn.Module): | |
| """ | |
| Language Model Target | |
| """ | |
| def __init__(self, args, vocab_size): | |
| super(LmTarget, self).__init__() | |
| self.vocab_size = vocab_size | |
| self.hidden_size = args.hidden_size | |
| if "label_smoothing" in args: | |
| self.label_smoothing = args.label_smoothing | |
| else: | |
| self.label_smoothing = None | |
| if "ignore_index" in args and args.ignore_index: | |
| self.ignore_index = args.tokenizer.vocab.get(PAD_TOKEN) | |
| else: | |
| self.ignore_index = None | |
| self.output_layer = nn.Linear(self.hidden_size, self.vocab_size, bias=args.has_lmtarget_bias) | |
| self.softmax = nn.LogSoftmax(dim=-1) | |
| self.criterion = nn.NLLLoss() | |
| def lm(self, memory_bank, tgt_lm): | |
| # Language modeling (LM) with full softmax prediction. | |
| tgt_lm = tgt_lm.contiguous().view(-1) | |
| memory_bank = memory_bank.contiguous().view(-1, self.hidden_size) | |
| memory_bank = memory_bank[tgt_lm > 0, :] | |
| tgt_lm = tgt_lm[tgt_lm > 0] | |
| output = self.output_layer(memory_bank) | |
| output = self.softmax(output) | |
| denominator = torch.tensor(output.size(0) + 1e-6) | |
| if output.size(0) == 0: | |
| correct = torch.tensor(0.0) | |
| else: | |
| correct = torch.sum((output.argmax(dim=-1).eq(tgt_lm)).float()) | |
| if self.label_smoothing is None: | |
| loss = self.criterion(output, tgt_lm) | |
| else: | |
| if tgt_lm.dim() == output.dim() - 1: | |
| tgt_lm = tgt_lm.unsqueeze(-1) | |
| nll_loss = -output.gather(dim=-1, index=tgt_lm) | |
| smooth_loss = -output.sum(dim=-1, keepdim=True) | |
| if self.ignore_index is not None: | |
| pad_mask = tgt_lm.eq(self.ignore_index) | |
| nll_loss.masked_fill_(pad_mask, 0.0) | |
| smooth_loss.masked_fill_(pad_mask, 0.0) | |
| else: | |
| nll_loss = nll_loss.squeeze(-1) | |
| smooth_loss = smooth_loss.squeeze(-1) | |
| nll_loss = nll_loss.mean() | |
| smooth_loss = smooth_loss.mean() | |
| eps_i = self.label_smoothing / (output.size(-1) - 1) | |
| loss = (1.0 - self.label_smoothing - eps_i) * nll_loss + eps_i * smooth_loss | |
| return loss, correct, denominator | |
| def forward(self, memory_bank, tgt, seg): | |
| """ | |
| Args: | |
| memory_bank: [batch_size x seq_length x hidden_size] | |
| tgt: [batch_size x seq_length] | |
| Returns: | |
| loss: Language modeling loss. | |
| correct: Number of words that are predicted correctly. | |
| denominator: Number of predicted words. | |
| """ | |
| # Language modeling (LM) with full softmax prediction. | |
| loss, correct, denominator = self.lm(memory_bank, tgt) | |
| return loss, correct, denominator | |