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        UpperCamelCase						(				lowerCAmelCase__				):
						a__		:List[str]      =			42


class        UpperCamelCase						(				nn.Module				):





						def __init__(self			,     __UpperCamelCase=3			,     __UpperCamelCase=3			,     __UpperCamelCase=("DownEncoderBlock2D",)			,     __UpperCamelCase=(64,)			,     __UpperCamelCase=2			,     __UpperCamelCase=32			,     __UpperCamelCase="silu"			,     __UpperCamelCase=True			,     )				->				Optional[Any]:
												super().__init__()
												UpperCamelCase_		:							Any  =		layers_per_block

												UpperCamelCase_		:							str  =		torch.nn.Convad(
												    __UpperCAmelCase			,     block_out_channels[0]			,     kernel_size=3			,     stride=1			,     padding=1			,     )

												UpperCamelCase_		:							Optional[int]  =		None
												UpperCamelCase_		:							Optional[int]  =		nn.ModuleList([]						)

												# down
												UpperCamelCase_		:							List[Any]  =		block_out_channels[0]
												for i, down_block_type in enumerate(__UpperCAmelCase						):
																		UpperCamelCase_		:							Optional[int]  =		output_channel
																		UpperCamelCase_		:							Tuple  =		block_out_channels[i]
																		UpperCamelCase_		:							Union[str, Any]  =		i == len(__UpperCAmelCase						) - 1

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

												# mid
												UpperCamelCase_		:							List[str]  =		UNetMidBlockaD(
												    in_channels=block_out_channels[-1]			,     resnet_eps=1E-6			,     resnet_act_fn=__UpperCAmelCase			,     output_scale_factor=1			,     resnet_time_scale_shift="""default"""			,     attention_head_dim=block_out_channels[-1]			,     resnet_groups=__UpperCAmelCase			,     temb_channels=__UpperCAmelCase			,     )

												# out
												UpperCamelCase_		:							Union[str, Any]  =		nn.GroupNorm(num_channels=block_out_channels[-1]			,     num_groups=__UpperCAmelCase			,     eps=1E-6						)
												UpperCamelCase_		:							List[Any]  =		nn.SiLU()

												UpperCamelCase_		:							Optional[Any]  =		2 * out_channels if double_z else out_channels
												UpperCamelCase_		:							Optional[Any]  =		nn.Convad(block_out_channels[-1]			,     __UpperCAmelCase			,     3			,     padding=1						)

												UpperCamelCase_		:							int  =		False





						def        A_							(self			,     __UpperCamelCase						)				->				Tuple:
												UpperCamelCase_		:							str  =		x
												UpperCamelCase_		:							int  =		self.conv_in(__UpperCAmelCase						)

												if self.training and self.gradient_checkpointing:

																		def create_custom_forward(__UpperCamelCase						):
																								def custom_forward(*__UpperCamelCase						):
																														return module(*__UpperCAmelCase						)

																								return custom_forward

																		# down
																		if is_torch_version(""">="""			,     """1.11.0"""						):
																								for down_block in self.down_blocks:
																														UpperCamelCase_		:							List[str]  =		torch.utils.checkpoint.checkpoint(
																														    create_custom_forward(__UpperCAmelCase						)			,     __UpperCAmelCase			,     use_reentrant=__UpperCAmelCase						)
																								# middle
																								UpperCamelCase_		:							int  =		torch.utils.checkpoint.checkpoint(
																								    create_custom_forward(self.mid_block						)			,     __UpperCAmelCase			,     use_reentrant=__UpperCAmelCase						)
																		else:
																								for down_block in self.down_blocks:
																														UpperCamelCase_		:							int  =		torch.utils.checkpoint.checkpoint(create_custom_forward(__UpperCAmelCase						)			,     __UpperCAmelCase						)
																								# middle
																								UpperCamelCase_		:							Any  =		torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block						)			,     __UpperCAmelCase						)

												else:
																		# down
																		for down_block in self.down_blocks:
																								UpperCamelCase_		:							int  =		down_block(__UpperCAmelCase						)

																		# middle
																		UpperCamelCase_		:							List[str]  =		self.mid_block(__UpperCAmelCase						)

												# post-process
												UpperCamelCase_		:							Union[str, Any]  =		self.conv_norm_out(__UpperCAmelCase						)
												UpperCamelCase_		:							List[str]  =		self.conv_act(__UpperCAmelCase						)
												UpperCamelCase_		:							Any  =		self.conv_out(__UpperCAmelCase						)

												return sample


