





from dataclasses import dataclass
from typing import Optional

import numpy as np
import torch
import torch.nn as nn

from ..utils import BaseOutput, is_torch_version, randn_tensor
from .attention_processor import SpatialNorm
from .unet_ad_blocks import UNetMidBlockaD, get_down_block, get_up_block

@dataclass
class        __SCREAMING_SNAKE_CASE					(     _a       ):
							snake_case     :      torch.FloatTensor






class        __SCREAMING_SNAKE_CASE					(     nn.Module       ):






							def __init__(						self  , __lowerCAmelCase=3  , __lowerCAmelCase=3  , __lowerCAmelCase=("DownEncoderBlock2D",)  , __lowerCAmelCase=(64,)  , __lowerCAmelCase=2  , __lowerCAmelCase=32  , __lowerCAmelCase="silu"  , __lowerCAmelCase=True  , ):
														super().__init__()
														UpperCamelCase__				     =			layers_per_block

														UpperCamelCase__				     =			torch.nn.Convad(
														    __lowerCAmelCase  , block_out_channels[0]  , kernel_size=3  , stride=1  , padding=1  , )

														UpperCamelCase__				     =			None
														UpperCamelCase__				     =			nn.ModuleList([]  )

														# down
														UpperCamelCase__				     =			block_out_channels[0]
														for i, down_block_type in enumerate(__lowerCAmelCase  ):
																					UpperCamelCase__				     =			output_channel
																					UpperCamelCase__				     =			block_out_channels[i]
																					UpperCamelCase__				     =			i == len(__lowerCAmelCase  ) - 1

																					UpperCamelCase__				     =			get_down_block(
																					    __lowerCAmelCase  , num_layers=self.layers_per_block  , in_channels=__lowerCAmelCase  , out_channels=__lowerCAmelCase  , add_downsample=not is_final_block  , resnet_eps=1E-6  , downsample_padding=0  , resnet_act_fn=__lowerCAmelCase  , resnet_groups=__lowerCAmelCase  , attention_head_dim=__lowerCAmelCase  , temb_channels=__lowerCAmelCase  , )
																					self.down_blocks.append(__lowerCAmelCase  )

														# mid
														UpperCamelCase__				     =			UNetMidBlockaD(
														    in_channels=block_out_channels[-1]  , resnet_eps=1E-6  , resnet_act_fn=__lowerCAmelCase  , output_scale_factor=1  , resnet_time_scale_shift="""default"""  , attention_head_dim=block_out_channels[-1]  , resnet_groups=__lowerCAmelCase  , temb_channels=__lowerCAmelCase  , )

														# out
														UpperCamelCase__				     =			nn.GroupNorm(num_channels=block_out_channels[-1]  , num_groups=__lowerCAmelCase  , eps=1E-6  )
														UpperCamelCase__				     =			nn.SiLU()

														UpperCamelCase__				     =			2 * out_channels if double_z else out_channels
														UpperCamelCase__				     =			nn.Convad(block_out_channels[-1]  , __lowerCAmelCase  , 3  , padding=1  )

														UpperCamelCase__				     =			False







							def 				_lowerCamelCase		(						self  , __lowerCAmelCase  ):
														UpperCamelCase__				     =			x
														UpperCamelCase__				     =			self.conv_in(__lowerCAmelCase  )

														if self.training and self.gradient_checkpointing:

																					def create_custom_forward(__lowerCAmelCase  ):
																												def custom_forward(*__lowerCAmelCase  ):
																																			return module(*__lowerCAmelCase  )

																												return custom_forward

																					# down
																					if is_torch_version(""">="""  , """1.11.0"""  ):
																												for down_block in self.down_blocks:
																																			UpperCamelCase__				     =			torch.utils.checkpoint.checkpoint(
																																			    create_custom_forward(__lowerCAmelCase  )  , __lowerCAmelCase  , use_reentrant=__lowerCAmelCase  )
																												# middle
																												UpperCamelCase__				     =			torch.utils.checkpoint.checkpoint(
																												    create_custom_forward(self.mid_block  )  , __lowerCAmelCase  , use_reentrant=__lowerCAmelCase  )
																					else:
																												for down_block in self.down_blocks:
																																			UpperCamelCase__				     =			torch.utils.checkpoint.checkpoint(create_custom_forward(__lowerCAmelCase  )  , __lowerCAmelCase  )
																												# middle
																												UpperCamelCase__				     =			torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block  )  , __lowerCAmelCase  )

														else:
																					# down
																					for down_block in self.down_blocks:
																												UpperCamelCase__				     =			down_block(__lowerCAmelCase  )

																					# middle
																					UpperCamelCase__				     =			self.mid_block(__lowerCAmelCase  )

														# post-process
														UpperCamelCase__				     =			self.conv_norm_out(__lowerCAmelCase  )
														UpperCamelCase__				     =			self.conv_act(__lowerCAmelCase  )
														UpperCamelCase__				     =			self.conv_out(__lowerCAmelCase  )

														return sample






