| 130 | |
| 131 | |
| 132 | class CLIPTextModel_(torch.nn.Module): |
| 133 | def __init__(self, config_dict, dtype, device): |
| 134 | num_layers = config_dict["num_hidden_layers"] |
| 135 | embed_dim = config_dict["hidden_size"] |
| 136 | heads = config_dict["num_attention_heads"] |
| 137 | intermediate_size = config_dict["intermediate_size"] |
| 138 | intermediate_activation = config_dict["hidden_act"] |
| 139 | super().__init__() |
| 140 | self.embeddings = CLIPEmbeddings(embed_dim, dtype=torch.float32, device=device) |
| 141 | self.encoder = CLIPEncoder(num_layers, embed_dim, heads, intermediate_size, intermediate_activation, dtype, device) |
| 142 | self.final_layer_norm = nn.LayerNorm(embed_dim, dtype=dtype, device=device) |
| 143 | |
| 144 | def forward(self, input_tokens, intermediate_output=None, final_layer_norm_intermediate=True): |
| 145 | x = self.embeddings(input_tokens) |
| 146 | causal_mask = torch.empty(x.shape[1], x.shape[1], dtype=x.dtype, device=x.device).fill_(float("-inf")).triu_(1) |
| 147 | x, i = self.encoder(x, mask=causal_mask, intermediate_output=intermediate_output) |
| 148 | x = self.final_layer_norm(x) |
| 149 | if i is not None and final_layer_norm_intermediate: |
| 150 | i = self.final_layer_norm(i) |
| 151 | pooled_output = x[torch.arange(x.shape[0], device=x.device), input_tokens.to(dtype=torch.int, device=x.device).argmax(dim=-1),] |
| 152 | return x, i, pooled_output |
| 153 | |
| 154 | |
| 155 | class CLIPTextModel(torch.nn.Module): |