





"""simple docstring"""


import re
import tempfile
import unittest
from pathlib import Path

from datasets.utils.metadata import DatasetMetadata

def     UpperCamelCase       (SCREAMING_SNAKE_CASE    ):
						UpperCamelCase							:		List[str]							      = min(re.search("""^ +"""       ,    SCREAMING_SNAKE_CASE    ).end() if t.startswith(""" """    ) else 0 for t in string.splitlines()    )
						return "\n".join([line[indent_level:] for line in string.splitlines() if indent_level < len(SCREAMING_SNAKE_CASE    )]    )


__magic_name__					:				Dict	= """\
---
language:
- zh
- en
task_ids:
- sentiment-classification
---
# Begin of markdown

Some cool dataset card
"""

__magic_name__					:				Any	= """\
---
---
# Begin of markdown

Some cool dataset card
"""


__magic_name__					:				str	= """\
# Begin of markdown

Some cool dataset card
"""






class        lowercase__ (      unittest.TestCase  ):

			"""simple docstring"""







			def      _a			(			self						):






									'''simple docstring'''







									with tempfile.TemporaryDirectory() as tmp_dir:
															UpperCamelCase							:		List[str]							      = Path(_A						) / """README.md"""
															with open(_A	,      """w+"""						) as readme_file:
																					readme_file.write(_A						)
															UpperCamelCase							:		Tuple							      = DatasetMetadata.from_readme(_A						)
															self.assertDictEqual(_A	,      {"""language""": ["""zh""", """en"""], """task_ids""": ["""sentiment-classification"""]}						)

															with open(_A	,      """w+"""						) as readme_file:
																					readme_file.write(_A						)
															UpperCamelCase							:		Dict							      = DatasetMetadata.from_readme(_A						)
															self.assertDictEqual(_A	,      {}						)

															with open(_A	,      """w+"""						) as readme_file:
																					readme_file.write(_A						)
															UpperCamelCase							:		Any							      = DatasetMetadata.from_readme(_A						)
															self.assertEqual(_A	,      {}						)







			def      _a			(			self						):






									'''simple docstring'''







									UpperCamelCase							:		List[str]							      = _dedent(
									    """\
            annotations_creators:
            - found
            language_creators:
            - found
            language:
            - en
            license:
            - unknown
            multilinguality:
            - monolingual
            pretty_name: Test Dataset
            size_categories:
            - 10K<n<100K
            source_datasets:
            - extended|other-yahoo-webscope-l6
            task_categories:
            - question-answering
            task_ids:
            - open-domain-qa
            """						)
									assert DatasetMetadata.from_yaml_string(_A						)

									UpperCamelCase							:		List[Any]							      = _dedent(
									    """\
            annotations_creators:
            - found
            language:
            - en
            license:
            - unknown
            multilinguality:
            - monolingual
            pretty_name: Test Dataset
            size_categories:
            - 10K<n<100K
            source_datasets:
            - extended|other-yahoo-webscope-l6
            task_categories:
            - question-answering
            task_ids:
            - open-domain-qa
            task_ids:
            - open-domain-qa
            """						)
									with self.assertRaises(_A						):
															DatasetMetadata.from_yaml_string(_A						)

									UpperCamelCase							:		List[str]							      = _dedent(
									    """\
            annotations_creators:
            - found
            language_creators:
            - found
            language:
            - en
            license:
            - unknown
            multilinguality:
            - monolingual
            pretty_name: Test Dataset
            size_categories:
            - 10K<n<100K
            source_datasets:
            - extended|other-yahoo-webscope-l6
            task_categories:
            - text-classification
            task_ids:
            - multi-class-classification
            paperswithcode_id:
            - squad
            configs:
            - en
            train-eval-index:
            - config: en
              task: text-classification
              task_id: multi_class_classification
              splits:
                train_split: train
                eval_split: test
              col_mapping:
                text: text
                label: target
              metrics:
                - type: accuracy
                  name: Accuracy
            extra_gated_prompt: |
              By clicking on “Access repository” below, you also agree to ImageNet Terms of Access:
              [RESEARCHER_FULLNAME] (the \"Researcher\") has requested permission to use the ImageNet database (the \"Database\") at Princeton University and Stanford University. In exchange for such permission, Researcher hereby agrees to the following terms and conditions:
              1. Researcher shall use the Database only for non-commercial research and educational purposes.
            extra_gated_fields:
              Company: text
              Country: text
              I agree to use this model for non-commerical use ONLY: checkbox
            """						)
									assert DatasetMetadata.from_yaml_string(_A						)