class        __SCREAMING_SNAKE_CASE					(     nn.Module       ):






							def __init__(						self  , __lowerCAmelCase=3  , __lowerCAmelCase=3  , __lowerCAmelCase=("UpDecoderBlock2D",)  , __lowerCAmelCase=(64,)  , __lowerCAmelCase=2  , __lowerCAmelCase=32  , __lowerCAmelCase="silu"  , __lowerCAmelCase="group"  , ):
														super().__init__()
														UpperCamelCase__				     =			layers_per_block

														UpperCamelCase__				     =			nn.Convad(
														    __lowerCAmelCase  , block_out_channels[-1]  , kernel_size=3  , stride=1  , padding=1  , )

														UpperCamelCase__				     =			None
														UpperCamelCase__				     =			nn.ModuleList([]  )

														UpperCamelCase__				     =			in_channels if norm_type == """spatial""" else None

														# mid
														UpperCamelCase__				     =			UNetMidBlockaD(
														    in_channels=block_out_channels[-1]  , resnet_eps=1E-6  , resnet_act_fn=__lowerCAmelCase  , output_scale_factor=1  , resnet_time_scale_shift="""default""" if norm_type == """group""" else norm_type  , attention_head_dim=block_out_channels[-1]  , resnet_groups=__lowerCAmelCase  , temb_channels=__lowerCAmelCase  , )

														# up
														UpperCamelCase__				     =			list(reversed(__lowerCAmelCase  )  )
														UpperCamelCase__				     =			reversed_block_out_channels[0]
														for i, up_block_type in enumerate(__lowerCAmelCase  ):
																					UpperCamelCase__				     =			output_channel
																					UpperCamelCase__				     =			reversed_block_out_channels[i]

																					UpperCamelCase__				     =			i == len(__lowerCAmelCase  ) - 1

																					UpperCamelCase__				     =			get_up_block(
																					    __lowerCAmelCase  , num_layers=self.layers_per_block + 1  , in_channels=__lowerCAmelCase  , out_channels=__lowerCAmelCase  , prev_output_channel=__lowerCAmelCase  , add_upsample=not is_final_block  , resnet_eps=1E-6  , resnet_act_fn=__lowerCAmelCase  , resnet_groups=__lowerCAmelCase  , attention_head_dim=__lowerCAmelCase  , temb_channels=__lowerCAmelCase  , resnet_time_scale_shift=__lowerCAmelCase  , )
																					self.up_blocks.append(__lowerCAmelCase  )
																					UpperCamelCase__				     =			output_channel

														# out
														if norm_type == "spatial":
																					UpperCamelCase__				     =			SpatialNorm(block_out_channels[0]  , __lowerCAmelCase  )
														else:
																					UpperCamelCase__				     =			nn.GroupNorm(num_channels=block_out_channels[0]  , num_groups=__lowerCAmelCase  , eps=1E-6  )
														UpperCamelCase__				     =			nn.SiLU()
														UpperCamelCase__				     =			nn.Convad(block_out_channels[0]  , __lowerCAmelCase  , 3  , padding=1  )

														UpperCamelCase__				     =			False







							def 				_lowerCamelCase		(						self  , __lowerCAmelCase  , __lowerCAmelCase=None  ):
														UpperCamelCase__				     =			z
														UpperCamelCase__				     =			self.conv_in(__lowerCAmelCase  )

														UpperCamelCase__				     =			next(iter(self.up_blocks.parameters()  )  ).dtype
														if self.training and self.gradient_checkpointing:

																					def create_custom_forward(__lowerCAmelCase  ):
																												def custom_forward(*__lowerCAmelCase  ):
																																			return module(*__lowerCAmelCase  )

																												return custom_forward

																					if is_torch_version(""">="""  , """1.11.0"""  ):
																												# middle
																												UpperCamelCase__				     =			torch.utils.checkpoint.checkpoint(
																												    create_custom_forward(self.mid_block  )  , __lowerCAmelCase  , __lowerCAmelCase  , use_reentrant=__lowerCAmelCase  )
																												UpperCamelCase__				     =			sample.to(__lowerCAmelCase  )

																												# up
																												for up_block in self.up_blocks:
																																			UpperCamelCase__				     =			torch.utils.checkpoint.checkpoint(
																																			    create_custom_forward(__lowerCAmelCase  )  , __lowerCAmelCase  , __lowerCAmelCase  , use_reentrant=__lowerCAmelCase  )
																					else:
																												# middle
																												UpperCamelCase__				     =			torch.utils.checkpoint.checkpoint(
																												    create_custom_forward(self.mid_block  )  , __lowerCAmelCase  , __lowerCAmelCase  )
																												UpperCamelCase__				     =			sample.to(__lowerCAmelCase  )

																												# up
																												for up_block in self.up_blocks:
																																			UpperCamelCase__				     =			torch.utils.checkpoint.checkpoint(create_custom_forward(__lowerCAmelCase  )  , __lowerCAmelCase  , __lowerCAmelCase  )
														else:
																					# middle
																					UpperCamelCase__				     =			self.mid_block(__lowerCAmelCase  , __lowerCAmelCase  )
																					UpperCamelCase__				     =			sample.to(__lowerCAmelCase  )

																					# up
																					for up_block in self.up_blocks:
																												UpperCamelCase__				     =			up_block(__lowerCAmelCase  , __lowerCAmelCase  )

        # post-process
														if latent_embeds is None:
																					UpperCamelCase__				     =			self.conv_norm_out(__lowerCAmelCase  )
														else:
																					UpperCamelCase__				     =			self.conv_norm_out(__lowerCAmelCase  , __lowerCAmelCase  )
														UpperCamelCase__				     =			self.conv_act(__lowerCAmelCase  )
														UpperCamelCase__				     =			self.conv_out(__lowerCAmelCase  )

														return sample






