A image-text autoregression Transformer model.
| 412 | |
| 413 | |
| 414 | class _Model(nn.Module): |
| 415 | """A image-text autoregression Transformer model.""" |
| 416 | num_classes: int = None |
| 417 | width: int = 512 |
| 418 | depth: int = 12 |
| 419 | mlp_dim: Optional[int] = None # Defaults to 4x input dim |
| 420 | num_heads: int = 12 |
| 421 | dropout: float = 0.0 |
| 422 | remat_policy: str = 'none' |
| 423 | fusion_style: str = 'cross_attn' |
| 424 | scan_mlp: bool = False |
| 425 | scan_attn: bool = False |
| 426 | mlp_chunck: int = 128 |
| 427 | casual_mask: bool = True |
| 428 | use_flash_attn: bool = False |
| 429 | dtype: Optional[Dtype] = jnp.float32 |
| 430 | param_dtype: Dtype = jnp.float32 |
| 431 | mesh: Any = None |
| 432 | drop_path: float = 0.0 |
| 433 | num_learnable_tokens: int = 80 |
| 434 | drop_token: int = 0 |
| 435 | |
| 436 | @nn.compact |
| 437 | def __call__(self, image_embeds, text_embeds, *, train=False): |
| 438 | out = {} |
| 439 | |
| 440 | if self.drop_token > 0: |
| 441 | image_embeds = image_embeds[:, :image_embeds.shape[1] - self.drop_token + 1] |
| 442 | |
| 443 | ni, li, di = image_embeds.shape |
| 444 | nt, lt, dt = text_embeds.shape |
| 445 | assert ni == nt, f'the image embed is {image_embeds.shape} and text embed is {text_embeds.shape}' |
| 446 | |
| 447 | image_projection_layer = nn.Dense( |
| 448 | self.width, |
| 449 | name="image_projection_layer", |
| 450 | use_bias=False, |
| 451 | kernel_init=nn.initializers.normal(stddev=di ** -0.5)) |
| 452 | image_embeds = out["projected_image_embeds"] = image_projection_layer(image_embeds) |
| 453 | |
| 454 | text_projection_layer = nn.Dense( |
| 455 | self.width, |
| 456 | name="text_projection_layer", |
| 457 | use_bias=False, |
| 458 | kernel_init=nn.initializers.normal(stddev=dt ** -0.5)) |
| 459 | text_embeds = out["projected_text_embeds"] = text_projection_layer(text_embeds) |
| 460 | |
| 461 | learnable_tokens = self.param( |
| 462 | 'learnable_tokens', |
| 463 | nn.initializers.normal(stddev=1.0), |
| 464 | (self.num_learnable_tokens, self.width) |
| 465 | ) |
| 466 | learnable_tokens = jnp.tile(learnable_tokens[None, :, :], (ni, 1, 1)) |
| 467 | |
| 468 | image_embeds = nn.with_logical_constraint(image_embeds, ("activation_batch", "activation_length", "activation_embed")) |
| 469 | text_embeds = nn.with_logical_constraint(text_embeds, ("activation_batch", "activation_length", "activation_embed")) |
| 470 | |
| 471 | # concatenate image_embeds and learnable_tokens in token dimension |