class        UpperCamelCase						(				nn.Module				):





						def __init__(self			,     __UpperCamelCase=3			,     __UpperCamelCase=3			,     __UpperCamelCase=("UpDecoderBlock2D",)			,     __UpperCamelCase=(64,)			,     __UpperCamelCase=2			,     __UpperCamelCase=32			,     __UpperCamelCase="silu"			,     __UpperCamelCase="group"			,     )				->				Any:
												super().__init__()
												UpperCamelCase_		:							str  =		layers_per_block

												UpperCamelCase_		:							int  =		nn.Convad(
												    __UpperCAmelCase			,     block_out_channels[-1]			,     kernel_size=3			,     stride=1			,     padding=1			,     )

												UpperCamelCase_		:							Tuple  =		None
												UpperCamelCase_		:							Union[str, Any]  =		nn.ModuleList([]						)

												UpperCamelCase_		:							List[str]  =		in_channels if norm_type == """spatial""" else None

												# mid
												UpperCamelCase_		:							Tuple  =		UNetMidBlockaD(
												    in_channels=block_out_channels[-1]			,     resnet_eps=1E-6			,     resnet_act_fn=__UpperCAmelCase			,     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=__UpperCAmelCase			,     temb_channels=__UpperCAmelCase			,     )

												# up
												UpperCamelCase_		:							Any  =		list(reversed(__UpperCAmelCase						)						)
												UpperCamelCase_		:							Union[str, Any]  =		reversed_block_out_channels[0]
												for i, up_block_type in enumerate(__UpperCAmelCase						):
																		UpperCamelCase_		:							str  =		output_channel
																		UpperCamelCase_		:							Union[str, Any]  =		reversed_block_out_channels[i]

																		UpperCamelCase_		:							Union[str, Any]  =		i == len(__UpperCAmelCase						) - 1

																		UpperCamelCase_		:							Any  =		get_up_block(
																		    __UpperCAmelCase			,     num_layers=self.layers_per_block + 1			,     in_channels=__UpperCAmelCase			,     out_channels=__UpperCAmelCase			,     prev_output_channel=__UpperCAmelCase			,     add_upsample=not is_final_block			,     resnet_eps=1E-6			,     resnet_act_fn=__UpperCAmelCase			,     resnet_groups=__UpperCAmelCase			,     attention_head_dim=__UpperCAmelCase			,     temb_channels=__UpperCAmelCase			,     resnet_time_scale_shift=__UpperCAmelCase			,     )
																		self.up_blocks.append(__UpperCAmelCase						)
																		UpperCamelCase_		:							Optional[Any]  =		output_channel

												# out
												if norm_type == "spatial":
																		UpperCamelCase_		:							Optional[int]  =		SpatialNorm(block_out_channels[0]			,     __UpperCAmelCase						)
												else:
																		UpperCamelCase_		:							Any  =		nn.GroupNorm(num_channels=block_out_channels[0]			,     num_groups=__UpperCAmelCase			,     eps=1E-6						)
												UpperCamelCase_		:							str  =		nn.SiLU()
												UpperCamelCase_		:							str  =		nn.Convad(block_out_channels[0]			,     __UpperCAmelCase			,     3			,     padding=1						)

												UpperCamelCase_		:							List[Any]  =		False





						def        A_							(self			,     __UpperCamelCase			,     __UpperCamelCase=None						)				->				str:
												UpperCamelCase_		:							Any  =		z
												UpperCamelCase_		:							str  =		self.conv_in(__UpperCAmelCase						)

												UpperCamelCase_		:							Union[str, Any]  =		next(iter(self.up_blocks.parameters()						)						).dtype
												if self.training and self.gradient_checkpointing:

																		def create_custom_forward(__UpperCamelCase						):
																								def custom_forward(*__UpperCamelCase						):
																														return module(*__UpperCAmelCase						)

																								return custom_forward

																		if is_torch_version(""">="""			,     """1.11.0"""						):
																								# middle
																								UpperCamelCase_		:							List[str]  =		torch.utils.checkpoint.checkpoint(
																								    create_custom_forward(self.mid_block						)			,     __UpperCAmelCase			,     __UpperCAmelCase			,     use_reentrant=__UpperCAmelCase						)
																								UpperCamelCase_		:							List[str]  =		sample.to(__UpperCAmelCase						)

																								# up
																								for up_block in self.up_blocks:
																														UpperCamelCase_		:							Optional[Any]  =		torch.utils.checkpoint.checkpoint(
																														    create_custom_forward(__UpperCAmelCase						)			,     __UpperCAmelCase			,     __UpperCAmelCase			,     use_reentrant=__UpperCAmelCase						)
																		else:
																								# middle
																								UpperCamelCase_		:							str  =		torch.utils.checkpoint.checkpoint(
																								    create_custom_forward(self.mid_block						)			,     __UpperCAmelCase			,     __UpperCAmelCase						)
																								UpperCamelCase_		:							Tuple  =		sample.to(__UpperCAmelCase						)

																								# up
																								for up_block in self.up_blocks:
																														UpperCamelCase_		:							int  =		torch.utils.checkpoint.checkpoint(create_custom_forward(__UpperCAmelCase						)			,     __UpperCAmelCase			,     __UpperCAmelCase						)
												else:
																		# middle
																		UpperCamelCase_		:							Optional[Any]  =		self.mid_block(__UpperCAmelCase			,     __UpperCAmelCase						)
																		UpperCamelCase_		:							int  =		sample.to(__UpperCAmelCase						)

																		# up
																		for up_block in self.up_blocks:
																								UpperCamelCase_		:							Optional[int]  =		up_block(__UpperCAmelCase			,     __UpperCAmelCase						)

																		 # post-process
												if latent_embeds is None:
																		UpperCamelCase_		:							List[Any]  =		self.conv_norm_out(__UpperCAmelCase						)
												else:
																		UpperCamelCase_		:							str  =		self.conv_norm_out(__UpperCAmelCase			,     __UpperCAmelCase						)
												UpperCamelCase_		:							List[Any]  =		self.conv_act(__UpperCAmelCase						)
												UpperCamelCase_		:							Optional[Any]  =		self.conv_out(__UpperCAmelCase						)

												return sample


