741 lines
32 KiB
Python
741 lines
32 KiB
Python
# Copyright 2022 The HuggingFace Team. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import inspect
|
|
from typing import List, Optional, Tuple, Union
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.utils.checkpoint
|
|
|
|
from transformers.activations import ACT2FN
|
|
from transformers.configuration_utils import PretrainedConfig
|
|
from transformers.modeling_outputs import BaseModelOutput
|
|
from transformers.modeling_utils import PreTrainedModel
|
|
from transformers.tokenization_utils import PreTrainedTokenizer
|
|
from transformers.utils import logging
|
|
|
|
from ...models import AutoencoderKL, UNet2DConditionModel, UNet2DModel, VQModel
|
|
from ...pipeline_utils import DiffusionPipeline, ImagePipelineOutput
|
|
from ...schedulers import DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler
|
|
|
|
|
|
class LDMTextToImagePipeline(DiffusionPipeline):
|
|
r"""
|
|
This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
|
|
library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
|
|
|
|
Parameters:
|
|
vqvae ([`VQModel`]):
|
|
Vector-quantized (VQ) Model to encode and decode images to and from latent representations.
|
|
bert ([`LDMBertModel`]):
|
|
Text-encoder model based on [BERT](https://huggingface.co/docs/transformers/model_doc/bert) architecture.
|
|
tokenizer (`transformers.BertTokenizer`):
|
|
Tokenizer of class
|
|
[BertTokenizer](https://huggingface.co/docs/transformers/model_doc/bert#transformers.BertTokenizer).
|
|
unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents.
|
|
scheduler ([`SchedulerMixin`]):
|
|
A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of
|
|
[`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`].
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
vqvae: Union[VQModel, AutoencoderKL],
|
|
bert: PreTrainedModel,
|
|
tokenizer: PreTrainedTokenizer,
|
|
unet: Union[UNet2DModel, UNet2DConditionModel],
|
|
scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler],
|
|
):
|
|
super().__init__()
|
|
self.register_modules(vqvae=vqvae, bert=bert, tokenizer=tokenizer, unet=unet, scheduler=scheduler)
|
|
self.vae_scale_factor = 2 ** (len(self.vqvae.config.block_out_channels) - 1)
|
|
|
|
@torch.no_grad()
|
|
def __call__(
|
|
self,
|
|
prompt: Union[str, List[str]],
|
|
height: Optional[int] = None,
|
|
width: Optional[int] = None,
|
|
num_inference_steps: Optional[int] = 50,
|
|
guidance_scale: Optional[float] = 1.0,
|
|
eta: Optional[float] = 0.0,
|
|
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
|
|
latents: Optional[torch.FloatTensor] = None,
|
|
output_type: Optional[str] = "pil",
|
|
return_dict: bool = True,
|
|
**kwargs,
|
|
) -> Union[Tuple, ImagePipelineOutput]:
|
|
r"""
|
|
Args:
|
|
prompt (`str` or `List[str]`):
|
|
The prompt or prompts to guide the image generation.
|
|
height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
|
|
The height in pixels of the generated image.
|
|
width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
|
|
The width in pixels of the generated image.
|
|
num_inference_steps (`int`, *optional*, defaults to 50):
|
|
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
|
|
expense of slower inference.
|
|
guidance_scale (`float`, *optional*, defaults to 1.0):
|
|
Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
|
|
`guidance_scale` is defined as `w` of equation 2. of [Imagen
|
|
Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
|
|
1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt` at
|
|
the, usually at the expense of lower image quality.
|
|
generator (`torch.Generator`, *optional*):
|
|
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
|
|
to make generation deterministic.
|
|
latents (`torch.FloatTensor`, *optional*):
|
|
Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
|
|
generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
|
|
tensor will ge generated by sampling using the supplied random `generator`.
|
|
output_type (`str`, *optional*, defaults to `"pil"`):
|
|
The output format of the generate image. Choose between
|
|
[PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
|
|
return_dict (`bool`, *optional*):
|
|
Whether or not to return a [`~pipeline_utils.ImagePipelineOutput`] instead of a plain tuple.
|
|
|
|
Returns:
|
|
[`~pipeline_utils.ImagePipelineOutput`] or `tuple`: [`~pipelines.utils.ImagePipelineOutput`] if
|
|
`return_dict` is True, otherwise a `tuple. When returning a tuple, the first element is a list with the
|
|
generated images.
|
|
"""
|
|
# 0. Default height and width to unet
|
|
height = height or self.unet.config.sample_size * self.vae_scale_factor
|
|
width = width or self.unet.config.sample_size * self.vae_scale_factor
|
|
|
|
if isinstance(prompt, str):
|
|
batch_size = 1
|
|
elif isinstance(prompt, list):
|
|
batch_size = len(prompt)
|
|
else:
|
|
raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
|
|
|
|
if height % 8 != 0 or width % 8 != 0:
|
|
raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
|
|
|
|
# get unconditional embeddings for classifier free guidance
|
|
if guidance_scale != 1.0:
|
|
uncond_input = self.tokenizer(
|
|
[""] * batch_size, padding="max_length", max_length=77, truncation=True, return_tensors="pt"
|
|
)
|
|
uncond_embeddings = self.bert(uncond_input.input_ids.to(self.device))[0]
|
|
|
|
# get prompt text embeddings
|
|
text_input = self.tokenizer(prompt, padding="max_length", max_length=77, truncation=True, return_tensors="pt")
|
|
text_embeddings = self.bert(text_input.input_ids.to(self.device))[0]
|
|
|
|
# get the initial random noise unless the user supplied it
|
|
latents_shape = (batch_size, self.unet.in_channels, height // 8, width // 8)
|
|
if isinstance(generator, list) and len(generator) != batch_size:
|
|
raise ValueError(
|
|
f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
|
|
f" size of {batch_size}. Make sure the batch size matches the length of the generators."
|
|
)
|
|
|
|
if latents is None:
|
|
rand_device = "cpu" if self.device.type == "mps" else self.device
|
|
|
|
if isinstance(generator, list):
|
|
latents_shape = (1,) + latents_shape[1:]
|
|
latents = [
|
|
torch.randn(latents_shape, generator=generator[i], device=rand_device, dtype=text_embeddings.dtype)
|
|
for i in range(batch_size)
|
|
]
|
|
latents = torch.cat(latents, dim=0)
|
|
else:
|
|
latents = torch.randn(
|
|
latents_shape, generator=generator, device=rand_device, dtype=text_embeddings.dtype
|
|
)
|
|
latents = latents.to(self.device)
|
|
else:
|
|
if latents.shape != latents_shape:
|
|
raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {latents_shape}")
|
|
latents = latents.to(self.device)
|
|
|
|
self.scheduler.set_timesteps(num_inference_steps)
|
|
|
|
# prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
|
|
accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
|
|
|
|
extra_kwargs = {}
|
|
if accepts_eta:
|
|
extra_kwargs["eta"] = eta
|
|
|
|
for t in self.progress_bar(self.scheduler.timesteps):
|
|
if guidance_scale == 1.0:
|
|
# guidance_scale of 1 means no guidance
|
|
latents_input = latents
|
|
context = text_embeddings
|
|
else:
|
|
# For classifier free guidance, we need to do two forward passes.
|
|
# Here we concatenate the unconditional and text embeddings into a single batch
|
|
# to avoid doing two forward passes
|
|
latents_input = torch.cat([latents] * 2)
|
|
context = torch.cat([uncond_embeddings, text_embeddings])
|
|
|
|
# predict the noise residual
|
|
noise_pred = self.unet(latents_input, t, encoder_hidden_states=context).sample
|
|
# perform guidance
|
|
if guidance_scale != 1.0:
|
|
noise_pred_uncond, noise_prediction_text = noise_pred.chunk(2)
|
|
noise_pred = noise_pred_uncond + guidance_scale * (noise_prediction_text - noise_pred_uncond)
|
|
|
|
# compute the previous noisy sample x_t -> x_t-1
|
|
latents = self.scheduler.step(noise_pred, t, latents, **extra_kwargs).prev_sample
|
|
|
|
# scale and decode the image latents with vae
|
|
latents = 1 / 0.18215 * latents
|
|
image = self.vqvae.decode(latents).sample
|
|
|
|
image = (image / 2 + 0.5).clamp(0, 1)
|
|
image = image.cpu().permute(0, 2, 3, 1).numpy()
|
|
if output_type == "pil":
|
|
image = self.numpy_to_pil(image)
|
|
|
|
if not return_dict:
|
|
return (image,)
|
|
|
|
return ImagePipelineOutput(images=image)
|
|
|
|
|
|
################################################################################
|
|
# Code for the text transformer model
|
|
################################################################################
|
|
""" PyTorch LDMBERT model."""
|
|
|
|
|
|
logger = logging.get_logger(__name__)
|
|
|
|
LDMBERT_PRETRAINED_MODEL_ARCHIVE_LIST = [
|
|
"ldm-bert",
|
|
# See all LDMBert models at https://huggingface.co/models?filter=ldmbert
|
|
]
|
|
|
|
|
|
LDMBERT_PRETRAINED_CONFIG_ARCHIVE_MAP = {
|
|
"ldm-bert": "https://huggingface.co/valhalla/ldm-bert/blob/main/config.json",
|
|
}
|
|
|
|
|
|
""" LDMBERT model configuration"""
|
|
|
|
|
|
class LDMBertConfig(PretrainedConfig):
|
|
model_type = "ldmbert"
|
|
keys_to_ignore_at_inference = ["past_key_values"]
|
|
attribute_map = {"num_attention_heads": "encoder_attention_heads", "hidden_size": "d_model"}
|
|
|
|
def __init__(
|
|
self,
|
|
vocab_size=30522,
|
|
max_position_embeddings=77,
|
|
encoder_layers=32,
|
|
encoder_ffn_dim=5120,
|
|
encoder_attention_heads=8,
|
|
head_dim=64,
|
|
encoder_layerdrop=0.0,
|
|
activation_function="gelu",
|
|
d_model=1280,
|
|
dropout=0.1,
|
|
attention_dropout=0.0,
|
|
activation_dropout=0.0,
|
|
init_std=0.02,
|
|
classifier_dropout=0.0,
|
|
scale_embedding=False,
|
|
use_cache=True,
|
|
pad_token_id=0,
|
|
**kwargs,
|
|
):
|
|
self.vocab_size = vocab_size
|
|
self.max_position_embeddings = max_position_embeddings
|
|
self.d_model = d_model
|
|
self.encoder_ffn_dim = encoder_ffn_dim
|
|
self.encoder_layers = encoder_layers
|
|
self.encoder_attention_heads = encoder_attention_heads
|
|
self.head_dim = head_dim
|
|
self.dropout = dropout
|
|
self.attention_dropout = attention_dropout
|
|
self.activation_dropout = activation_dropout
|
|
self.activation_function = activation_function
|
|
self.init_std = init_std
|
|
self.encoder_layerdrop = encoder_layerdrop
|
|
self.classifier_dropout = classifier_dropout
|
|
self.use_cache = use_cache
|
|
self.num_hidden_layers = encoder_layers
|
|
self.scale_embedding = scale_embedding # scale factor will be sqrt(d_model) if True
|
|
|
|
super().__init__(pad_token_id=pad_token_id, **kwargs)
|
|
|
|
|
|
def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
|
|
"""
|
|
Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
|
|
"""
|
|
bsz, src_len = mask.size()
|
|
tgt_len = tgt_len if tgt_len is not None else src_len
|
|
|
|
expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)
|
|
|
|
inverted_mask = 1.0 - expanded_mask
|
|
|
|
return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min)
|
|
|
|
|
|
# Copied from transformers.models.bart.modeling_bart.BartAttention with Bart->LDMBert
|
|
class LDMBertAttention(nn.Module):
|
|
"""Multi-headed attention from 'Attention Is All You Need' paper"""
|
|
|
|
def __init__(
|
|
self,
|
|
embed_dim: int,
|
|
num_heads: int,
|
|
head_dim: int,
|
|
dropout: float = 0.0,
|
|
is_decoder: bool = False,
|
|
bias: bool = False,
|
|
):
|
|
super().__init__()
|
|
self.embed_dim = embed_dim
|
|
self.num_heads = num_heads
|
|
self.dropout = dropout
|
|
self.head_dim = head_dim
|
|
self.inner_dim = head_dim * num_heads
|
|
|
|
self.scaling = self.head_dim**-0.5
|
|
self.is_decoder = is_decoder
|
|
|
|
self.k_proj = nn.Linear(embed_dim, self.inner_dim, bias=bias)
|
|
self.v_proj = nn.Linear(embed_dim, self.inner_dim, bias=bias)
|
|
self.q_proj = nn.Linear(embed_dim, self.inner_dim, bias=bias)
|
|
self.out_proj = nn.Linear(self.inner_dim, embed_dim)
|
|
|
|
def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
|
|
return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
key_value_states: Optional[torch.Tensor] = None,
|
|
past_key_value: Optional[Tuple[torch.Tensor]] = None,
|
|
attention_mask: Optional[torch.Tensor] = None,
|
|
layer_head_mask: Optional[torch.Tensor] = None,
|
|
output_attentions: bool = False,
|
|
) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
|
|
"""Input shape: Batch x Time x Channel"""
|
|
|
|
# if key_value_states are provided this layer is used as a cross-attention layer
|
|
# for the decoder
|
|
is_cross_attention = key_value_states is not None
|
|
|
|
bsz, tgt_len, _ = hidden_states.size()
|
|
|
|
# get query proj
|
|
query_states = self.q_proj(hidden_states) * self.scaling
|
|
# get key, value proj
|
|
if is_cross_attention and past_key_value is not None:
|
|
# reuse k,v, cross_attentions
|
|
key_states = past_key_value[0]
|
|
value_states = past_key_value[1]
|
|
elif is_cross_attention:
|
|
# cross_attentions
|
|
key_states = self._shape(self.k_proj(key_value_states), -1, bsz)
|
|
value_states = self._shape(self.v_proj(key_value_states), -1, bsz)
|
|
elif past_key_value is not None:
|
|
# reuse k, v, self_attention
|
|
key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
|
|
value_states = self._shape(self.v_proj(hidden_states), -1, bsz)
|
|
key_states = torch.cat([past_key_value[0], key_states], dim=2)
|
|
value_states = torch.cat([past_key_value[1], value_states], dim=2)
|
|
else:
|
|
# self_attention
|
|
key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
|
|
value_states = self._shape(self.v_proj(hidden_states), -1, bsz)
|
|
|
|
if self.is_decoder:
|
|
# if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states.
|
|
# Further calls to cross_attention layer can then reuse all cross-attention
|
|
# key/value_states (first "if" case)
|
|
# if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of
|
|
# all previous decoder key/value_states. Further calls to uni-directional self-attention
|
|
# can concat previous decoder key/value_states to current projected key/value_states (third "elif" case)
|
|
# if encoder bi-directional self-attention `past_key_value` is always `None`
|
|
past_key_value = (key_states, value_states)
|
|
|
|
proj_shape = (bsz * self.num_heads, -1, self.head_dim)
|
|
query_states = self._shape(query_states, tgt_len, bsz).view(*proj_shape)
|
|
key_states = key_states.view(*proj_shape)
|
|
value_states = value_states.view(*proj_shape)
|
|
|
|
src_len = key_states.size(1)
|
|
attn_weights = torch.bmm(query_states, key_states.transpose(1, 2))
|
|
|
|
if attn_weights.size() != (bsz * self.num_heads, tgt_len, src_len):
|
|
raise ValueError(
|
|
f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is"
|
|
f" {attn_weights.size()}"
|
|
)
|
|
|
|
if attention_mask is not None:
|
|
if attention_mask.size() != (bsz, 1, tgt_len, src_len):
|
|
raise ValueError(
|
|
f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is {attention_mask.size()}"
|
|
)
|
|
attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + attention_mask
|
|
attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
|
|
|
|
attn_weights = nn.functional.softmax(attn_weights, dim=-1)
|
|
|
|
if layer_head_mask is not None:
|
|
if layer_head_mask.size() != (self.num_heads,):
|
|
raise ValueError(
|
|
f"Head mask for a single layer should be of size {(self.num_heads,)}, but is"
|
|
f" {layer_head_mask.size()}"
|
|
)
|
|
attn_weights = layer_head_mask.view(1, -1, 1, 1) * attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
|
|
attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
|
|
|
|
if output_attentions:
|
|
# this operation is a bit awkward, but it's required to
|
|
# make sure that attn_weights keeps its gradient.
|
|
# In order to do so, attn_weights have to be reshaped
|
|
# twice and have to be reused in the following
|
|
attn_weights_reshaped = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
|
|
attn_weights = attn_weights_reshaped.view(bsz * self.num_heads, tgt_len, src_len)
|
|
else:
|
|
attn_weights_reshaped = None
|
|
|
|
attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)
|
|
|
|
attn_output = torch.bmm(attn_probs, value_states)
|
|
|
|
if attn_output.size() != (bsz * self.num_heads, tgt_len, self.head_dim):
|
|
raise ValueError(
|
|
f"`attn_output` should be of size {(bsz, self.num_heads, tgt_len, self.head_dim)}, but is"
|
|
f" {attn_output.size()}"
|
|
)
|
|
|
|
attn_output = attn_output.view(bsz, self.num_heads, tgt_len, self.head_dim)
|
|
attn_output = attn_output.transpose(1, 2)
|
|
|
|
# Use the `embed_dim` from the config (stored in the class) rather than `hidden_state` because `attn_output` can be
|
|
# partitioned across GPUs when using tensor-parallelism.
|
|
attn_output = attn_output.reshape(bsz, tgt_len, self.inner_dim)
|
|
|
|
attn_output = self.out_proj(attn_output)
|
|
|
|
return attn_output, attn_weights_reshaped, past_key_value
|
|
|
|
|
|
class LDMBertEncoderLayer(nn.Module):
|
|
def __init__(self, config: LDMBertConfig):
|
|
super().__init__()
|
|
self.embed_dim = config.d_model
|
|
self.self_attn = LDMBertAttention(
|
|
embed_dim=self.embed_dim,
|
|
num_heads=config.encoder_attention_heads,
|
|
head_dim=config.head_dim,
|
|
dropout=config.attention_dropout,
|
|
)
|
|
self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
|
|
self.dropout = config.dropout
|
|
self.activation_fn = ACT2FN[config.activation_function]
|
|
self.activation_dropout = config.activation_dropout
|
|
self.fc1 = nn.Linear(self.embed_dim, config.encoder_ffn_dim)
|
|
self.fc2 = nn.Linear(config.encoder_ffn_dim, self.embed_dim)
|
|
self.final_layer_norm = nn.LayerNorm(self.embed_dim)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.FloatTensor,
|
|
attention_mask: torch.FloatTensor,
|
|
layer_head_mask: torch.FloatTensor,
|
|
output_attentions: Optional[bool] = False,
|
|
) -> Tuple[torch.FloatTensor, Optional[torch.FloatTensor]]:
|
|
"""
|
|
Args:
|
|
hidden_states (`torch.FloatTensor`): input to the layer of shape `(seq_len, batch, embed_dim)`
|
|
attention_mask (`torch.FloatTensor`): attention mask of size
|
|
`(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
|
|
layer_head_mask (`torch.FloatTensor`): mask for attention heads in a given layer of size
|
|
`(encoder_attention_heads,)`.
|
|
output_attentions (`bool`, *optional*):
|
|
Whether or not to return the attentions tensors of all attention layers. See `attentions` under
|
|
returned tensors for more detail.
|
|
"""
|
|
residual = hidden_states
|
|
hidden_states = self.self_attn_layer_norm(hidden_states)
|
|
hidden_states, attn_weights, _ = self.self_attn(
|
|
hidden_states=hidden_states,
|
|
attention_mask=attention_mask,
|
|
layer_head_mask=layer_head_mask,
|
|
output_attentions=output_attentions,
|
|
)
|
|
hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
|
|
hidden_states = residual + hidden_states
|
|
|
|
residual = hidden_states
|
|
hidden_states = self.final_layer_norm(hidden_states)
|
|
hidden_states = self.activation_fn(self.fc1(hidden_states))
|
|
hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
|
|
hidden_states = self.fc2(hidden_states)
|
|
hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
|
|
hidden_states = residual + hidden_states
|
|
|
|
if hidden_states.dtype == torch.float16 and (
|
|
torch.isinf(hidden_states).any() or torch.isnan(hidden_states).any()
|
|
):
|
|
clamp_value = torch.finfo(hidden_states.dtype).max - 1000
|
|
hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)
|
|
|
|
outputs = (hidden_states,)
|
|
|
|
if output_attentions:
|
|
outputs += (attn_weights,)
|
|
|
|
return outputs
|
|
|
|
|
|
# Copied from transformers.models.bart.modeling_bart.BartPretrainedModel with Bart->LDMBert
|
|
class LDMBertPreTrainedModel(PreTrainedModel):
|
|
config_class = LDMBertConfig
|
|
base_model_prefix = "model"
|
|
_supports_gradient_checkpointing = True
|
|
_keys_to_ignore_on_load_unexpected = [r"encoder\.version", r"decoder\.version"]
|
|
|
|
def _init_weights(self, module):
|
|
std = self.config.init_std
|
|
if isinstance(module, nn.Linear):
|
|
module.weight.data.normal_(mean=0.0, std=std)
|
|
if module.bias is not None:
|
|
module.bias.data.zero_()
|
|
elif isinstance(module, nn.Embedding):
|
|
module.weight.data.normal_(mean=0.0, std=std)
|
|
if module.padding_idx is not None:
|
|
module.weight.data[module.padding_idx].zero_()
|
|
|
|
def _set_gradient_checkpointing(self, module, value=False):
|
|
if isinstance(module, (LDMBertEncoder,)):
|
|
module.gradient_checkpointing = value
|
|
|
|
@property
|
|
def dummy_inputs(self):
|
|
pad_token = self.config.pad_token_id
|
|
input_ids = torch.tensor([[0, 6, 10, 4, 2], [0, 8, 12, 2, pad_token]], device=self.device)
|
|
dummy_inputs = {
|
|
"attention_mask": input_ids.ne(pad_token),
|
|
"input_ids": input_ids,
|
|
}
|
|
return dummy_inputs
|
|
|
|
|
|
class LDMBertEncoder(LDMBertPreTrainedModel):
|
|
"""
|
|
Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a
|
|
[`LDMBertEncoderLayer`].
|
|
|
|
Args:
|
|
config: LDMBertConfig
|
|
embed_tokens (nn.Embedding): output embedding
|
|
"""
|
|
|
|
def __init__(self, config: LDMBertConfig):
|
|
super().__init__(config)
|
|
|
|
self.dropout = config.dropout
|
|
|
|
embed_dim = config.d_model
|
|
self.padding_idx = config.pad_token_id
|
|
self.max_source_positions = config.max_position_embeddings
|
|
|
|
self.embed_tokens = nn.Embedding(config.vocab_size, embed_dim)
|
|
self.embed_positions = nn.Embedding(config.max_position_embeddings, embed_dim)
|
|
self.layers = nn.ModuleList([LDMBertEncoderLayer(config) for _ in range(config.encoder_layers)])
|
|
self.layer_norm = nn.LayerNorm(embed_dim)
|
|
|
|
self.gradient_checkpointing = False
|
|
# Initialize weights and apply final processing
|
|
self.post_init()
|
|
|
|
def get_input_embeddings(self):
|
|
return self.embed_tokens
|
|
|
|
def set_input_embeddings(self, value):
|
|
self.embed_tokens = value
|
|
|
|
def forward(
|
|
self,
|
|
input_ids: torch.LongTensor = None,
|
|
attention_mask: Optional[torch.Tensor] = None,
|
|
position_ids: Optional[torch.LongTensor] = None,
|
|
head_mask: Optional[torch.Tensor] = None,
|
|
inputs_embeds: Optional[torch.FloatTensor] = None,
|
|
output_attentions: Optional[bool] = None,
|
|
output_hidden_states: Optional[bool] = None,
|
|
return_dict: Optional[bool] = None,
|
|
) -> Union[Tuple, BaseModelOutput]:
|
|
r"""
|
|
Args:
|
|
input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
|
|
Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you
|
|
provide it.
|
|
|
|
Indices can be obtained using [`BartTokenizer`]. See [`PreTrainedTokenizer.encode`] and
|
|
[`PreTrainedTokenizer.__call__`] for details.
|
|
|
|
[What are input IDs?](../glossary#input-ids)
|
|
attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
|
|
Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
|
|
|
|
- 1 for tokens that are **not masked**,
|
|
- 0 for tokens that are **masked**.
|
|
|
|
[What are attention masks?](../glossary#attention-mask)
|
|
head_mask (`torch.Tensor` of shape `(encoder_layers, encoder_attention_heads)`, *optional*):
|
|
Mask to nullify selected heads of the attention modules. Mask values selected in `[0, 1]`:
|
|
|
|
- 1 indicates the head is **not masked**,
|
|
- 0 indicates the head is **masked**.
|
|
|
|
inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
|
|
Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation.
|
|
This is useful if you want more control over how to convert `input_ids` indices into associated vectors
|
|
than the model's internal embedding lookup matrix.
|
|
output_attentions (`bool`, *optional*):
|
|
Whether or not to return the attentions tensors of all attention layers. See `attentions` under
|
|
returned tensors for more detail.
|
|
output_hidden_states (`bool`, *optional*):
|
|
Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
|
|
for more detail.
|
|
return_dict (`bool`, *optional*):
|
|
Whether or not to return a [`~utils.BaseModelOutput`] instead of a plain tuple.
|
|
"""
|
|
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
|
output_hidden_states = (
|
|
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
|
)
|
|
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
|
|
|
# retrieve input_ids and inputs_embeds
|
|
if input_ids is not None and inputs_embeds is not None:
|
|
raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
|
|
elif input_ids is not None:
|
|
input_shape = input_ids.size()
|
|
input_ids = input_ids.view(-1, input_shape[-1])
|
|
elif inputs_embeds is not None:
|
|
input_shape = inputs_embeds.size()[:-1]
|
|
else:
|
|
raise ValueError("You have to specify either input_ids or inputs_embeds")
|
|
|
|
if inputs_embeds is None:
|
|
inputs_embeds = self.embed_tokens(input_ids)
|
|
|
|
seq_len = input_shape[1]
|
|
if position_ids is None:
|
|
position_ids = torch.arange(seq_len, dtype=torch.long, device=inputs_embeds.device).expand((1, -1))
|
|
embed_pos = self.embed_positions(position_ids)
|
|
|
|
hidden_states = inputs_embeds + embed_pos
|
|
hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
|
|
|
|
# expand attention_mask
|
|
if attention_mask is not None:
|
|
# [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
|
|
attention_mask = _expand_mask(attention_mask, inputs_embeds.dtype)
|
|
|
|
encoder_states = () if output_hidden_states else None
|
|
all_attentions = () if output_attentions else None
|
|
|
|
# check if head_mask has a correct number of layers specified if desired
|
|
if head_mask is not None:
|
|
if head_mask.size()[0] != (len(self.layers)):
|
|
raise ValueError(
|
|
f"The head_mask should be specified for {len(self.layers)} layers, but it is for"
|
|
f" {head_mask.size()[0]}."
|
|
)
|
|
|
|
for idx, encoder_layer in enumerate(self.layers):
|
|
if output_hidden_states:
|
|
encoder_states = encoder_states + (hidden_states,)
|
|
if self.gradient_checkpointing and self.training:
|
|
|
|
def create_custom_forward(module):
|
|
def custom_forward(*inputs):
|
|
return module(*inputs, output_attentions)
|
|
|
|
return custom_forward
|
|
|
|
layer_outputs = torch.utils.checkpoint.checkpoint(
|
|
create_custom_forward(encoder_layer),
|
|
hidden_states,
|
|
attention_mask,
|
|
(head_mask[idx] if head_mask is not None else None),
|
|
)
|
|
else:
|
|
layer_outputs = encoder_layer(
|
|
hidden_states,
|
|
attention_mask,
|
|
layer_head_mask=(head_mask[idx] if head_mask is not None else None),
|
|
output_attentions=output_attentions,
|
|
)
|
|
|
|
hidden_states = layer_outputs[0]
|
|
|
|
if output_attentions:
|
|
all_attentions = all_attentions + (layer_outputs[1],)
|
|
|
|
hidden_states = self.layer_norm(hidden_states)
|
|
|
|
if output_hidden_states:
|
|
encoder_states = encoder_states + (hidden_states,)
|
|
|
|
if not return_dict:
|
|
return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None)
|
|
return BaseModelOutput(
|
|
last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions
|
|
)
|
|
|
|
|
|
class LDMBertModel(LDMBertPreTrainedModel):
|
|
_no_split_modules = []
|
|
|
|
def __init__(self, config: LDMBertConfig):
|
|
super().__init__(config)
|
|
self.model = LDMBertEncoder(config)
|
|
self.to_logits = nn.Linear(config.hidden_size, config.vocab_size)
|
|
|
|
def forward(
|
|
self,
|
|
input_ids=None,
|
|
attention_mask=None,
|
|
position_ids=None,
|
|
head_mask=None,
|
|
inputs_embeds=None,
|
|
output_attentions=None,
|
|
output_hidden_states=None,
|
|
return_dict=None,
|
|
):
|
|
outputs = self.model(
|
|
input_ids,
|
|
attention_mask=attention_mask,
|
|
position_ids=position_ids,
|
|
head_mask=head_mask,
|
|
inputs_embeds=inputs_embeds,
|
|
output_attentions=output_attentions,
|
|
output_hidden_states=output_hidden_states,
|
|
return_dict=return_dict,
|
|
)
|
|
return outputs
|