Core implementation of the image-text autoregressive Transformer. Attributes: fusion_style: Method to combine image and text features. 'concat': Concatenate image and text embeddings and process with a standard decoder. 'cross_attn': Use the
| 359 | |
| 360 | |
| 361 | class _Model(nn.Module): |
| 362 | """Core implementation of the image-text autoregressive Transformer. |
| 363 | |
| 364 | Attributes: |
| 365 | fusion_style: Method to combine image and text features. |
| 366 | 'concat': Concatenate image and text embeddings and process with a |
| 367 | standard decoder. |
| 368 | 'cross_attn': Use the interleaved self/cross-attention encoder. |
| 369 | vocab_size: The size of the text vocabulary. |
| 370 | """ |
| 371 | width: int = 512 |
| 372 | depth: int = 12 |
| 373 | mlp_dim: Optional[int] = None |
| 374 | num_heads: int = 12 |
| 375 | dropout: float = 0.0 |
| 376 | drop_path: float = 0.0 |
| 377 | remat_policy: str = 'none' |
| 378 | fusion_style: str = 'cross_attn' |
| 379 | casual_mask: bool = True |
| 380 | use_flash_attn: bool = False |
| 381 | dtype: Optional[Dtype] = jnp.float32 |
| 382 | param_dtype: Dtype = jnp.float32 |
| 383 | mesh: Optional[Any] = None |
| 384 | vocab_size: int = 32000 |
| 385 | |
| 386 | @nn.compact |
| 387 | def __call__(self, text_input: Array, context: Array, *, train: bool = False) -> Array: |
| 388 | out = {} |
| 389 | # The model predicts the next token, so we use tokens up to the second to last. |
| 390 | token_ids_in = text_input[:, :-1] |
| 391 | |
| 392 | embedding = nn.Embed( |
| 393 | num_embeddings=self.vocab_size, |
| 394 | features=self.width, |
| 395 | dtype=jnp.float32, # Use float32 for stability |
| 396 | param_dtype=self.param_dtype, |
| 397 | embedding_init=nn.with_logical_partitioning( |
| 398 | nn.initializers.normal(stddev=0.02), ('vocab', 'embed'))) |
| 399 | text_embeds = embedding(token_ids_in.astype("int32")) |
| 400 | |
| 401 | # Project image features to the same dimension as text embeddings. |
| 402 | _, _, img_dim = context.shape |
| 403 | image_projection_layer = nn.Dense( |
| 404 | self.width, |
| 405 | name="image_projection_layer", |
| 406 | use_bias=False, |
| 407 | kernel_init=nn.initializers.normal(stddev=img_dim ** -0.5)) |
| 408 | image_embeds = image_projection_layer(context) |
| 409 | |
| 410 | # --- FUSION OF IMAGE AND TEXT FEATURES --- |
| 411 | if self.fusion_style == 'concat': |
| 412 | img_len = image_embeds.shape[1] |
| 413 | image_text_embeds = jnp.concatenate((image_embeds, text_embeds), axis=1) |
| 414 | |
| 415 | # Use a standard causal decoder over the concatenated sequence. |
| 416 | decoder_blocks = Encoder( |
| 417 | name="Transformer", **self.get_transformer_kwargs()) |
| 418 | x, _ = decoder_blocks(image_text_embeds, deterministic=not train) |