| 347 | return self.visual(image.type(self.dtype)) |
| 348 | |
| 349 | def encode_text(self, text): |
| 350 | x = self.token_embedding(text).type(self.dtype) # [batch_size, n_ctx, d_model] |
| 351 | |
| 352 | x = x + self.positional_embedding.type(self.dtype) |
| 353 | x = x.permute(1, 0, 2) # NLD -> LND |
| 354 | x = self.transformer(x) |
| 355 | x = x.permute(1, 0, 2) # LND -> NLD |
| 356 | x = self.ln_final(x).type(self.dtype) |
| 357 | |
| 358 | # x.shape = [batch_size, n_ctx, transformer.width] |
| 359 | # take features from the eot embedding (eot_token is the highest number in each sequence) |
| 360 | x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection |
| 361 | |
| 362 | return x |
| 363 | |
| 364 | def forward(self, image, text): |
| 365 | image_features = self.encode_image(image) |