Takes a conditioning sequence (prompt) as input and continues to generate as many tokens as requested. The implementation of this function is modified from A. Karpathy's nanoGPT. Args: model: The model to use. idx: Tensor of shape (T) with indices of the prompt sequence.
(
model,
idx: torch.Tensor,
max_new_tokens: int,
*,
max_seq_length: Optional[int] = None,
temperature: float = 1.0,
top_k: Optional[int] = None,
eos_id: Optional[int] = None,
)
| 74 | # Referenced from https://github.com/Lightning-AI/lit-llama/blob/main/generate.py#L19 |
| 75 | @torch.no_grad() |
| 76 | def generate( |
| 77 | model, |
| 78 | idx: torch.Tensor, |
| 79 | max_new_tokens: int, |
| 80 | *, |
| 81 | max_seq_length: Optional[int] = None, |
| 82 | temperature: float = 1.0, |
| 83 | top_k: Optional[int] = None, |
| 84 | eos_id: Optional[int] = None, |
| 85 | ) -> torch.Tensor: |
| 86 | """Takes a conditioning sequence (prompt) as input and continues to generate as many tokens as requested. |
| 87 | |
| 88 | The implementation of this function is modified from A. Karpathy's nanoGPT. |
| 89 | |
| 90 | Args: |
| 91 | model: The model to use. |
| 92 | idx: Tensor of shape (T) with indices of the prompt sequence. |
| 93 | max_new_tokens: The number of new tokens to generate. |
| 94 | max_seq_length: The maximum sequence length allowed. |
| 95 | temperature: Scales the predicted logits by 1 / temperature |
| 96 | top_k: If specified, only sample among the tokens with the k highest probabilities |
| 97 | eos_id: If specified, stop generating any more token once the <eos> token is triggered |
| 98 | """ |
| 99 | # create an empty tensor of the expected final shape and fill in the current tokens |
| 100 | T = idx.size(0) |
| 101 | T_new = T + max_new_tokens |
| 102 | if max_seq_length is None: |
| 103 | max_seq_length = min(T_new, model.config.block_size) |
| 104 | |
| 105 | device, dtype = idx.device, idx.dtype |
| 106 | # create an empty tensor of the expected final shape and fill in the current tokens |
| 107 | empty = torch.empty(T_new, dtype=dtype, device=device) |
| 108 | empty[:T] = idx |
| 109 | idx = empty |
| 110 | input_pos = torch.arange(0, T, device=device) |
| 111 | |
| 112 | if idx.device.type == "xla": |
| 113 | import torch_xla.core.xla_model as xm |
| 114 | |
| 115 | xm.mark_step() |
| 116 | |
| 117 | # generate max_new_tokens tokens |
| 118 | for _ in range(max_new_tokens): |
| 119 | x = idx.index_select(0, input_pos).view(1, -1) |
| 120 | |
| 121 | # forward |
| 122 | logits = model(x, input_pos=input_pos, max_seq_length=max_seq_length) |
| 123 | logits = logits[0, -1] / temperature |
| 124 | |
| 125 | # optionally crop the logits to only the top k options |
| 126 | if top_k is not None: |
| 127 | v, _ = torch.topk(logits, min(top_k, logits.size(-1))) |
| 128 | logits = torch.where(logits < v[[-1]], -float("Inf"), logits) |
| 129 | |
| 130 | probs = torch.nn.functional.softmax(logits, dim=-1) |
| 131 | idx_next = torch.multinomial(probs, num_samples=1).to(dtype=dtype) |
| 132 | |
| 133 | # advance |