class        UpperCamelCase						(				nn.Module				):





						def __init__(self			,     __UpperCamelCase			,     __UpperCamelCase			,     __UpperCamelCase			,     __UpperCamelCase=None			,     __UpperCamelCase="random"			,     __UpperCamelCase=False			,     __UpperCamelCase=True						)				->				int:
												super().__init__()
												UpperCamelCase_		:							List[Any]  =		n_e
												UpperCamelCase_		:							Dict  =		vq_embed_dim
												UpperCamelCase_		:							Dict  =		beta
												UpperCamelCase_		:							Optional[int]  =		legacy

												UpperCamelCase_		:							Optional[int]  =		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_		:							Optional[int]  =		remap
												if self.remap is not None:
																		self.register_buffer("""used"""			,     torch.tensor(np.load(self.remap						)						)						)
																		UpperCamelCase_		:							str  =		self.used.shape[0]
																		UpperCamelCase_		:							int  =		unknown_index  # "random" or "extra" or integer
																		if self.unknown_index == "extra":
																								UpperCamelCase_		:							Optional[Any]  =		self.re_embed
																								UpperCamelCase_		:							List[str]  =		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_		:							Dict  =		n_e

												UpperCamelCase_		:							List[Any]  =		sane_index_shape





						def        A_							(self			,     __UpperCamelCase						)				->				Optional[Any]:
												UpperCamelCase_		:							Dict  =		inds.shape
												assert len(__UpperCAmelCase						) > 1
												UpperCamelCase_		:							List[str]  =		inds.reshape(ishape[0]			,     -1						)
												UpperCamelCase_		:							List[str]  =		self.used.to(__UpperCAmelCase						)
												UpperCamelCase_		:							str  =		(inds[:, :, None] == used[None, None, ...]).long()
												UpperCamelCase_		:							Union[str, Any]  =		match.argmax(-1						)
												UpperCamelCase_		:							Any  =		match.sum(2						) < 1
												if self.unknown_index == "random":
																		UpperCamelCase_		:							Any  =		torch.randint(0			,     self.re_embed			,     size=new[unknown].shape						).to(device=new.device						)
												else:
																		UpperCamelCase_		:							Dict  =		self.unknown_index
												return new.reshape(__UpperCAmelCase						)





						def        A_							(self			,     __UpperCamelCase						)				->				List[Any]:
												UpperCamelCase_		:							str  =		inds.shape
												assert len(__UpperCAmelCase						) > 1
												UpperCamelCase_		:							Any  =		inds.reshape(ishape[0]			,     -1						)
												UpperCamelCase_		:							List[str]  =		self.used.to(__UpperCAmelCase						)
												if self.re_embed > self.used.shape[0]:  # extra token
																		UpperCamelCase_		:							Any  =		0  # simply set to zero
												UpperCamelCase_		:							List[Any]  =		torch.gather(used[None, :][inds.shape[0] * [0], :]			,     1			,     __UpperCAmelCase						)
												return back.reshape(__UpperCAmelCase						)





						def        A_							(self			,     __UpperCamelCase						)				->				Tuple:
												UpperCamelCase_		:							Tuple  =		z.permute(0			,     2			,     3			,     1						).contiguous()
												UpperCamelCase_		:							Union[str, Any]  =		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_		:							str  =		torch.argmin(torch.cdist(__UpperCAmelCase			,     self.embedding.weight						)			,     dim=1						)

												UpperCamelCase_		:							Optional[Any]  =		self.embedding(__UpperCAmelCase						).view(z.shape						)
												UpperCamelCase_		:							Optional[Any]  =		None
												UpperCamelCase_		:							str  =		None

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

												# preserve gradients
												UpperCamelCase_		:							List[str]  =		z + (z_q - z).detach()

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

												if self.remap is not None:
																		UpperCamelCase_		:							Optional[Any]  =		min_encoding_indices.reshape(z.shape[0]			,     -1						)  # add batch axis
																		UpperCamelCase_		:							int  =		self.remap_to_used(__UpperCAmelCase						)
																		UpperCamelCase_		:							List[Any]  =		min_encoding_indices.reshape(-1			,     1						)  # flatten

												if self.sane_index_shape:
																		UpperCamelCase_		:							List[str]  =		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        A_							(self			,     __UpperCamelCase			,     __UpperCamelCase						)				->				Optional[Any]:
												# shape specifying (batch, height, width, channel)
												if self.remap is not None:
																		UpperCamelCase_		:							Tuple  =		indices.reshape(shape[0]			,     -1						)  # add batch axis
																		UpperCamelCase_		:							Optional[Any]  =		self.unmap_to_all(__UpperCAmelCase						)
																		UpperCamelCase_		:							int  =		indices.reshape(-1						)  # flatten again

												# get quantized latent vectors
												UpperCamelCase_		:							int  =		self.embedding(__UpperCAmelCase						)

												if shape is not None:
																		UpperCamelCase_		:							Dict  =		z_q.view(__UpperCAmelCase						)
																		# reshape back to match original input shape
																		UpperCamelCase_		:							Optional[int]  =		z_q.permute(0			,     3			,     1			,     2						).contiguous()

												return z_q


