





'''simple docstring'''



from dataclasses import dataclass
from typing import Optional, Tuple, Union

import torch
import torch.nn as nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .embeddings import GaussianFourierProjection, TimestepEmbedding, Timesteps
from .modeling_utils import ModelMixin
from .unet_ad_blocks import get_down_block, get_mid_block, get_out_block, get_up_block






@dataclass
class 							__lowercase    ( lowerCAmelCase__  ):







	'''simple docstring'''

	a      :			torch.FloatTensor






class 							__lowercase    ( lowerCAmelCase__	, lowerCAmelCase__  ):







	'''simple docstring'''

	@register_to_config
	def __init__(self       ,_lowerCamelCase = 65536       ,_lowerCamelCase = None       ,_lowerCamelCase = 2       ,_lowerCamelCase = 2       ,_lowerCamelCase = 0       ,_lowerCamelCase = "fourier"       ,_lowerCamelCase = True       ,_lowerCamelCase = False       ,_lowerCamelCase = 0.0       ,_lowerCamelCase = ("DownBlock1DNoSkip", "DownBlock1D", "AttnDownBlock1D")       ,_lowerCamelCase = ("AttnUpBlock1D", "UpBlock1D", "UpBlock1DNoSkip")       ,_lowerCamelCase = "UNetMidBlock1D"       ,_lowerCamelCase = None       ,_lowerCamelCase = (32, 32, 64)       ,_lowerCamelCase = None       ,_lowerCamelCase = 8       ,_lowerCamelCase = 1       ,_lowerCamelCase = False       ,)							->							Union[str, Any]:







				'''simple docstring'''





				super().__init__()
				__lowercase			  =							sample_size

				# time
				if time_embedding_type == "fourier":
							__lowercase			  =							GaussianFourierProjection(
							    embedding_size=8       ,set_W_to_weight=_lowerCamelCase       ,log=_lowerCamelCase       ,flip_sin_to_cos=_lowerCamelCase	)
							__lowercase			  =							2 * block_out_channels[0]
				elif time_embedding_type == "positional":
							__lowercase			  =							Timesteps(
							    block_out_channels[0]       ,flip_sin_to_cos=_lowerCamelCase       ,downscale_freq_shift=_lowerCamelCase	)
							__lowercase			  =							block_out_channels[0]

				if use_timestep_embedding:
							__lowercase			  =							block_out_channels[0] * 4
							__lowercase			  =							TimestepEmbedding(
							    in_channels=_lowerCamelCase       ,time_embed_dim=_lowerCamelCase       ,act_fn=_lowerCamelCase       ,out_dim=block_out_channels[0]       ,)

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

				# down
				__lowercase			  =							in_channels
				for i, down_block_type in enumerate(_lowerCamelCase	):
							__lowercase			  =							output_channel
							__lowercase			  =							block_out_channels[i]

							if i == 0:
										input_channel += extra_in_channels

							__lowercase			  =							i == len(_lowerCamelCase	) - 1

							__lowercase			  =							get_down_block(
							    _lowerCamelCase       ,num_layers=_lowerCamelCase       ,in_channels=_lowerCamelCase       ,out_channels=_lowerCamelCase       ,temb_channels=block_out_channels[0]       ,add_downsample=not is_final_block or downsample_each_block       ,)
							self.down_blocks.append(_lowerCamelCase	)

				# mid
				__lowercase			  =							get_mid_block(
				    _lowerCamelCase       ,in_channels=block_out_channels[-1]       ,mid_channels=block_out_channels[-1]       ,out_channels=block_out_channels[-1]       ,embed_dim=block_out_channels[0]       ,num_layers=_lowerCamelCase       ,add_downsample=_lowerCamelCase       ,)

				# up
				__lowercase			  =							list(reversed(_lowerCamelCase	)	)
				__lowercase			  =							reversed_block_out_channels[0]
				if out_block_type is None:
							__lowercase			  =							out_channels
				else:
							__lowercase			  =							block_out_channels[0]

				for i, up_block_type in enumerate(_lowerCamelCase	):
							__lowercase			  =							output_channel
							__lowercase			  =							(
							    reversed_block_out_channels[i + 1] if i < len(_lowerCamelCase	) - 1 else final_upsample_channels
							)

							__lowercase			  =							i == len(_lowerCamelCase	) - 1

							__lowercase			  =							get_up_block(
							    _lowerCamelCase       ,num_layers=_lowerCamelCase       ,in_channels=_lowerCamelCase       ,out_channels=_lowerCamelCase       ,temb_channels=block_out_channels[0]       ,add_upsample=not is_final_block       ,)
							self.up_blocks.append(_lowerCamelCase	)
							__lowercase			  =							output_channel

				# out
				__lowercase			  =							norm_num_groups if norm_num_groups is not None else min(block_out_channels[0] // 4       ,32	)
				__lowercase			  =							get_out_block(
				    out_block_type=_lowerCamelCase       ,num_groups_out=_lowerCamelCase       ,embed_dim=block_out_channels[0]       ,out_channels=_lowerCamelCase       ,act_fn=_lowerCamelCase       ,fc_dim=block_out_channels[-1] // 4       ,)


	def   _UpperCAmelCase       (self       ,_lowerCamelCase       ,_lowerCamelCase       ,_lowerCamelCase = True       ,)							->							Union[UNetaDOutput, Tuple]:







				'''simple docstring'''





				__lowercase			  =							timestep
				if not torch.is_tensor(_lowerCamelCase	):
							__lowercase			  =							torch.tensor([timesteps]       ,dtype=torch.long       ,device=sample.device	)
				elif torch.is_tensor(_lowerCamelCase	) and len(timesteps.shape	) == 0:
							__lowercase			  =							timesteps[None].to(sample.device	)

				__lowercase			  =							self.time_proj(_lowerCamelCase	)
				if self.config.use_timestep_embedding:
							__lowercase			  =							self.time_mlp(_lowerCamelCase	)
				else:
							__lowercase			  =							timestep_embed[..., None]
							__lowercase			  =							timestep_embed.repeat([1, 1, sample.shape[2]]	).to(sample.dtype	)
							__lowercase			  =							timestep_embed.broadcast_to((sample.shape[:1] + timestep_embed.shape[1:])	)

				# 2. down
				__lowercase			  =							()
				for downsample_block in self.down_blocks:
							__lowercase  ,						__lowercase			  =							downsample_block(hidden_states=_lowerCamelCase       ,temb=_lowerCamelCase	)
							down_block_res_samples += res_samples

				# 3. mid
				if self.mid_block:
							__lowercase			  =							self.mid_block(_lowerCamelCase       ,_lowerCamelCase	)

				# 4. up
				for i, upsample_block in enumerate(self.up_blocks	):
							__lowercase			  =							down_block_res_samples[-1:]
							__lowercase			  =							down_block_res_samples[:-1]
							__lowercase			  =							upsample_block(_lowerCamelCase       ,res_hidden_states_tuple=_lowerCamelCase       ,temb=_lowerCamelCase	)

				# 5. post-process
				if self.out_block:
							__lowercase			  =							self.out_block(_lowerCamelCase       ,_lowerCamelCase	)

				if not return_dict:
							return (sample,)

				return UNetaDOutput(sample=_lowerCamelCase	)












'''simple docstring'''



import json
from typing import List, Optional, Tuple

from tokenizers import normalizers

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from .tokenization_electra import ElectraTokenizer


_SCREAMING_SNAKE_CASE   =      {'''vocab_file''': '''vocab.txt''', '''tokenizer_file''': '''tokenizer.json'''}

_SCREAMING_SNAKE_CASE   =      {
    '''vocab_file''': {
        '''google/electra-small-generator''': (
            '''https://huggingface.co/google/electra-small-generator/resolve/main/vocab.txt'''
        ),
        '''google/electra-base-generator''': '''https://huggingface.co/google/electra-base-generator/resolve/main/vocab.txt''',
        '''google/electra-large-generator''': (
            '''https://huggingface.co/google/electra-large-generator/resolve/main/vocab.txt'''
        ),
        '''google/electra-small-discriminator''': (
            '''https://huggingface.co/google/electra-small-discriminator/resolve/main/vocab.txt'''
        ),
        '''google/electra-base-discriminator''': (
            '''https://huggingface.co/google/electra-base-discriminator/resolve/main/vocab.txt'''
        ),
        '''google/electra-large-discriminator''': (
            '''https://huggingface.co/google/electra-large-discriminator/resolve/main/vocab.txt'''
        ),
    },
    '''tokenizer_file''': {
        '''google/electra-small-generator''': (
            '''https://huggingface.co/google/electra-small-generator/resolve/main/tokenizer.json'''
        ),
        '''google/electra-base-generator''': (
            '''https://huggingface.co/google/electra-base-generator/resolve/main/tokenizer.json'''
        ),
        '''google/electra-large-generator''': (
            '''https://huggingface.co/google/electra-large-generator/resolve/main/tokenizer.json'''
        ),
        '''google/electra-small-discriminator''': (
            '''https://huggingface.co/google/electra-small-discriminator/resolve/main/tokenizer.json'''
        ),
        '''google/electra-base-discriminator''': (
            '''https://huggingface.co/google/electra-base-discriminator/resolve/main/tokenizer.json'''
        ),
        '''google/electra-large-discriminator''': (
            '''https://huggingface.co/google/electra-large-discriminator/resolve/main/tokenizer.json'''
        ),
    },
}

_SCREAMING_SNAKE_CASE   =      {
    '''google/electra-small-generator''': 5_1_2,
    '''google/electra-base-generator''': 5_1_2,
    '''google/electra-large-generator''': 5_1_2,
    '''google/electra-small-discriminator''': 5_1_2,
    '''google/electra-base-discriminator''': 5_1_2,
    '''google/electra-large-discriminator''': 5_1_2,
}

_SCREAMING_SNAKE_CASE   =      {
    '''google/electra-small-generator''': {'''do_lower_case''': True},
    '''google/electra-base-generator''': {'''do_lower_case''': True},
    '''google/electra-large-generator''': {'''do_lower_case''': True},
    '''google/electra-small-discriminator''': {'''do_lower_case''': True},
    '''google/electra-base-discriminator''': {'''do_lower_case''': True},
    '''google/electra-large-discriminator''': {'''do_lower_case''': True},
}






class 							__lowercase    ( lowerCAmelCase__  ):







	'''simple docstring'''

	a      :			str     =    VOCAB_FILES_NAMES
	a      :			Dict     =    PRETRAINED_VOCAB_FILES_MAP
	a      :			int     =    PRETRAINED_INIT_CONFIGURATION
	a      :			Tuple     =    PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
	a      :			List[Any]     =    ElectraTokenizer


	def __init__(self       ,_lowerCamelCase=None       ,_lowerCamelCase=None       ,_lowerCamelCase=True       ,_lowerCamelCase="[UNK]"       ,_lowerCamelCase="[SEP]"       ,_lowerCamelCase="[PAD]"       ,_lowerCamelCase="[CLS]"       ,_lowerCamelCase="[MASK]"       ,_lowerCamelCase=True       ,_lowerCamelCase=None       ,**_lowerCamelCase       ,)							->							List[Any]:







				'''simple docstring'''





				super().__init__(
				    _lowerCamelCase       ,tokenizer_file=_lowerCamelCase       ,do_lower_case=_lowerCamelCase       ,unk_token=_lowerCamelCase       ,sep_token=_lowerCamelCase       ,pad_token=_lowerCamelCase       ,cls_token=_lowerCamelCase       ,mask_token=_lowerCamelCase       ,tokenize_chinese_chars=_lowerCamelCase       ,strip_accents=_lowerCamelCase       ,**_lowerCamelCase       ,)

				__lowercase			  =							json.loads(self.backend_tokenizer.normalizer.__getstate__()	)
				if (
				    normalizer_state.get('''lowercase'''       ,_lowerCamelCase	) != do_lower_case
				    or normalizer_state.get('''strip_accents'''       ,_lowerCamelCase	) != strip_accents
				    or normalizer_state.get('''handle_chinese_chars'''       ,_lowerCamelCase	) != tokenize_chinese_chars
				):
							__lowercase			  =							getattr(_lowerCamelCase       ,normalizer_state.pop('''type'''	)	)
							__lowercase			  =							do_lower_case
							__lowercase			  =							strip_accents
							__lowercase			  =							tokenize_chinese_chars
							__lowercase			  =							normalizer_class(**_lowerCamelCase	)

				__lowercase			  =							do_lower_case


	def   _UpperCAmelCase       (self       ,_lowerCamelCase       ,_lowerCamelCase=None	)							->							str:







				'''simple docstring'''





				__lowercase			  =							[self.cls_token_id] + token_ids_a + [self.sep_token_id]

				if token_ids_a:
							output += token_ids_a + [self.sep_token_id]

				return output


	def   _UpperCAmelCase       (self       ,_lowerCamelCase       ,_lowerCamelCase = None	)							->							List[int]:







				'''simple docstring'''





				__lowercase			  =							[self.sep_token_id]
				__lowercase			  =							[self.cls_token_id]
				if token_ids_a is None:
							return len(cls + token_ids_a + sep	) * [0]
				return len(cls + token_ids_a + sep	) * [0] + len(token_ids_a + sep	) * [1]


	def   _UpperCAmelCase       (self       ,_lowerCamelCase       ,_lowerCamelCase = None	)							->							Tuple[str]:







				'''simple docstring'''





				__lowercase			  =							self._tokenizer.model.save(_lowerCamelCase       ,name=_lowerCamelCase	)
				return tuple(_lowerCamelCase	)





