Spaces:
Runtime error
Runtime error
| import os | |
| import math | |
| import logging | |
| import json | |
| from pathlib import Path | |
| from typing import Tuple, Optional, Union | |
| import torch | |
| import torch.nn as nn | |
| from torch.nn import functional as F | |
| from transformers import AutoModelForCausalLM, AutoTokenizer, GPTJForCausalLM | |
| logging.basicConfig(level=os.environ.get("LOGLEVEL", "INFO")) | |
| log = logging.getLogger(__name__) | |
| def load_weights(self, Module, path, name, default_name, prev_name=None, **kwargs): | |
| hparams = None | |
| assert isinstance(default_name, str), f'invalid default transformer name: {default_name}' | |
| model = get_transformer_module(Module, default_name, **kwargs) | |
| setattr(self, name, model) | |
| return hparams | |
| def get_transformer_module(Module, default_name, **kwargs): | |
| if default_name == 'EleutherAI/gpt-j-6B': | |
| kwargs = {**kwargs, **dict(revision="float16", torch_dtype=torch.float16, low_cpu_mem_usage=True)} | |
| model = Module.from_pretrained(default_name, **kwargs) | |
| return model | |
| class MLP(nn.Module): | |
| def __init__(self, sizes: Tuple[int, ...], bias=True, act=nn.Tanh): | |
| super(MLP, self).__init__() | |
| self.divider = math.sqrt(sizes[-1] / sizes[0]) | |
| layers = [] | |
| for i in range(len(sizes) - 1): | |
| layers.append(nn.Linear(sizes[i], sizes[i + 1], bias=bias)) | |
| if i < len(sizes) - 2: | |
| layers.append(act()) | |
| self.model = nn.Sequential(*layers) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| x = x / self.divider # scaling for the initial stability | |
| x = self.model(x) | |
| return x | |
| class MlpTransformer(nn.Module): | |
| def __init__(self, in_dim, h_dim, out_d: Optional[int] = None, act=F.relu, dropout=0.): | |
| super().__init__() | |
| out_d = out_d if out_d is not None else in_dim | |
| self.fc1 = nn.Linear(in_dim, h_dim) | |
| self.act = act | |
| self.fc2 = nn.Linear(h_dim, out_d) | |
| self.dropout = nn.Dropout(dropout) | |
| def forward(self, x): | |
| x = self.fc1(x) | |
| x = self.act(x) | |
| x = self.dropout(x) | |
| x = self.fc2(x) | |
| x = self.dropout(x) | |
| return x | |
| class MultiHeadAttention(nn.Module): | |
| def __init__(self, dim_self, dim_ref, num_heads, bias=True, dropout=0.): | |
| super().__init__() | |
| self.num_heads = num_heads | |
| head_dim = dim_self // num_heads | |
| self.scale = head_dim ** -0.5 | |
| self.to_queries = nn.Linear(dim_self, dim_self, bias=bias) | |
| self.to_keys_values = nn.Linear(dim_ref, dim_self * 2, bias=bias) | |
| self.project = nn.Linear(dim_self, dim_self) | |
| self.dropout = nn.Dropout(dropout) | |
| def forward(self, x, y=None, mask=None): | |
| y = y if y is not None else x | |
| b, n, c = x.shape | |
| _, m, d = y.shape | |
| # b n h dh | |
| queries = self.to_queries(x).reshape(b, n, self.num_heads, c // self.num_heads) | |
| # b m 2 h dh | |
| keys_values = self.to_keys_values(y).reshape(b, m, 2, self.num_heads, c // self.num_heads) | |
| keys, values = keys_values[:, :, 0], keys_values[:, :, 1] | |
| attention = torch.einsum('bnhd,bmhd->bnmh', queries, keys) * self.scale | |
| if mask is not None: | |
| if mask.dim() == 2: | |
| mask = mask.unsqueeze(1) | |
| attention = attention.masked_fill(mask.unsqueeze(3), float("-inf")) | |
| attention = attention.softmax(dim=2) | |
| out = torch.einsum('bnmh,bmhd->bnhd', attention, values).reshape(b, n, c) | |
| out = self.project(out) | |
| return out, attention | |
| class TransformerLayer(nn.Module): | |
| def __init__(self, dim_self, dim_ref, num_heads, mlp_ratio=4., bias=False, dropout=0., act=F.relu, | |
| norm_layer: nn.Module = nn.LayerNorm): | |
| super().__init__() | |
| self.norm1 = norm_layer(dim_self) | |
| self.attn = MultiHeadAttention(dim_self, dim_ref, num_heads, bias=bias, dropout=dropout) | |
| self.norm2 = norm_layer(dim_self) | |
| self.mlp = MlpTransformer(dim_self, int(dim_self * mlp_ratio), act=act, dropout=dropout) | |
| def forward_with_attention(self, x, y=None, mask=None): | |
| x_, attention = self.attn(self.norm1(x), y, mask) | |
| x = x + x_ | |
| x = x + self.mlp(self.norm2(x)) | |
| return x, attention | |
| def forward(self, x, y=None, mask=None): | |
| x = x + self.attn(self.norm1(x), y, mask)[0] | |
| x = x + self.mlp(self.norm2(x)) | |
| return x | |
| class Transformer(nn.Module): | |
| def forward_with_attention(self, x, y=None, mask=None): | |
| attentions = [] | |
| for layer in self.layers: | |
| x, att = layer.forward_with_attention(x, y, mask) | |
| attentions.append(att) | |
| return x, attentions | |
| def forward(self, x, y=None, mask=None): | |
| for i, layer in enumerate(self.layers): | |
| if i % 2 == 0 and self.enc_dec: # cross | |
| x = layer(x, y) | |
| elif self.enc_dec: # self | |
| x = layer(x, x, mask) | |
| else: # self or cross | |
| x = layer(x, y, mask) | |
| return x | |
| def __init__(self, dim_self: int, num_heads: int, num_layers: int, dim_ref: Optional[int] = None, | |
| mlp_ratio: float = 2., act=F.relu, norm_layer: nn.Module = nn.LayerNorm, enc_dec: bool = False): | |
| super(Transformer, self).__init__() | |
| dim_ref = dim_ref if dim_ref is not None else dim_self | |
| self.enc_dec = enc_dec | |
| if enc_dec: | |
| num_layers = num_layers * 2 | |
| layers = [] | |
| for i in range(num_layers): | |
| if i % 2 == 0 and enc_dec: # cross | |
| layers.append(TransformerLayer(dim_self, dim_ref, num_heads, mlp_ratio, act=act, norm_layer=norm_layer)) | |
| elif enc_dec: # self | |
| layers.append(TransformerLayer(dim_self, dim_self, num_heads, mlp_ratio, act=act, norm_layer=norm_layer)) | |
| else: # self or cross | |
| layers.append(TransformerLayer(dim_self, dim_ref, num_heads, mlp_ratio, act=act, norm_layer=norm_layer)) | |
| self.layers = nn.ModuleList(layers) | |
| class TransformerMapper(nn.Module): | |
| def __init__(self, dim_clip: int, dim_embedding: int, prefix_length: int = 10, | |
| clip_length: int = 10, num_layers: int = 8): | |
| super(TransformerMapper, self).__init__() | |
| self.clip_length = clip_length | |
| self.transformer = Transformer(dim_embedding, 8, num_layers) | |
| self.linear = nn.Linear(dim_clip, clip_length * dim_embedding) | |
| self.prefix_const = nn.Parameter(torch.randn(prefix_length, dim_embedding), requires_grad=True) | |
| def forward(self, x): | |
| x = self.linear(x).view(x.shape[0], self.clip_length, -1) | |
| prefix = self.prefix_const.unsqueeze(0).expand(x.shape[0], *self.prefix_const.shape) | |
| prefix = torch.cat((x, prefix), dim=1) | |
| out = self.transformer(prefix)[:, self.clip_length:] | |
| return out | |
| class ClipCap(nn.Module): | |
| def __init__(self, model_name, device, prefix_length: int = 10, clip_length: int = 40, prefix_size: int = 512, | |
| num_layers: int = 1, model_path: str = '', fix_gpt: bool = False, | |
| use_label_prefix: bool = False, label_path: str = '', label_length: int = 10, | |
| use_transformer_mapper: bool = False, use_ptuning_v2: bool = False, | |
| dropout: float = 0, | |
| model_weight: str = '', scalar_output: bool = False): | |
| super(ClipCap, self).__init__() | |
| self.prefix_length = prefix_length | |
| self.prefix_size = prefix_size | |
| self.label_length = label_length | |
| self.scalar_output = scalar_output | |
| self.num_layers = num_layers | |
| self.use_transformer_mapper = use_transformer_mapper | |
| self.use_ptuning_v2 = use_ptuning_v2 | |
| self.dropout = nn.Dropout(dropout) | |
| hparams = load_weights(self, AutoModelForCausalLM, model_weight, 'gpt', model_name, | |
| prev_name='model') | |
| self.device = device | |
| self.gpt = self.gpt.to(self.device) | |
| config = self.gpt.config | |
| self.match_n_layer = getattr(config, 'n_layer', getattr(config, 'num_layers', None)) # gpt2 vs. gpt_neo | |
| self.match_n_head = getattr(config, 'n_head', getattr(config, 'num_heads', None)) | |
| self.n_embd = getattr(config, 'n_embd', getattr(config, 'hidden_size', None)) | |
| self.match_n_embd = self.n_embd // self.match_n_head | |
| self.clip_project = self.get_mapper() | |
| if Path(label_path).is_file(): | |
| with open(label_path) as f: | |
| labels = json.load(f) | |
| self.labels = {i: v for v, i in labels.items()} | |
| if not use_label_prefix: | |
| log.info("adding label projections") | |
| self.label_project = nn.Sequential( | |
| nn.Embedding(len(self.labels), self.prefix_size), | |
| self.get_mapper() | |
| ) | |
| if os.path.isfile(model_path): | |
| log.info(f"loading model from {model_path}") | |
| weight = torch.load(model_path, map_location=torch.device('cpu')) | |
| weight = {k[len('clip_project.'):]: v for k, v in weight.items() | |
| if k.startswith('clip_project.')} | |
| self.clip_project.load_state_dict(weight) | |
| if fix_gpt: | |
| log.info("fixing gpt parameters") | |
| for param in self.gpt.parameters(): | |
| param.requires_grad_(False) | |
| if self.scalar_output: | |
| self.gpt.lm_head = nn.Linear(self.gpt.transformer.embed_dim, 1).to(self.device) | |
| self.clip_project = self.clip_project.to(self.device) | |
| if hasattr(self, 'label_project'): | |
| self.label_project = self.label_project.to(self.device) | |
| def get_mapper(self): | |
| if self.use_ptuning_v2: | |
| total_embd = self.match_n_layer * 2 * self.n_embd | |
| module = MLP((self.prefix_size, | |
| *[self.prefix_size | |
| for i in range(self.num_layers)], | |
| total_embd * self.prefix_length)) | |
| elif self.use_transformer_mapper: | |
| log.info("using transformer mapper") | |
| module = TransformerMapper(self.prefix_size, self.n_embd, | |
| self.prefix_length, self.prefix_length, num_layers=self.num_layers) # 8) | |
| else: | |
| module = MLP((self.prefix_size, | |
| *[(self.n_embd * self.prefix_length) // 2 | |
| for i in range(self.num_layers)], | |
| self.n_embd * self.prefix_length)) | |
| return module | |
| def get_encoder_loss(self, input_ids: torch.Tensor, features: torch.Tensor, | |
| device = None): | |
| input_ids = input_ids[:, :self.prefix_length].to(device) | |
| embedding = self.gpt.transformer.wte(input_ids) | |
| features = features.to(device) | |
| prefix_projections = self.clip_project(features.type_as(embedding)).reshape(-1, self.prefix_length, self.n_embd) | |
| fct = nn.MSELoss() | |
| loss = fct(prefix_projections, embedding.detach()) | |
| return loss | |
| def forward(self, *args, **kwargs): | |
| if self.use_ptuning_v2: | |
| return self.forward_prefix(*args, **kwargs) | |
| else: | |
| return self.forward_embedding(*args, **kwargs) | |
| def forward_embedding(self, input_ids: torch.Tensor, features: torch.Tensor, | |
| attention_mask: Optional[torch.Tensor] = None, | |
| labels: Optional[torch.Tensor] = None, | |
| past_key_values = None, device = None, **kwargs): | |
| if device is None: | |
| device = self.device | |
| input_ids = input_ids.to(device) | |
| if features is not None: | |
| features = features.to(device) | |
| if attention_mask is not None: | |
| attention_mask = attention_mask.to(device) | |
| if labels is not None: | |
| labels = labels.to(device) | |
| use_labels = labels is not None and hasattr(self, 'label_project') | |
| embedding = self.gpt.transformer.wte(input_ids) | |
| embed_txt = embedding | |
| prefix_length = self.prefix_length | |
| if use_labels: | |
| prefix_length += self.label_length | |
| if past_key_values is None: | |
| prefix_projections = self.clip_project(features.type_as(embedding)).reshape(-1, self.prefix_length, self.n_embd) | |
| if use_labels: | |
| label_projections = self.label_project(labels.long()).reshape(-1, self.label_length, self.n_embd) | |
| prefix_projections = torch.cat((prefix_projections, label_projections), dim=1) | |
| embedding = torch.cat((prefix_projections.to(embedding.dtype), embedding), dim=1) | |
| if torch.is_tensor(attention_mask): | |
| prefix_mask = torch.ones_like(attention_mask)[:, :1].repeat(1, prefix_length) | |
| attention_mask = torch.cat([prefix_mask, attention_mask], dim=1) | |
| outputs = self.gpt(inputs_embeds=embedding, attention_mask=attention_mask, | |
| past_key_values=past_key_values, | |
| return_dict=True, | |
| output_attentions=False, | |
| output_hidden_states=True) | |
| if past_key_values is None: | |
| outputs.logits = outputs.logits[:, prefix_length:] | |
| return outputs | |
| def forward_prefix(self, input_ids: torch.Tensor, features: torch.Tensor, | |
| attention_mask: Optional[torch.Tensor] = None, | |
| labels: Optional[torch.Tensor] = None, | |
| past_key_values = None, device = None, **kwargs): | |
| if device is None: | |
| device = self.device | |
| input_ids = input_ids.to(device) | |
| if features is not None: | |
| features = features.to(device) | |
| if attention_mask is not None: | |
| attention_mask = attention_mask.to(device) | |
| if labels is not None: | |
| labels = labels.to(device) | |
| use_labels = labels is not None and hasattr(self, 'label_project') | |
| prefix_length = self.prefix_length | |
| if use_labels: | |
| prefix_length += self.label_length | |
| if past_key_values is None: | |
| prefix_projections = self.clip_project(features.type_as(self.clip_project.model[0].weight)) | |
| prefix_projections = prefix_projections.reshape(-1, self.prefix_length, | |
| self.match_n_layer * 2, self.match_n_head, self.match_n_embd) | |
| if use_labels: | |
| label_projections = self.label_project(labels.long()) | |
| label_projections = label_projections.reshape(-1, self.label_length, | |
| self.match_n_layer * 2, self.match_n_head, self.match_n_embd) | |
| prefix_projections = torch.cat((prefix_projections, label_projections), dim=1) | |
| temp_control = prefix_projections | |
| temp_control = self.dropout(temp_control) | |
| past_key_values = temp_control.permute([2, 0, 3, 1, 4]).split(2) | |
| if torch.is_tensor(attention_mask): | |
| prefix_mask = torch.ones_like(attention_mask)[:, :1].repeat(1, prefix_length) | |
| attention_mask = torch.cat([prefix_mask, attention_mask], dim=1) | |
| outputs = self.gpt(input_ids=input_ids, attention_mask=attention_mask, | |
| past_key_values=past_key_values, | |
| return_dict=True, | |
| output_attentions=False, | |
| output_hidden_states=True) | |
| if past_key_values is None: | |
| outputs.logits = outputs.logits[:, prefix_length:] | |
| return outputs | |
| def prepare_inputs_for_generation(self, input_ids, past=None, **kwargs): | |
| token_type_ids = kwargs.get("token_type_ids", None) | |
| # only last token for inputs_ids if past is defined in kwargs | |
| if past: | |
| input_ids = input_ids[:, -1].unsqueeze(-1) | |
| if token_type_ids is not None: | |
| token_type_ids = token_type_ids[:, -1].unsqueeze(-1) | |
| attention_mask = kwargs.get("attention_mask", None) | |
| position_ids = kwargs.get("position_ids", None) | |
| features = kwargs.get("features", None) | |
| labels = kwargs.get("labels", None) | |
| if attention_mask is not None and position_ids is None: | |
| # create position_ids on the fly for batch generation | |
| position_ids = attention_mask.long().cumsum(-1) - 1 | |
| position_ids.masked_fill_(attention_mask == 0, 1) | |
| if past: | |
| position_ids = position_ids[:, -1].unsqueeze(-1) | |
| else: | |
| position_ids = None | |
| return { | |
| "input_ids": input_ids, | |
| "past_key_values": past, | |
| "use_cache": kwargs.get("use_cache"), | |
| "position_ids": position_ids, | |
| "attention_mask": attention_mask, | |
| "token_type_ids": token_type_ids, | |
| "features": features, | |
| "labels": labels, | |
| } | |