Generates autoregressive sequences based on the given start tokens. Args: start_tokens (torch.Tensor): The initial tokens to start the generation. seq_len (int): The length of the generated sequence. eos_token (int, optional): The end-of-sequence
(
self,
start_tokens,
seq_len,
eos_token=None,
temperature=1.0,
filter_thres=0.9,
**kwargs,
)
| 51 | @torch.no_grad() |
| 52 | @eval_decorator |
| 53 | def generate( |
| 54 | self, |
| 55 | start_tokens, |
| 56 | seq_len, |
| 57 | eos_token=None, |
| 58 | temperature=1.0, |
| 59 | filter_thres=0.9, |
| 60 | **kwargs, |
| 61 | ): |
| 62 | """ |
| 63 | Generates autoregressive sequences based on the given start tokens. |
| 64 | |
| 65 | Args: |
| 66 | start_tokens (torch.Tensor): The initial tokens to start the generation. |
| 67 | seq_len (int): The length of the generated sequence. |
| 68 | eos_token (int, optional): The end-of-sequence token. If provided, generation will stop when this token is generated. Defaults to None. |
| 69 | temperature (float, optional): The temperature value for controlling the randomness of the generation. Higher values result in more randomness. Defaults to 1.0. |
| 70 | filter_thres (float, optional): The threshold value for filtering logits during generation. Only logits above this threshold will be considered. Defaults to 0.9. |
| 71 | **kwargs: Additional keyword arguments to be passed to the underlying network. |
| 72 | |
| 73 | Returns: |
| 74 | torch.Tensor: The generated sequence. |
| 75 | """ |
| 76 | |
| 77 | b, t, device = *start_tokens.shape, start_tokens.device |
| 78 | |
| 79 | out = start_tokens |
| 80 | |
| 81 | for _ in range(seq_len): |
| 82 | logits = self.net(out, **kwargs)[:, -1, :] |
| 83 | |
| 84 | filtered_logits = top_k(logits, thres=filter_thres) |
| 85 | probs = F.softmax(filtered_logits / temperature, dim=-1) |
| 86 | |
| 87 | sample = torch.multinomial(probs, 1) |
| 88 | |
| 89 | out = torch.cat((out, sample), dim=-1) |
| 90 | |
| 91 | if exists(eos_token): |
| 92 | is_eos_token = out == eos_token |
| 93 | |
| 94 | if is_eos_token.any(dim=-1).all(): |
| 95 | # mask out everything after the eos tokens |
| 96 | shifted_is_eos_tokens = F.pad(is_eos_token, (1, -1)) |
| 97 | mask = shifted_is_eos_tokens.float().cumsum(dim=-1) >= 1 |
| 98 | out = out.masked_fill(mask, self.pad_value) |
| 99 | break |
| 100 | |
| 101 | out = out[:, t:] |
| 102 | return out |
| 103 | |
| 104 | def forward(self, x, **kwargs): |
| 105 | x_inp, x_labels = x[:, :-1], x[:, 1:] |