






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

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

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .attention import BasicTransformerBlock
from .attention_processor import AttentionProcessor, AttnProcessor
from .embeddings import TimestepEmbedding, Timesteps
from .modeling_utils import ModelMixin
@dataclass
class       _UpperCAmelCase					(       __snake_case	):

				"""simple docstring"""




				lowercase__											=       42



class       _UpperCAmelCase					(       __snake_case		,__snake_case	):

				"""simple docstring"""




				@register_to_config
				def __init__(     self						:    List[Any],  lowerCamelCase						:    List[Any] = 32,  lowerCamelCase						:    str = 64,  lowerCamelCase						:    str = 20,  lowerCamelCase						:    Any = 768,  lowerCamelCase						:    List[Any]=77,  lowerCamelCase						:    Dict=4,  lowerCamelCase						:    Any = 0.0,  lowerCamelCase						:    List[Any] = "silu",  lowerCamelCase						:    List[Any] = None,  lowerCamelCase						:    Union[str, Any] = None,  lowerCamelCase						:    int = "linear",  lowerCamelCase						:    Optional[int] = "prd",  lowerCamelCase						:    Optional[Any] = None,  lowerCamelCase						:    Any = None,  lowerCamelCase						:    List[Any] = None,  ):

								'''simple docstring'''
								super().__init__()
								lowercase__    =			num_attention_heads
								lowercase__    =			attention_head_dim
								lowercase__    =			num_attention_heads * attention_head_dim
								lowercase__    =			additional_embeddings

								lowercase__    =			time_embed_dim or inner_dim
								lowercase__    =			embedding_proj_dim or embedding_dim
								lowercase__    =			clip_embed_dim or embedding_dim

								lowercase__    =			Timesteps(A_,  A_,  0  )
								lowercase__    =			TimestepEmbedding(A_,  A_,  out_dim=A_,  act_fn=A_  )

								lowercase__    =			nn.Linear(A_,  A_  )

								if embedding_proj_norm_type is None:
												lowercase__    =			None
								elif embedding_proj_norm_type == "layer":
												lowercase__    =			nn.LayerNorm(A_  )
								else:
												raise ValueError(F"""unsupported embedding_proj_norm_type: {embedding_proj_norm_type}"""  )

								lowercase__    =			nn.Linear(A_,  A_  )

								if encoder_hid_proj_type is None:
												lowercase__    =			None
								elif encoder_hid_proj_type == "linear":
												lowercase__    =			nn.Linear(A_,  A_  )
								else:
												raise ValueError(F"""unsupported encoder_hid_proj_type: {encoder_hid_proj_type}"""  )

								lowercase__    =			nn.Parameter(torch.zeros(1,  num_embeddings + additional_embeddings,  A_  )  )

								if added_emb_type == "prd":
												lowercase__    =			nn.Parameter(torch.zeros(1,  1,  A_  )  )
								elif added_emb_type is None:
												lowercase__    =			None
								else:
												raise ValueError(
												    F"""`added_emb_type`: {added_emb_type} is not supported. Make sure to choose one of `'prd'` or `None`."""  )

								lowercase__    =			nn.ModuleList(
								    [
								        BasicTransformerBlock(
								            A_,  A_,  A_,  dropout=A_,  activation_fn='''gelu''',  attention_bias=A_,  )
								        for d in range(A_  )
								    ]  )

								if norm_in_type == "layer":
												lowercase__    =			nn.LayerNorm(A_  )
								elif norm_in_type is None:
												lowercase__    =			None
								else:
												raise ValueError(F"""Unsupported norm_in_type: {norm_in_type}."""  )

								lowercase__    =			nn.LayerNorm(A_  )

								lowercase__    =			nn.Linear(A_,  A_  )

								lowercase__    =			torch.full(
								    [num_embeddings + additional_embeddings, num_embeddings + additional_embeddings],  -10_000.0  )
								causal_attention_mask.triu_(1  )
								lowercase__    =			causal_attention_mask[None, ...]
								self.register_buffer('''causal_attention_mask''',  A_,  persistent=A_  )

								lowercase__    =			nn.Parameter(torch.zeros(1,  A_  )  )
								lowercase__    =			nn.Parameter(torch.zeros(1,  A_  )  )




				@property
				# Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
				def 				lowercase__   (     self						:    List[Any]  ):

								'''simple docstring'''
								lowercase__    =			{}

								def fn_recursive_add_processors(lowerCamelCase						:    str,  lowerCamelCase						:    int,  lowerCamelCase						:    Dict  ):
												if hasattr(A_,  '''set_processor'''  ):
																lowercase__    =			module.processor

												for sub_name, child in module.named_children():
																fn_recursive_add_processors(F"""{name}.{sub_name}""",  A_,  A_  )

												return processors

								for name, module in self.named_children():
												fn_recursive_add_processors(A_,  A_,  A_  )

								return processors




				def 				lowercase__   (     self						:    str,  lowerCamelCase						:    List[Any]  ):

								'''simple docstring'''
								lowercase__    =			len(self.attn_processors.keys()  )

								if isinstance(A_,  A_  ) and len(A_  ) != count:
												raise ValueError(
												    F"""A dict of processors was passed, but the number of processors {len(A_  )} does not match the"""
												    F""" number of attention layers: {count}. Please make sure to pass {count} processor classes."""  )

								def fn_recursive_attn_processor(lowerCamelCase						:    List[Any],  lowerCamelCase						:    Any,  lowerCamelCase						:    int  ):
												if hasattr(A_,  '''set_processor'''  ):
																if not isinstance(A_,  A_  ):
																				module.set_processor(A_  )
																else:
																				module.set_processor(processor.pop(F"""{name}.processor"""  )  )

												for sub_name, child in module.named_children():
																fn_recursive_attn_processor(F"""{name}.{sub_name}""",  A_,  A_  )

								for name, module in self.named_children():
												fn_recursive_attn_processor(A_,  A_,  A_  )




				def 				lowercase__   (     self						:    Tuple  ):

								'''simple docstring'''
								self.set_attn_processor(AttnProcessor()  )




				def 				lowercase__   (     self						:    int,  lowerCamelCase						:    Tuple,  lowerCamelCase						:    Optional[Any],  lowerCamelCase						:    List[Any],  lowerCamelCase						:    str = None,  lowerCamelCase						:    int = None,  lowerCamelCase						:    Any = True,  ):

								'''simple docstring'''
								lowercase__    =			hidden_states.shape[0]

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

								# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
								lowercase__    =			timesteps * torch.ones(A_,  dtype=timesteps.dtype,  device=timesteps.device  )

								lowercase__    =			self.time_proj(A_  )

								# timesteps does not contain any weights and will always return f32 tensors
								# but time_embedding might be fp16, so we need to cast here.
								lowercase__    =			timesteps_projected.to(dtype=self.dtype  )
								lowercase__    =			self.time_embedding(A_  )

								if self.embedding_proj_norm is not None:
												lowercase__    =			self.embedding_proj_norm(A_  )

								lowercase__    =			self.embedding_proj(A_  )
								if self.encoder_hidden_states_proj is not None and encoder_hidden_states is not None:
												lowercase__    =			self.encoder_hidden_states_proj(A_  )
								elif self.encoder_hidden_states_proj is not None and encoder_hidden_states is None:
												raise ValueError('''`encoder_hidden_states_proj` requires `encoder_hidden_states` to be set'''  )

								lowercase__    =			self.proj_in(A_  )

								lowercase__    =			self.positional_embedding.to(hidden_states.dtype  )

								lowercase__    =			[]
								lowercase__    =			0

								if encoder_hidden_states is not None:
												additional_embeds.append(A_  )
												additional_embeddings_len += encoder_hidden_states.shape[1]

								if len(proj_embeddings.shape  ) == 2:
												lowercase__    =			proj_embeddings[:, None, :]

								if len(hidden_states.shape  ) == 2:
												lowercase__    =			hidden_states[:, None, :]

								lowercase__    =			additional_embeds + [
								    proj_embeddings,
								    time_embeddings[:, None, :],
								    hidden_states,
								]

								if self.prd_embedding is not None:
												lowercase__    =			self.prd_embedding.to(hidden_states.dtype  ).expand(A_,  -1,  -1  )
												additional_embeds.append(A_  )

								lowercase__    =			torch.cat(
								    A_,  dim=1,  )

								# Allow positional_embedding to not include the `addtional_embeddings` and instead pad it with zeros for these additional tokens
								lowercase__    =			additional_embeddings_len + proj_embeddings.shape[1] + 1
								if positional_embeddings.shape[1] < hidden_states.shape[1]:
												lowercase__    =			F.pad(
												    A_,  (
												        0,
												        0,
												        additional_embeddings_len,
												        self.prd_embedding.shape[1] if self.prd_embedding is not None else 0,
												    ),  value=0.0,  )

								lowercase__    =			hidden_states + positional_embeddings

								if attention_mask is not None:
												lowercase__    =			(1 - attention_mask.to(hidden_states.dtype  )) * -10_000.0
												lowercase__    =			F.pad(A_,  (0, self.additional_embeddings),  value=0.0  )
												lowercase__    =			(attention_mask[:, None, :] + self.causal_attention_mask).to(hidden_states.dtype  )
												lowercase__    =			attention_mask.repeat_interleave(self.config.num_attention_heads,  dim=0  )

								if self.norm_in is not None:
												lowercase__    =			self.norm_in(A_  )

								for block in self.transformer_blocks:
												lowercase__    =			block(A_,  attention_mask=A_  )

								lowercase__    =			self.norm_out(A_  )

								if self.prd_embedding is not None:
												lowercase__    =			hidden_states[:, -1]
								else:
												lowercase__    =			hidden_states[:, additional_embeddings_len:]

								lowercase__    =			self.proj_to_clip_embeddings(A_  )

								if not return_dict:
												return (predicted_image_embedding,)

								return PriorTransformerOutput(predicted_image_embedding=A_  )




				def 				lowercase__   (     self						:    Optional[Any],  lowerCamelCase						:    int  ):

								'''simple docstring'''
								lowercase__    =			(prior_latents * self.clip_std) + self.clip_mean
								return prior_latents