"""simple docstring"""


from dataclasses import dataclass
from typing import Any, Dict, Optional

import torch
import torch.nn.functional as F
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..models.embeddings import ImagePositionalEmbeddings
from ..utils import BaseOutput, deprecate
from .attention import BasicTransformerBlock
from .embeddings import PatchEmbed
from .modeling_utils import ModelMixin






@dataclass
class        lowercase__ (      __SCREAMING_SNAKE_CASE  ):

			"""simple docstring"""
			__lowerCAmelCase    :							torch.FloatTensor






class        lowercase__ (      __SCREAMING_SNAKE_CASE				,    __SCREAMING_SNAKE_CASE  ):

			"""simple docstring"""







			@register_to_config
			def __init__(			self	,      _A = 1_6	,      _A = 8_8	,      _A = None	,      _A = None	,      _A = 1	,      _A = 0.0	,      _A = 3_2	,      _A = None	,      _A = False	,      _A = None	,      _A = None	,      _A = None	,      _A = "geglu"	,      _A = None	,      _A = False	,      _A = False	,      _A = False	,      _A = "layer_norm"	,      _A = True	,      ):






									'''simple docstring'''







									super().__init__()
									UpperCamelCase							:		int							      = use_linear_projection
									UpperCamelCase							:		str							      = num_attention_heads
									UpperCamelCase							:		Union[str, Any]							      = attention_head_dim
									UpperCamelCase							:		List[Any]							      = num_attention_heads * attention_head_dim

									# 1. Transformer2DModel can process both standard continuous images of shape `(batch_size, num_channels, width, height)` as well as quantized image embeddings of shape `(batch_size, num_image_vectors)`
									# Define whether input is continuous or discrete depending on configuration
									UpperCamelCase							:		int							      = (in_channels is not None) and (patch_size is None)
									UpperCamelCase							:		str							      = num_vector_embeds is not None
									UpperCamelCase							:		Tuple							      = in_channels is not None and patch_size is not None

									if norm_type == "layer_norm" and num_embeds_ada_norm is not None:
															UpperCamelCase							:		Tuple							      = (
															    f"""The configuration file of this model: {self.__class__} is outdated. `norm_type` is either not set or"""
															    """ incorrectly set to `'layer_norm'`.Make sure to set `norm_type` to `'ada_norm'` in the config."""
															    """ Please make sure to update the config accordingly as leaving `norm_type` might led to incorrect"""
															    """ results in future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it"""
															    """ would be very nice if you could open a Pull request for the `transformer/config.json` file"""
															)
															deprecate("""norm_type!=num_embeds_ada_norm"""	,      """1.0.0"""	,      _A	,      standard_warn=_A						)
															UpperCamelCase							:		Optional[int]							      = """ada_norm"""

									if self.is_input_continuous and self.is_input_vectorized:
															raise ValueError(
															    f"""Cannot define both `in_channels`: {in_channels} and `num_vector_embeds`: {num_vector_embeds}. Make"""
															    """ sure that either `in_channels` or `num_vector_embeds` is None."""						)
									elif self.is_input_vectorized and self.is_input_patches:
															raise ValueError(
															    f"""Cannot define both `num_vector_embeds`: {num_vector_embeds} and `patch_size`: {patch_size}. Make"""
															    """ sure that either `num_vector_embeds` or `num_patches` is None."""						)
									elif not self.is_input_continuous and not self.is_input_vectorized and not self.is_input_patches:
															raise ValueError(
															    f"""Has to define `in_channels`: {in_channels}, `num_vector_embeds`: {num_vector_embeds}, or patch_size:"""
															    f""" {patch_size}. Make sure that `in_channels`, `num_vector_embeds` or `num_patches` is not None."""						)

									# 2. Define input layers
									if self.is_input_continuous:
															UpperCamelCase							:		Union[str, Any]							      = in_channels

															UpperCamelCase							:		List[str]							      = torch.nn.GroupNorm(num_groups=_A	,      num_channels=_A	,      eps=1e-6	,      affine=_A						)
															if use_linear_projection:
																					UpperCamelCase							:		Dict							      = nn.Linear(_A	,      _A						)
															else:
																					UpperCamelCase							:		int							      = nn.Convad(_A	,      _A	,      kernel_size=1	,      stride=1	,      padding=0						)
									elif self.is_input_vectorized:
															assert sample_size is not None, "Transformer2DModel over discrete input must provide sample_size"
															assert num_vector_embeds is not None, "Transformer2DModel over discrete input must provide num_embed"

															UpperCamelCase							:		Tuple							      = sample_size
															UpperCamelCase							:		str							      = sample_size
															UpperCamelCase							:		int							      = num_vector_embeds
															UpperCamelCase							:		Any							      = self.height * self.width

															UpperCamelCase							:		Any							      = ImagePositionalEmbeddings(
															    num_embed=_A	,      embed_dim=_A	,      height=self.height	,      width=self.width						)
									elif self.is_input_patches:
															assert sample_size is not None, "Transformer2DModel over patched input must provide sample_size"

															UpperCamelCase							:		Optional[Any]							      = sample_size
															UpperCamelCase							:		int							      = sample_size

															UpperCamelCase							:		int							      = patch_size
															UpperCamelCase							:		Optional[Any]							      = PatchEmbed(
															    height=_A	,      width=_A	,      patch_size=_A	,      in_channels=_A	,      embed_dim=_A	,      )

									# 3. Define transformers blocks
									UpperCamelCase							:		List[Any]							      = nn.ModuleList(
									    [
									        BasicTransformerBlock(
									            _A	,      _A	,      _A	,      dropout=_A	,      cross_attention_dim=_A	,      activation_fn=_A	,      num_embeds_ada_norm=_A	,      attention_bias=_A	,      only_cross_attention=_A	,      upcast_attention=_A	,      norm_type=_A	,      norm_elementwise_affine=_A	,      )
									        for d in range(_A						)
									    ]						)

									# 4. Define output layers
									UpperCamelCase							:		Union[str, Any]							      = in_channels if out_channels is None else out_channels
									if self.is_input_continuous:
															# TODO: should use out_channels for continuous projections
															if use_linear_projection:
																					UpperCamelCase							:		Optional[int]							      = nn.Linear(_A	,      _A						)
															else:
																					UpperCamelCase							:		str							      = nn.Convad(_A	,      _A	,      kernel_size=1	,      stride=1	,      padding=0						)
									elif self.is_input_vectorized:
															UpperCamelCase							:		str							      = nn.LayerNorm(_A						)
															UpperCamelCase							:		Tuple							      = nn.Linear(_A	,      self.num_vector_embeds - 1						)
									elif self.is_input_patches:
															UpperCamelCase							:		Any							      = nn.LayerNorm(_A	,      elementwise_affine=_A	,      eps=1e-6						)
															UpperCamelCase							:		Optional[Any]							      = nn.Linear(_A	,      2 * inner_dim						)
															UpperCamelCase							:		List[Any]							      = nn.Linear(_A	,      patch_size * patch_size * self.out_channels						)







			def      _a			(			self	,      _A	,      _A = None	,      _A = None	,      _A = None	,      _A = None	,      _A = None	,      _A = None	,      _A = True	,      ):






									'''simple docstring'''







									if attention_mask is not None and attention_mask.ndim == 2:
															# assume that mask is expressed as:
															#   (1 = keep,      0 = discard)
															# convert mask into a bias that can be added to attention scores:
															#       (keep = +0,     discard = -10000.0)
															UpperCamelCase							:		List[str]							      = (1 - attention_mask.to(hidden_states.dtype						)) * -1_00_00.0
															UpperCamelCase							:		Any							      = attention_mask.unsqueeze(1						)

									# convert encoder_attention_mask to a bias the same way we do for attention_mask
									if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2:
															UpperCamelCase							:		Union[str, Any]							      = (1 - encoder_attention_mask.to(hidden_states.dtype						)) * -1_00_00.0
															UpperCamelCase							:		Dict							      = encoder_attention_mask.unsqueeze(1						)

									# 1. Input
									if self.is_input_continuous:
															UpperCamelCase      ,	UpperCamelCase      ,	UpperCamelCase      ,	UpperCamelCase							:		int							      = hidden_states.shape
															UpperCamelCase							:		int							      = hidden_states

															UpperCamelCase							:		List[Any]							      = self.norm(_A						)
															if not self.use_linear_projection:
																					UpperCamelCase							:		Union[str, Any]							      = self.proj_in(_A						)
																					UpperCamelCase							:		Any							      = hidden_states.shape[1]
																					UpperCamelCase							:		Optional[int]							      = hidden_states.permute(0	,      2	,      3	,      1						).reshape(_A	,      height * width	,      _A						)
															else:
																					UpperCamelCase							:		Optional[int]							      = hidden_states.shape[1]
																					UpperCamelCase							:		Optional[Any]							      = hidden_states.permute(0	,      2	,      3	,      1						).reshape(_A	,      height * width	,      _A						)
																					UpperCamelCase							:		Optional[int]							      = self.proj_in(_A						)
									elif self.is_input_vectorized:
															UpperCamelCase							:		Dict							      = self.latent_image_embedding(_A						)
									elif self.is_input_patches:
															UpperCamelCase							:		str							      = self.pos_embed(_A						)

									# 2. Blocks
									for block in self.transformer_blocks:
															UpperCamelCase							:		Union[str, Any]							      = block(
															    _A	,      attention_mask=_A	,      encoder_hidden_states=_A	,      encoder_attention_mask=_A	,      timestep=_A	,      cross_attention_kwargs=_A	,      class_labels=_A	,      )

									# 3. Output
									if self.is_input_continuous:
															if not self.use_linear_projection:
																					UpperCamelCase							:		Optional[int]							      = hidden_states.reshape(_A	,      _A	,      _A	,      _A						).permute(0	,      3	,      1	,      2						).contiguous()
																					UpperCamelCase							:		Tuple							      = self.proj_out(_A						)
															else:
																					UpperCamelCase							:		List[Any]							      = self.proj_out(_A						)
																					UpperCamelCase							:		str							      = hidden_states.reshape(_A	,      _A	,      _A	,      _A						).permute(0	,      3	,      1	,      2						).contiguous()

															UpperCamelCase							:		int							      = hidden_states + residual
									elif self.is_input_vectorized:
															UpperCamelCase							:		List[str]							      = self.norm_out(_A						)
															UpperCamelCase							:		Dict							      = self.out(_A						)
															# (batch, self.num_vector_embeds - 1, self.num_latent_pixels)
															UpperCamelCase							:		Tuple							      = logits.permute(0	,      2	,      1						)

															# log(p(x_0))
															UpperCamelCase							:		str							      = F.log_softmax(logits.double()	,      dim=1						).float()
									elif self.is_input_patches:
															# TODO: cleanup!
															UpperCamelCase							:		Union[str, Any]							      = self.transformer_blocks[0].norma.emb(
															    _A	,      _A	,      hidden_dtype=hidden_states.dtype						)
															UpperCamelCase      ,	UpperCamelCase							:		Dict							      = self.proj_out_a(F.silu(_A						)						).chunk(2	,      dim=1						)
															UpperCamelCase							:		Union[str, Any]							      = self.norm_out(_A						) * (1 + scale[:, None]) + shift[:, None]
															UpperCamelCase							:		Dict							      = self.proj_out_a(_A						)

															# unpatchify
															UpperCamelCase							:		Tuple							      = int(hidden_states.shape[1] ** 0.5						)
															UpperCamelCase							:		Dict							      = hidden_states.reshape(
															    shape=(-1, height, width, self.patch_size, self.patch_size, self.out_channels)						)
															UpperCamelCase							:		Optional[Any]							      = torch.einsum("""nhwpqc->nchpwq"""	,      _A						)
															UpperCamelCase							:		Union[str, Any]							      = hidden_states.reshape(
															    shape=(-1, self.out_channels, height * self.patch_size, width * self.patch_size)						)

									if not return_dict:
															return (output,)

									return TransformeraDModelOutput(sample=_A						)



