| import torch |
| from torch import nn |
| from transformers import PreTrainedModel |
| from transformers.models.dinov2.modeling_dinov2 import Dinov2Encoder |
|
|
| from .configuration_radzero import AlignTransformerConfig |
|
|
|
|
| def build_align_transformer(config): |
| if config.model_type == "align_transformer": |
| model = AlignTransformer(config) |
| else: |
| raise NotImplementedError() |
|
|
| return model |
|
|
|
|
| class AlignTransformer(PreTrainedModel): |
| def __init__(self, config: AlignTransformerConfig): |
| super().__init__(config) |
|
|
| self.projector = None |
|
|
| if config.num_hidden_layers: |
| self.transformer_layers = Dinov2Encoder(config) |
| else: |
| self.transformer_layers = None |
|
|
| if config.use_layer_norm: |
| self.layer_norm = nn.LayerNorm(config.hidden_size) |
| else: |
| self.layer_norm = None |
|
|
| def forward(self, vision_tokens): |
|
|
| if self.projector is not None: |
|
|
| cls_token = vision_tokens[:, :1] |
| patch_tokens = vision_tokens[:, 1:] |
|
|
| patch_tokens = self.projector(patch_tokens)["last_hidden_state"] |
| vision_tokens = torch.cat([cls_token, patch_tokens], dim=1) |
|
|
| if self.transformer_layers is not None: |
| vision_tokens = self.transformer_layers(vision_tokens)["last_hidden_state"] |
|
|
| if self.layer_norm is not None: |
| vision_tokens = self.layer_norm(vision_tokens) |
|
|
| return vision_tokens |
|
|