class        UpperCamelCase						(				lowerCAmelCase__				):





						def __init__(self			,     __UpperCamelCase			,     __UpperCamelCase=False						)				->				List[str]:
												UpperCamelCase_		:							List[str]  =		parameters
												UpperCamelCase_,UpperCamelCase_		:							str  =		torch.chunk(__UpperCAmelCase			,     2			,     dim=1						)
												UpperCamelCase_		:							Tuple  =		torch.clamp(self.logvar			,     -30.0			,     20.0						)
												UpperCamelCase_		:							Union[str, Any]  =		deterministic
												UpperCamelCase_		:							str  =		torch.exp(0.5 * self.logvar						)
												UpperCamelCase_		:							int  =		torch.exp(self.logvar						)
												if self.deterministic:
																		UpperCamelCase_		:							Dict  =		torch.zeros_like(
																		    self.mean			,     device=self.parameters.device			,     dtype=self.parameters.dtype						)





						def        A_							(self			,     __UpperCamelCase = None						)				->				List[str]:
												UpperCamelCase_		:							List[Any]  =		randn_tensor(
												    self.mean.shape			,     generator=__UpperCAmelCase			,     device=self.parameters.device			,     dtype=self.parameters.dtype						)
												UpperCamelCase_		:							Optional[int]  =		self.mean + self.std * sample
												return x





						def        A_							(self			,     __UpperCamelCase=None						)				->				List[str]:
												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        A_							(self			,     __UpperCamelCase			,     __UpperCamelCase=[1, 2, 3]						)				->				str:
												if self.deterministic:
																		return torch.Tensor([0.0]						)
												UpperCamelCase_		:							Optional[Any]  =		np.log(2.0 * np.pi						)
												return 0.5 * torch.sum(logtwopi + self.logvar + torch.pow(sample - self.mean			,     2						) / self.var			,     dim=__UpperCAmelCase						)





						def        A_							(self						)				->				Union[str, Any]:
												return self.mean