import unittest

from transformers import LiltConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
    import torch

    from transformers import (
        LiltForQuestionAnswering,
        LiltForSequenceClassification,
        LiltForTokenClassification,
        LiltModel,
    )
    from transformers.models.lilt.modeling_lilt import LILT_PRETRAINED_MODEL_ARCHIVE_LIST



class 					A__   :





      def __init__( self       ,							A_       ,							A_=13       ,							A_=7       ,							A_=True       ,							A_=True       ,							A_=True       ,							A_=True       ,							A_=99       ,							A_=24       ,							A_=2       ,							A_=6       ,							A_=37       ,							A_="gelu"       ,							A_=0.1       ,							A_=0.1       ,							A_=512       ,							A_=16       ,							A_=2       ,							A_=0.02       ,							A_=3       ,							A_=None       ,							A_=1000       ,							):
       '''simple docstring'''




       UpperCamelCase	:					Union[str, Any]       	=   parent
       UpperCamelCase	:					List[Any]       	=   batch_size
       UpperCamelCase	:					Dict       	=   seq_length
       UpperCamelCase	:					Tuple       	=   is_training
       UpperCamelCase	:					Union[str, Any]       	=   use_input_mask
       UpperCamelCase	:					Tuple       	=   use_token_type_ids
       UpperCamelCase	:					Optional[Any]       	=   use_labels
       UpperCamelCase	:					str       	=   vocab_size
       UpperCamelCase	:					Optional[int]       	=   hidden_size
       UpperCamelCase	:					Any       	=   num_hidden_layers
       UpperCamelCase	:					Optional[Any]       	=   num_attention_heads
       UpperCamelCase	:					Optional[Any]       	=   intermediate_size
       UpperCamelCase	:					Optional[Any]       	=   hidden_act
       UpperCamelCase	:					Union[str, Any]       	=   hidden_dropout_prob
       UpperCamelCase	:					Union[str, Any]       	=   attention_probs_dropout_prob
       UpperCamelCase	:					List[Any]       	=   max_position_embeddings
       UpperCamelCase	:					str       	=   type_vocab_size
       UpperCamelCase	:					Optional[int]       	=   type_sequence_label_size
       UpperCamelCase	:					Dict       	=   initializer_range
       UpperCamelCase	:					int       	=   num_labels
       UpperCamelCase	:					Optional[int]       	=   scope
       UpperCamelCase	:					int       	=   range_bbox





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					int       	=   ids_tensor([self.batch_size, self.seq_length]       ,							self.vocab_size   )

       UpperCamelCase	:					Any       	=   ids_tensor([self.batch_size, self.seq_length, 4]       ,							self.range_bbox   )
       # Ensure that bbox is legal
       for i in range(bbox.shape[0]   ):
        for j in range(bbox.shape[1]   ):
         if bbox[i, j, 3] < bbox[i, j, 1]:
          UpperCamelCase	:					Union[str, Any]       	=   bbox[i, j, 3]
          UpperCamelCase	:					int       	=   bbox[i, j, 1]
          UpperCamelCase	:					int       	=   t
         if bbox[i, j, 2] < bbox[i, j, 0]:
          UpperCamelCase	:					List[str]       	=   bbox[i, j, 2]
          UpperCamelCase	:					Optional[int]       	=   bbox[i, j, 0]
          UpperCamelCase	:					Optional[Any]       	=   t

       UpperCamelCase	:					Dict       	=   None
       if self.use_input_mask:
        UpperCamelCase	:					Optional[int]       	=   ids_tensor([self.batch_size, self.seq_length]       ,							vocab_size=2   )

       UpperCamelCase	:					str       	=   None
       if self.use_token_type_ids:
        UpperCamelCase	:					str       	=   ids_tensor([self.batch_size, self.seq_length]       ,							self.type_vocab_size   )

       UpperCamelCase	:					Dict       	=   None
       UpperCamelCase	:					int       	=   None
       if self.use_labels:
        UpperCamelCase	:					List[Any]       	=   ids_tensor([self.batch_size]       ,							self.type_sequence_label_size   )
        UpperCamelCase	:					int       	=   ids_tensor([self.batch_size, self.seq_length]       ,							self.num_labels   )

       UpperCamelCase	:					List[Any]       	=   self.get_config()

       return config, input_ids, bbox, token_type_ids, input_mask, sequence_labels, token_labels





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       return LiltConfig(
           vocab_size=self.vocab_size       ,							hidden_size=self.hidden_size       ,							num_hidden_layers=self.num_hidden_layers       ,							num_attention_heads=self.num_attention_heads       ,							intermediate_size=self.intermediate_size       ,							hidden_act=self.hidden_act       ,							hidden_dropout_prob=self.hidden_dropout_prob       ,							attention_probs_dropout_prob=self.attention_probs_dropout_prob       ,							max_position_embeddings=self.max_position_embeddings       ,							type_vocab_size=self.type_vocab_size       ,							initializer_range=self.initializer_range       ,							)





      def        __UpperCamelCase( self       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							):
       '''simple docstring'''




       UpperCamelCase	:					Any       	=   LiltModel(config=A_   )
       model.to(A_   )
       model.eval()
       UpperCamelCase	:					str       	=   model(A_       ,							bbox=A_       ,							attention_mask=A_       ,							token_type_ids=A_   )
       UpperCamelCase	:					Optional[int]       	=   model(A_       ,							bbox=A_       ,							token_type_ids=A_   )
       UpperCamelCase	:					Any       	=   model(A_       ,							bbox=A_   )

       self.parent.assertEqual(result.last_hidden_state.shape       ,							(self.batch_size, self.seq_length, self.hidden_size)   )
       self.parent.assertEqual(result.pooler_output.shape       ,							(self.batch_size, self.hidden_size)   )





      def        __UpperCamelCase( self       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							):
       '''simple docstring'''




       UpperCamelCase	:					Any       	=   self.num_labels
       UpperCamelCase	:					Dict       	=   LiltForTokenClassification(config=A_   )
       model.to(A_   )
       model.eval()
       UpperCamelCase	:					Dict       	=   model(
           A_       ,							bbox=A_       ,							attention_mask=A_       ,							token_type_ids=A_       ,							labels=A_   )
       self.parent.assertEqual(result.logits.shape       ,							(self.batch_size, self.seq_length, self.num_labels)   )





      def        __UpperCamelCase( self       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							):
       '''simple docstring'''




       UpperCamelCase	:					Dict       	=   LiltForQuestionAnswering(config=A_   )
       model.to(A_   )
       model.eval()
       UpperCamelCase	:					List[str]       	=   model(
           A_       ,							bbox=A_       ,							attention_mask=A_       ,							token_type_ids=A_       ,							start_positions=A_       ,							end_positions=A_       ,							)
       self.parent.assertEqual(result.start_logits.shape       ,							(self.batch_size, self.seq_length)   )
       self.parent.assertEqual(result.end_logits.shape       ,							(self.batch_size, self.seq_length)   )





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Any       	=   self.prepare_config_and_inputs()
       (
           (
           UpperCamelCase
       )	,    (
           UpperCamelCase
       )	,    (
           UpperCamelCase
       )	,    (
           UpperCamelCase
       )	,    (
           UpperCamelCase
       )	,    (
           UpperCamelCase
       )	,    (
           UpperCamelCase
       )	,    
       )	:					Tuple       	=   config_and_inputs
       UpperCamelCase	:					Tuple       	=   {
           "input_ids": input_ids,
           "bbox": bbox,
           "token_type_ids": token_type_ids,
           "attention_mask": input_mask,
       }
       return config, inputs_dict