class        __SCREAMING_SNAKE_CASE					(     nn.Module       ):






							def __init__(						self  , __lowerCAmelCase  , __lowerCAmelCase  , __lowerCAmelCase  , __lowerCAmelCase=None  , __lowerCAmelCase="random"  , __lowerCAmelCase=False  , __lowerCAmelCase=True  ):
														super().__init__()
														UpperCamelCase__				     =			n_e
														UpperCamelCase__				     =			vq_embed_dim
														UpperCamelCase__				     =			beta
														UpperCamelCase__				     =			legacy

														UpperCamelCase__				     =			nn.Embedding(self.n_e  , self.vq_embed_dim  )
														self.embedding.weight.data.uniform_(-1.0 / self.n_e  , 1.0 / self.n_e  )

														UpperCamelCase__				     =			remap
														if self.remap is not None:
																					self.register_buffer("""used"""  , torch.tensor(np.load(self.remap  )  )  )
																					UpperCamelCase__				     =			self.used.shape[0]
																					UpperCamelCase__				     =			unknown_index  # "random" or "extra" or integer
																					if self.unknown_index == "extra":
																												UpperCamelCase__				     =			self.re_embed
																												UpperCamelCase__				     =			self.re_embed + 1
																					print(
																					    f"""Remapping {self.n_e} indices to {self.re_embed} indices. """
																					    f"""Using {self.unknown_index} for unknown indices."""  )
														else:
																					UpperCamelCase__				     =			n_e

														UpperCamelCase__				     =			sane_index_shape






							def 				_lowerCamelCase		(						self  , __lowerCAmelCase  ):
														UpperCamelCase__				     =			inds.shape
														assert len(__lowerCAmelCase  ) > 1
														UpperCamelCase__				     =			inds.reshape(ishape[0]  , -1  )
														UpperCamelCase__				     =			self.used.to(__lowerCAmelCase  )
														UpperCamelCase__				     =			(inds[:, :, None] == used[None, None, ...]).long()
														UpperCamelCase__				     =			match.argmax(-1  )
														UpperCamelCase__				     =			match.sum(2  ) < 1
														if self.unknown_index == "random":
																					UpperCamelCase__				     =			torch.randint(0  , self.re_embed  , size=new[unknown].shape  ).to(device=new.device  )
														else:
																					UpperCamelCase__				     =			self.unknown_index
														return new.reshape(__lowerCAmelCase  )






							def 				_lowerCamelCase		(						self  , __lowerCAmelCase  ):
														UpperCamelCase__				     =			inds.shape
														assert len(__lowerCAmelCase  ) > 1
														UpperCamelCase__				     =			inds.reshape(ishape[0]  , -1  )
														UpperCamelCase__				     =			self.used.to(__lowerCAmelCase  )
														if self.re_embed > self.used.shape[0]:  # extra token
																					UpperCamelCase__				     =			0  # simply set to zero
														UpperCamelCase__				     =			torch.gather(used[None, :][inds.shape[0] * [0], :]  , 1  , __lowerCAmelCase  )
														return back.reshape(__lowerCAmelCase  )






							def 				_lowerCamelCase		(						self  , __lowerCAmelCase  ):
														# reshape z -> (batch, height, width, channel) and flatten
														UpperCamelCase__				     =			z.permute(0  , 2  , 3  , 1  ).contiguous()
														UpperCamelCase__				     =			z.view(-1  , self.vq_embed_dim  )

														# distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
														UpperCamelCase__				     =			torch.argmin(torch.cdist(__lowerCAmelCase  , self.embedding.weight  )  , dim=1  )

														UpperCamelCase__				     =			self.embedding(__lowerCAmelCase  ).view(z.shape  )
														UpperCamelCase__				     =			None
														UpperCamelCase__				     =			None

														# compute loss for embedding
														if not self.legacy:
																					UpperCamelCase__				     =			self.beta * torch.mean((z_q.detach() - z) ** 2  ) + torch.mean((z_q - z.detach()) ** 2  )
														else:
																					UpperCamelCase__				     =			torch.mean((z_q.detach() - z) ** 2  ) + self.beta * torch.mean((z_q - z.detach()) ** 2  )

														# preserve gradients
														UpperCamelCase__				     =			z + (z_q - z).detach()

														# reshape back to match original input shape
														UpperCamelCase__				     =			z_q.permute(0  , 3  , 1  , 2  ).contiguous()

														if self.remap is not None:
																					UpperCamelCase__				     =			min_encoding_indices.reshape(z.shape[0]  , -1  )  # add batch axis
																					UpperCamelCase__				     =			self.remap_to_used(__lowerCAmelCase  )
																					UpperCamelCase__				     =			min_encoding_indices.reshape(-1  , 1  )  # flatten

														if self.sane_index_shape:
																					UpperCamelCase__				     =			min_encoding_indices.reshape(z_q.shape[0]  , z_q.shape[2]  , z_q.shape[3]  )

														return z_q, loss, (perplexity, min_encodings, min_encoding_indices)







							def 				_lowerCamelCase		(						self  , __lowerCAmelCase  , __lowerCAmelCase  ):
														# shape specifying (batch, height, width, channel)
														if self.remap is not None:
																					UpperCamelCase__				     =			indices.reshape(shape[0]  , -1  )  # add batch axis
																					UpperCamelCase__				     =			self.unmap_to_all(__lowerCAmelCase  )
																					UpperCamelCase__				     =			indices.reshape(-1  )  # flatten again

														# get quantized latent vectors
														UpperCamelCase__				     =			self.embedding(__lowerCAmelCase  )

														if shape is not None:
																					UpperCamelCase__				     =			z_q.view(__lowerCAmelCase  )
																					# reshape back to match original input shape
																					UpperCamelCase__				     =			z_q.permute(0  , 3  , 1  , 2  ).contiguous()

														return z_q