import json
import os
import unittest

from transformers.models.xlm.tokenization_xlm import VOCAB_FILES_NAMES, XLMTokenizer
from transformers.testing_utils import slow

from ...test_tokenization_common import TokenizerTesterMixin



class 			__lowerCAmelCase  (						lowerCAmelCase__	,						unittest.TestCase ):
 lowerCAmelCase__							       =	XLMTokenizer
 lowerCAmelCase__							       =	False




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    super().setUp()

    # Adapted from Sennrich et al. 2015 and https://github.com/rsennrich/subword-nmt
    __lowerCamelCase    				=     [
        '''l''',
        '''o''',
        '''w''',
        '''e''',
        '''r''',
        '''s''',
        '''t''',
        '''i''',
        '''d''',
        '''n''',
        '''w</w>''',
        '''r</w>''',
        '''t</w>''',
        '''lo''',
        '''low''',
        '''er</w>''',
        '''low</w>''',
        '''lowest</w>''',
        '''newer</w>''',
        '''wider</w>''',
        '''<unk>''',
    ]
    __lowerCamelCase    				=     dict(zip(__UpperCAmelCase						,					range(len(__UpperCAmelCase					)					)					)					)
    __lowerCamelCase    				=     ['''l o 123''', '''lo w 1456''', '''e r</w> 1789''', '''''']

    __lowerCamelCase    				=     os.path.join(self.tmpdirname						,					VOCAB_FILES_NAMES['''vocab_file''']					)
    __lowerCamelCase    				=     os.path.join(self.tmpdirname						,					VOCAB_FILES_NAMES['''merges_file''']					)
    with open(self.vocab_file						,					'''w'''					) as fp:
       fp.write(json.dumps(__UpperCAmelCase					)					)
    with open(self.merges_file						,					'''w'''					) as fp:
       fp.write('''\n'''.join(__UpperCAmelCase					)					)




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase					):



    '''simple docstring'''
    __lowerCamelCase    				=     '''lower newer'''
    __lowerCamelCase    				=     '''lower newer'''
    return input_text, output_text




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     XLMTokenizer(self.vocab_file						,					self.merges_file					)

    __lowerCamelCase    				=     '''lower'''
    __lowerCamelCase    				=     ['''low''', '''er</w>''']
    __lowerCamelCase    				=     tokenizer.tokenize(__UpperCAmelCase					)
    self.assertListEqual(__UpperCAmelCase						,					__UpperCAmelCase					)

    __lowerCamelCase    				=     tokens + ['''<unk>''']
    __lowerCamelCase    				=     [14, 15, 20]
    self.assertListEqual(tokenizer.convert_tokens_to_ids(__UpperCAmelCase					)						,					__UpperCAmelCase					)




 @slow
 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     XLMTokenizer.from_pretrained('''xlm-mlm-en-2048'''					)

    __lowerCamelCase    				=     tokenizer.encode('''sequence builders'''						,					add_special_tokens=__UpperCAmelCase					)
    __lowerCamelCase    				=     tokenizer.encode('''multi-sequence build'''						,					add_special_tokens=__UpperCAmelCase					)

    __lowerCamelCase    				=     tokenizer.build_inputs_with_special_tokens(__UpperCAmelCase					)
    __lowerCamelCase    				=     tokenizer.build_inputs_with_special_tokens(__UpperCAmelCase						,					__UpperCAmelCase					)

    assert encoded_sentence == [0] + text + [1]
    assert encoded_pair == [0] + text + [1] + text_a + [1]






