Call the meliad model.
(
self,
inputs: np.ndarray,
dstate: tuple[dict[str, np.ndarray], ...] = None,
eos: np.ndarray = None,
mask: np.ndarray = None,
)
| 101 | return [x[0] for x in result] |
| 102 | |
| 103 | def call( |
| 104 | self, |
| 105 | inputs: np.ndarray, |
| 106 | dstate: tuple[dict[str, np.ndarray], ...] = None, |
| 107 | eos: np.ndarray = None, |
| 108 | mask: np.ndarray = None, |
| 109 | ) -> MetricsOutput: |
| 110 | """Call the meliad model.""" |
| 111 | batch_size, length = inputs.shape |
| 112 | inputs = jax.numpy.pad(inputs, [(0, 0), (0, 1024 - length)]) |
| 113 | |
| 114 | if eos is None: |
| 115 | eos = self.eos |
| 116 | if mask is None: |
| 117 | mask = self.mask |
| 118 | |
| 119 | x = {'targets': inputs, 'length': length, 'eos': eos, 'mask': mask} |
| 120 | |
| 121 | if dstate is not None: |
| 122 | x['start_of_sequence'] = jax.numpy.array([False] * batch_size) |
| 123 | else: |
| 124 | dstate = tuple( |
| 125 | [{ # this dummy value will never be used. |
| 126 | 'current_index': np.array([0] * batch_size, dtype=np.int32), |
| 127 | 'keys': np.zeros( |
| 128 | (batch_size, 2048, self.n, self.h), dtype=np.bfloat16 |
| 129 | ), |
| 130 | 'values': np.zeros( |
| 131 | (batch_size, 2048, self.n, self.h), dtype=np.bfloat16 |
| 132 | ), |
| 133 | 'recurrent_kvq': None, |
| 134 | 'relative_position_bias': np.zeros( |
| 135 | (batch_size, self.n, 1, 1024), dtype=np.bfloat16 |
| 136 | ), |
| 137 | }] |
| 138 | * 12 |
| 139 | ) |
| 140 | x['start_of_sequence'] = jax.numpy.array([True] * batch_size) |
| 141 | |
| 142 | x['dstate'] = dstate |
| 143 | _, metrics_np = self.task.run_step(self.tstate, x, 0) |
| 144 | return metrics_np |
| 145 | |
| 146 | def beam_decode( |
| 147 | self, |