(self, text: torch.Tensor, normalize: bool = False)
| 293 | return F.normalize(features, dim=-1) if normalize else features |
| 294 | |
| 295 | def encode_text(self, text: torch.Tensor, normalize: bool = False) -> torch.Tensor: |
| 296 | cast_dtype = self.transformer.get_cast_dtype() |
| 297 | |
| 298 | x = self.token_embedding(text).to(cast_dtype) |
| 299 | |
| 300 | x = x + self.positional_embedding.to(cast_dtype) |
| 301 | x = self.transformer(x, attn_mask=self.attn_mask) |
| 302 | x = self.ln_final(x) |
| 303 | |
| 304 | x = text_global_pool(x, text, self.text_pool_type) |
| 305 | |
| 306 | if self.text_projection is not None: |
| 307 | if isinstance(self.text_projection, nn.Linear): |
| 308 | x = self.text_projection(x) |
| 309 | else: |
| 310 | x = x @ self.text_projection |
| 311 | |
| 312 | return F.normalize(x, dim=-1) if normalize else x |
| 313 | |
| 314 | def get_logits(self, image: torch.Tensor, text: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: |
| 315 | image_features = self.encode_image(image, normalize=True) |
no test coverage detected