@require_torch
class 					A__   (       __snake_case       ,   __snake_case       ,   __snake_case       ,   unittest.TestCase			):
      _UpperCAmelCase	:Union[str, Any]						    =   (
          (
              LiltModel,
              LiltForSequenceClassification,
              LiltForTokenClassification,
              LiltForQuestionAnswering,
          )
          if is_torch_available()
          else ()
      )
      _UpperCAmelCase	:Optional[Any]						    =   (
          {
              'feature-extraction': LiltModel,
              'question-answering': LiltForQuestionAnswering,
              'text-classification': LiltForSequenceClassification,
              'token-classification': LiltForTokenClassification,
              'zero-shot': LiltForSequenceClassification,
          }
          if is_torch_available()
          else {}
      )
      _UpperCAmelCase	:Dict						    =   False
      _UpperCAmelCase	:Union[str, Any]						    =   False





      def        __UpperCamelCase( self       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_   ):
       '''simple docstring'''




       return True





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Optional[int]       	=   LiltModelTester(self   )
       UpperCamelCase	:					Optional[int]       	=   ConfigTester(self       ,							config_class=A_       ,							hidden_size=37   )





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       self.config_tester.run_common_tests()





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Tuple       	=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*A_   )





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Tuple       	=   self.model_tester.prepare_config_and_inputs()
       for type in ["absolute", "relative_key", "relative_key_query"]:
        UpperCamelCase	:					Union[str, Any]       	=   type
        self.model_tester.create_and_check_model(*A_   )





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					int       	=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_token_classification(*A_   )





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					List[Any]       	=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_question_answering(*A_   )





      @slow
      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       for model_name in LILT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        UpperCamelCase	:					Dict       	=   LiltModel.from_pretrained(A_   )
        self.assertIsNotNone(A_   )



@require_torch
@slow
class 					A__   (       unittest.TestCase			):





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					int       	=   LiltModel.from_pretrained("SCUT-DLVCLab/lilt-roberta-en-base"   ).to(A_   )

       UpperCamelCase	:					Tuple       	=   torch.tensor([[1, 2]]       ,							device=A_   )
       UpperCamelCase	:					List[str]       	=   torch.tensor([[[1, 2, 3, 4], [5, 6, 7, 8]]]       ,							device=A_   )

       # forward pass
       with torch.no_grad():
        UpperCamelCase	:					Optional[int]       	=   model(input_ids=A_       ,							bbox=A_   )

       UpperCamelCase	:					List[str]       	=   torch.Size([1, 2, 768]   )
       UpperCamelCase	:					Any       	=   torch.tensor(
           [[-0.06_53, 0.09_50, -0.00_61], [-0.05_45, 0.09_26, -0.03_24]]       ,							device=A_       ,							)

       self.assertTrue(outputs.last_hidden_state.shape       ,							A_   )
       self.assertTrue(torch.allclose(outputs.last_hidden_state[0, :, :3]       ,							A_       ,							atol=1e-3   )   )
