Generate an output sequence. Args: inputs: the same as argument to _call_. Returns: An array of generated tokens of shape (batch_size, sequence_length).
(
self, inputs: ...
)
| 115 | } |
| 116 | |
| 117 | def generate( |
| 118 | self, inputs: ... |
| 119 | ) -> tuple[tuple[dict[str, jnp.ndarray, ...], ...], jnp.ndarray]: |
| 120 | """Generate an output sequence. |
| 121 | |
| 122 | Args: |
| 123 | inputs: the same as argument to _call_. |
| 124 | |
| 125 | Returns: |
| 126 | An array of generated tokens of shape (batch_size, sequence_length). |
| 127 | """ |
| 128 | input_tokens = inputs['targets'] # [b,seq_len] |
| 129 | start_of_sequence = inputs['start_of_sequence'] # [b] |
| 130 | target_tokens = jnp.pad(input_tokens[:, 1:], [(0, 0), (0, 1)]) |
| 131 | batch_size = target_tokens.shape[0] |
| 132 | |
| 133 | # Assuming all sequences start at the same time. |
| 134 | start0 = inputs['start_of_sequence'][0] |
| 135 | dstate = jax.lax.cond( |
| 136 | start0, |
| 137 | lambda: self.decoder.init_decoder_state_vanilla( # pylint: disable=g-long-lambda |
| 138 | 1024, start_of_sequence |
| 139 | ), |
| 140 | lambda: inputs['dstate'], |
| 141 | ) |
| 142 | |
| 143 | first_token = input_tokens[:, 0:1] |
| 144 | no_start_of_seq = jnp.array([False] * batch_size, dtype=jnp.bool_) |
| 145 | temperature = 1 |
| 146 | if 'temperature' in inputs: |
| 147 | temperature = inputs['temperature'] |
| 148 | |
| 149 | num_steps = inputs['length'] |
| 150 | if self.mode == 'beam_search': |
| 151 | num_steps -= 1 |
| 152 | |
| 153 | def cond_fn(scan_state) -> jnp.bool_: |
| 154 | _, _, i, _ = scan_state |
| 155 | return i < num_steps |
| 156 | |
| 157 | def loop_fn(scan_state: Any) -> Tuple[Any, Any, Any, Any]: |
| 158 | (dstate, input_token, i, _) = scan_state |
| 159 | |
| 160 | (logits, dstate, _) = self.decoder( |
| 161 | input_tokens=input_token, |
| 162 | target_tokens=None, |
| 163 | start_of_sequence=no_start_of_seq, |
| 164 | decoder_state=dstate, |
| 165 | ) |
| 166 | |
| 167 | logits = logits / temperature |
| 168 | output_token = jax.lax.dynamic_slice_in_dim(target_tokens, i, 1, axis=1) |
| 169 | |
| 170 | return (dstate, output_token, i + 1, logits) |
| 171 | |
| 172 | # Scan over the sequence length. |
| 173 | dummy_logits = jnp.zeros((batch_size, 1, 1024)) |
| 174 | initial_scan_state = (dstate, first_token, 0, dummy_logits) |
no test coverage detected