(self, config: CLIPTextModelConfig)
| 63 | """Implements the text encoder transformer from CLIP.""" |
| 64 | |
| 65 | def __init__(self, config: CLIPTextModelConfig): |
| 66 | super().__init__() |
| 67 | |
| 68 | self.token_embedding = nn.Embedding(config.vocab_size, config.model_dims) |
| 69 | self.position_embedding = nn.Embedding(config.max_length, config.model_dims) |
| 70 | self.layers = [ |
| 71 | CLIPEncoderLayer(config.model_dims, config.num_heads, config.hidden_act) |
| 72 | for i in range(config.num_layers) |
| 73 | ] |
| 74 | self.final_layer_norm = nn.LayerNorm(config.model_dims) |
| 75 | |
| 76 | if config.projection_dim is not None: |
| 77 | self.text_projection = nn.Linear( |
| 78 | config.model_dims, config.projection_dim, bias=False |
| 79 | ) |
| 80 | |
| 81 | def _get_mask(self, N, dtype): |
| 82 | indices = mx.arange(N) |
no test coverage detected