-
Notifications
You must be signed in to change notification settings - Fork 2.4k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Canary greedy and temperature decoding #8885
Changes from all commits
60cafd9
b3b23dd
e5a82a2
5ad0e69
8e6b220
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -15,6 +15,7 @@ | |
from contextlib import contextmanager | ||
|
||
import torch | ||
from torch.distributions import Categorical | ||
|
||
from nemo.collections.common.parts import NEG_INF, mask_padded_tokens | ||
|
||
|
@@ -34,8 +35,8 @@ class GreedySequenceGenerator: | |
Args: | ||
embedding: nn.Module, transforms input_ids into vector embeddings | ||
decoder: nn.Module, takes embeddings and produces hidden_states | ||
log_softmax: nn.Module, takes hidden_states and produces log_probs | ||
which correspond to probability distribution of tokens (ids) | ||
classifier: nn.Module, takes hidden_states and produces | ||
logits or log-probability distribution of tokens (ids) | ||
pad: index of padding token in the vocabulary | ||
bos: index of beginning of sequence token in the vocabulary | ||
eos: index of end of sequence token in the vocabulary | ||
|
@@ -51,22 +52,26 @@ def __init__( | |
self, | ||
embedding, | ||
decoder, | ||
log_softmax, | ||
classifier, | ||
pad=0, | ||
bos=1, | ||
eos=2, | ||
max_sequence_length=512, | ||
max_delta_length=20, | ||
batch_size=1, | ||
n_samples=1, | ||
temperature=None, | ||
): | ||
super().__init__() | ||
self.embedding = embedding | ||
self.decoder = decoder | ||
self.log_softmax = log_softmax | ||
self.classifier = classifier.set_log_softmax_enabled(False) | ||
self.pad, self.bos, self.eos = pad, bos, eos | ||
self.max_seq_length = max_sequence_length | ||
self.max_delta_len = max_delta_length | ||
self.batch_size = batch_size | ||
self.n_samples = n_samples | ||
self.temperature = temperature | ||
|
||
def _one_step_forward( | ||
self, | ||
|
@@ -107,8 +112,8 @@ def _one_step_forward( | |
decoder_mems_list = self.decoder.forward( | ||
decoder_hidden_states, decoder_input_mask, decoder_mems_list, return_mems=True | ||
) | ||
log_probs = self.log_softmax.forward(hidden_states=decoder_mems_list[-1][:, -1:]) | ||
return log_probs, decoder_mems_list | ||
logits = self.classifier.forward(hidden_states=decoder_mems_list[-1][:, -1:], temperature=self.temperature) | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. shall we add a check to see if the forward function has |
||
return logits, decoder_mems_list | ||
|
||
def _prepare_for_search(self, decoder_input_ids=None, encoder_hidden_states=None): | ||
""" | ||
|
@@ -145,30 +150,49 @@ def _forward( | |
self, decoder_input_ids=None, encoder_hidden_states=None, encoder_input_mask=None, return_beam_scores=False | ||
): | ||
assert not return_beam_scores | ||
is_sampling = self.temperature is not None and self.n_samples > 1 | ||
|
||
tgt, batch_size, max_generation_length = self._prepare_for_search(decoder_input_ids, encoder_hidden_states) | ||
if is_sampling: | ||
tgt = torch.repeat_interleave(tgt, self.n_samples, dim=0) | ||
encoder_hidden_states = torch.repeat_interleave(encoder_hidden_states, self.n_samples, dim=0) | ||
encoder_input_mask = torch.repeat_interleave(encoder_input_mask, self.n_samples, dim=0) | ||
orig_batch_size = batch_size | ||
batch_size = batch_size * self.n_samples | ||
|
||
# pad profile tracks sequences ending with <eos> token to replace | ||
# everything after <eos> with <pad> token | ||
decoder_parameter = next(self.decoder.parameters()) | ||
pad_profile = torch.zeros(batch_size, 1).long().to(decoder_parameter.device) | ||
pad_profile = torch.zeros(batch_size).long().to(decoder_parameter.device) | ||
|
||
decoder_mems_list = None | ||
for i in range(max_generation_length): | ||
|
||
log_probs, decoder_mems_list = self._one_step_forward( | ||
tgt[:, -1:], encoder_hidden_states, encoder_input_mask, decoder_mems_list, i | ||
if i == 0: | ||
input_ids = tgt | ||
else: | ||
input_ids = tgt[:, -1:] | ||
logits, decoder_mems_list = self._one_step_forward( | ||
input_ids, encoder_hidden_states, encoder_input_mask, decoder_mems_list, i | ||
) | ||
if self.temperature is None: | ||
next_tokens = torch.argmax(logits[:, -1], dim=-1) | ||
else: | ||
next_tokens = Categorical(logits=logits[:, -1] / self.temperature).sample() | ||
|
||
next_tokens = torch.argmax(log_probs[:, -1], dim=-1, keepdim=True) | ||
next_tokens = self.pad * pad_profile + next_tokens * (1 - pad_profile) | ||
pad_profile = torch.max(pad_profile, (next_tokens == self.eos).long()) | ||
tgt = torch.cat((tgt, next_tokens), dim=-1) | ||
tgt = torch.cat((tgt, next_tokens.unsqueeze(1)), dim=-1) | ||
|
||
# abort generation if all sequences end with <eos> | ||
if pad_profile.sum() == batch_size: | ||
break | ||
|
||
return tgt | ||
samples = None | ||
if is_sampling: | ||
samples = list(tgt.view(orig_batch_size, self.n_samples, -1)) | ||
tgt = tgt[:: self.n_samples] | ||
|
||
return tgt, samples | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. since we're adding an additional output |
||
|
||
def __call__( | ||
self, decoder_input_ids=None, encoder_hidden_states=None, encoder_input_mask=None, return_beam_scores=False | ||
|
@@ -196,9 +220,9 @@ def freeze(self) -> None: | |
for param in self.decoder.parameters(): | ||
param.requires_grad = False | ||
self.decoder.eval() | ||
for param in self.log_softmax.parameters(): | ||
for param in self.classifier.parameters(): | ||
param.requires_grad = False | ||
self.log_softmax.eval() | ||
self.classifier.eval() | ||
|
||
def unfreeze(self) -> None: | ||
"""Unfreeze weights of embedding, decoder, and classification layers. | ||
|
@@ -209,14 +233,14 @@ def unfreeze(self) -> None: | |
for param in self.decoder.parameters(): | ||
param.requires_grad = True | ||
self.decoder.train() | ||
for param in self.log_softmax.parameters(): | ||
for param in self.classifier.parameters(): | ||
param.requires_grad = True | ||
self.log_softmax.train() | ||
self.classifier.train() | ||
|
||
@contextmanager | ||
def as_frozen(self): | ||
""" | ||
Context manager which temporarily freezes embedding, decoder, and log_softmax modules, | ||
Context manager which temporarily freezes embedding, decoder, and classifier modules, | ||
yields control and finally unfreezes the modules. | ||
""" | ||
self.freeze() | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Shall we add a check to see if the classifier has the set_log_softmax_enabled() function, and if not, we default to not using temperature sampling and print out a warning?