| 59 | return fake_input_dict |
| 60 | |
| 61 | def __call__(self, inputs: ...) -> tuple[Any, dict[str, Any]]: |
| 62 | # Make sure this code is not used on untested cases. |
| 63 | if self.mode not in ['init', 'beam_search']: |
| 64 | raise ValueError(f'{type(self)} cannot do mode {self.mode}') |
| 65 | if self.decoder.supports_generate(): |
| 66 | raise ValueError(f'{type(self)}.decoder cannot supports_generate()') |
| 67 | |
| 68 | self.decoder( |
| 69 | input_tokens=inputs['targets'][:, 0:1], |
| 70 | target_tokens=None, |
| 71 | start_of_sequence=inputs['start_of_sequence'], |
| 72 | ) |
| 73 | |
| 74 | b = inputs['targets'].shape[0] |
| 75 | no_start_of_seq = jnp.array([False] * b, dtype=jnp.bool_) |
| 76 | |
| 77 | # This fn is used in both beam_search or topk_sampling. |
| 78 | def tokens_to_logits_fn( |
| 79 | input_token: jnp.ndarray, dstate: tuple[dict[str, jnp.ndarray], ...] |
| 80 | ) -> tuple[jnp.ndarray, tuple[dict[str, jnp.ndarray], ...]]: |
| 81 | (logits, dstate, _) = self.decoder( |
| 82 | input_tokens=input_token, |
| 83 | target_tokens=None, |
| 84 | start_of_sequence=no_start_of_seq, |
| 85 | decoder_state=dstate, |
| 86 | ) |
| 87 | return logits[:, -1, :], dstate |
| 88 | |
| 89 | last_token = jax.lax.dynamic_slice_in_dim( |
| 90 | inputs['targets'], inputs['length'] - 1, 1, axis=1 |
| 91 | ) |
| 92 | |
| 93 | # last token is used to seed beam_search |
| 94 | inputs['targets'] = inputs['targets'][:, 0:-1] |
| 95 | dstate = jax.lax.cond( |
| 96 | inputs['start_of_sequence'][0], |
| 97 | lambda: self.generate(inputs)[0], |
| 98 | lambda: inputs['dstate'], |
| 99 | ) |
| 100 | |
| 101 | # Then we run beam search, init with last_token & dstate. |
| 102 | finished_seqs, finished_scores, dstate = beam_search.beam_search_flat( |
| 103 | last_token, |
| 104 | dstate, |
| 105 | tokens_to_logits_fn, |
| 106 | max_decode_len=512, |
| 107 | eos=inputs['eos'].reshape((1, 1, -1)), |
| 108 | mask=inputs['mask'].reshape((1, 1, -1)), |
| 109 | ) |
| 110 | |
| 111 | return 0.0, { |
| 112 | 'finished_seqs': finished_seqs, |
| 113 | 'finished_scores': finished_scores, |
| 114 | 'dstate': dstate, |
| 115 | } |
| 116 | |
| 117 | def generate( |
| 118 | self, inputs: ... |