Initialises any state which a generation step transforms.
(
self,
*args, # pylint: disable=unused-argument
rng: PRNGKeyType | None = None,
**kwargs, # pylint: disable=unused-argument
)
| 1519 | raise ValueError(f"Unsupported tokenizer type: {metadata.tokenizer_type}") |
| 1520 | |
| 1521 | def init_decode_state( |
| 1522 | self, |
| 1523 | *args, # pylint: disable=unused-argument |
| 1524 | rng: PRNGKeyType | None = None, |
| 1525 | **kwargs, # pylint: disable=unused-argument |
| 1526 | ) -> DecodeState: |
| 1527 | """Initialises any state which a generation step transforms.""" |
| 1528 | if rng is None: |
| 1529 | rng = jax.random.PRNGKey(0) |
| 1530 | page_state = None |
| 1531 | if self.config.attention == "paged" and self.page_manager is not None: |
| 1532 | page_state = self.page_manager.get_initial_page_state() # pytype: disable=attribute-error |
| 1533 | |
| 1534 | # pylint: disable=unused-argument |
| 1535 | def init(abstract_params, page_state): |
| 1536 | x = jnp.ones( |
| 1537 | (int(self.config.per_device_batch_size * self.mesh.size), 1), |
| 1538 | dtype=jnp.int32, |
| 1539 | ) |
| 1540 | dummy_image = jnp.ones( |
| 1541 | multimodal_utils.get_dummy_image_shape_for_init( |
| 1542 | self.config.model_name, batch_size=self.config.micro_batch_size_to_train_on |
| 1543 | ), |
| 1544 | dtype=jnp.int32, |
| 1545 | ) |
| 1546 | _, cache = self.model.apply( |
| 1547 | abstract_params, |
| 1548 | x, |
| 1549 | x, |
| 1550 | encoder_images=dummy_image if self.config.use_multimodal else None, |
| 1551 | enable_dropout=False, |
| 1552 | model_mode=MODEL_MODE_AUTOREGRESSIVE, |
| 1553 | rngs={"params": rng}, |
| 1554 | mutable=["cache"], |
| 1555 | page_state=page_state, |
| 1556 | slot=0, |
| 1557 | ) |
| 1558 | |
| 1559 | next_pos = jnp.zeros( |
| 1560 | (int(self.config.per_device_batch_size * self.mesh.size), 1), |
| 1561 | dtype=jnp.int32, |
| 1562 | ) |
| 1563 | generated_tokens = jnp.zeros( |
| 1564 | (int(self.config.per_device_batch_size * self.mesh.size), 1), |
| 1565 | dtype=jnp.int32, |
| 1566 | ) |
| 1567 | tokens = jnp.zeros( |
| 1568 | (int(self.config.per_device_batch_size * self.mesh.size), 1), |
| 1569 | dtype=jnp.int32, |
| 1570 | ) |
| 1571 | token_logp = jnp.zeros( |
| 1572 | (int(self.config.per_device_batch_size * self.mesh.size), 1), |
| 1573 | dtype=jnp.float32, |
| 1574 | ) |
| 1575 | return { |
| 1576 | "logits": jnp.zeros( |
| 1577 | ( |
| 1578 | int(self.config.per_device_batch_size * self.mesh.size), |