class        __SCREAMING_SNAKE_CASE					(     _a       ):






							def __init__(						self  , __lowerCAmelCase  , __lowerCAmelCase=False  ):
														UpperCamelCase__				     =			parameters
														UpperCamelCase__     ,     UpperCamelCase__				     =			torch.chunk(__lowerCAmelCase  , 2  , dim=1  )
														UpperCamelCase__				     =			torch.clamp(self.logvar  , -30.0  , 20.0  )
														UpperCamelCase__				     =			deterministic
														UpperCamelCase__				     =			torch.exp(0.5 * self.logvar  )
														UpperCamelCase__				     =			torch.exp(self.logvar  )
														if self.deterministic:
																					UpperCamelCase__				     =			UpperCamelCase__				     =			torch.zeros_like(
																					    self.mean  , device=self.parameters.device  , dtype=self.parameters.dtype  )






							def 				_lowerCamelCase		(						self  , __lowerCAmelCase = None  ):
														# make sure sample is on the same device as the parameters and has same dtype
														UpperCamelCase__				     =			randn_tensor(
														    self.mean.shape  , generator=__lowerCAmelCase  , device=self.parameters.device  , dtype=self.parameters.dtype  )
														UpperCamelCase__				     =			self.mean + self.std * sample
														return x






							def 				_lowerCamelCase		(						self  , __lowerCAmelCase=None  ):
														if self.deterministic:
																					return torch.Tensor([0.0]  )
														else:
																					if other is None:
																												return 0.5 * torch.sum(torch.pow(self.mean  , 2  ) + self.var - 1.0 - self.logvar  , dim=[1, 2, 3]  )
																					else:
																												return 0.5 * torch.sum(
																												    torch.pow(self.mean - other.mean  , 2  ) / other.var
																												    + self.var / other.var
																												    - 1.0
																												    - self.logvar
																												    + other.logvar  , dim=[1, 2, 3]  , )






							def 				_lowerCamelCase		(						self  , __lowerCAmelCase  , __lowerCAmelCase=[1, 2, 3]  ):
														if self.deterministic:
																					return torch.Tensor([0.0]  )
														UpperCamelCase__				     =			np.log(2.0 * np.pi  )
														return 0.5 * torch.sum(logtwopi + self.logvar + torch.pow(sample - self.mean  , 2  ) / self.var  , dim=__lowerCAmelCase  )







							def 				_lowerCamelCase		(						self  ):
														return self.mean











from ...configuration_utils import PretrainedConfig
from ...utils import logging


UpperCamelCase__											= logging.get_logger(__name__)

UpperCamelCase__											= {
    "studio-ousia/luke-base": "https://huggingface.co/studio-ousia/luke-base/resolve/main/config.json",
    "studio-ousia/luke-large": "https://huggingface.co/studio-ousia/luke-large/resolve/main/config.json",
}

class        __SCREAMING_SNAKE_CASE					(     _a       ):
							snake_case     :      Optional[int]						     =       """luke"""






							def __init__(						self  , __lowerCAmelCase=50267  , __lowerCAmelCase=500000  , __lowerCAmelCase=768  , __lowerCAmelCase=256  , __lowerCAmelCase=12  , __lowerCAmelCase=12  , __lowerCAmelCase=3072  , __lowerCAmelCase="gelu"  , __lowerCAmelCase=0.1  , __lowerCAmelCase=0.1  , __lowerCAmelCase=512  , __lowerCAmelCase=2  , __lowerCAmelCase=0.02  , __lowerCAmelCase=1E-12  , __lowerCAmelCase=True  , __lowerCAmelCase=None  , __lowerCAmelCase=1  , __lowerCAmelCase=0  , __lowerCAmelCase=2  , **__lowerCAmelCase  , ):
														super().__init__(pad_token_id=__lowerCAmelCase  , bos_token_id=__lowerCAmelCase  , eos_token_id=__lowerCAmelCase  , **__lowerCAmelCase  )

														UpperCamelCase__				     =			vocab_size
														UpperCamelCase__				     =			entity_vocab_size
														UpperCamelCase__				     =			hidden_size
														UpperCamelCase__				     =			entity_emb_size
														UpperCamelCase__				     =			num_hidden_layers
														UpperCamelCase__				     =			num_attention_heads
														UpperCamelCase__				     =			hidden_act
														UpperCamelCase__				     =			intermediate_size
														UpperCamelCase__				     =			hidden_dropout_prob
														UpperCamelCase__				     =			attention_probs_dropout_prob
														UpperCamelCase__				     =			max_position_embeddings
														UpperCamelCase__				     =			type_vocab_size
														UpperCamelCase__				     =			initializer_range
														UpperCamelCase__				     =			layer_norm_eps
														UpperCamelCase__				     =			use_entity_aware_attention
														UpperCamelCase__				     =			classifier_dropout




