




"""simple docstring"""

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 				lowercase_       :



 '''simple docstring'''

 def __init__( self       :       Tuple    ,							_UpperCAmelCase       :       Optional[int]    ,							_UpperCAmelCase       :       Union[str, Any]=13    ,							_UpperCAmelCase       :       Any=7    ,							_UpperCAmelCase       :       Optional[Any]=True    ,							_UpperCAmelCase       :       Union[str, Any]=True    ,							_UpperCAmelCase       :       Tuple=True    ,							_UpperCAmelCase       :       Tuple=True    ,							_UpperCAmelCase       :       Dict=99    ,							_UpperCAmelCase       :       Dict=24    ,							_UpperCAmelCase       :       List[str]=2    ,							_UpperCAmelCase       :       Any=6    ,							_UpperCAmelCase       :       Optional[Any]=37    ,							_UpperCAmelCase       :       Any="gelu"    ,							_UpperCAmelCase       :       Tuple=0.1    ,							_UpperCAmelCase       :       Any=0.1    ,							_UpperCAmelCase       :       Any=512    ,							_UpperCAmelCase       :       Union[str, Any]=16    ,							_UpperCAmelCase       :       Tuple=2    ,							_UpperCAmelCase       :       Optional[Any]=0.02    ,							_UpperCAmelCase       :       Tuple=3    ,							_UpperCAmelCase       :       Dict=None    ,							_UpperCAmelCase       :       List[str]=1_000    ,							):
       _A					   =   parent
       _A					   =   batch_size
       _A					   =   seq_length
       _A					   =   is_training
       _A					   =   use_input_mask
       _A					   =   use_token_type_ids
       _A					   =   use_labels
       _A					   =   vocab_size
       _A					   =   hidden_size
       _A					   =   num_hidden_layers
       _A					   =   num_attention_heads
       _A					   =   intermediate_size
       _A					   =   hidden_act
       _A					   =   hidden_dropout_prob
       _A					   =   attention_probs_dropout_prob
       _A					   =   max_position_embeddings
       _A					   =   type_vocab_size
       _A					   =   type_sequence_label_size
       _A					   =   initializer_range
       _A					   =   num_labels
       _A					   =   scope
       _A					   =   range_bbox

 def 			lowerCAmelCase_							( self       :       Optional[int]		):
       _A					   =   ids_tensor([self.batch_size, self.seq_length]    ,							self.vocab_size		)

       _A					   =   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]:
                         _A					   =   bbox[i, j, 3]
                         _A					   =   bbox[i, j, 1]
                         _A					   =   t
                   if bbox[i, j, 2] < bbox[i, j, 0]:
                         _A					   =   bbox[i, j, 2]
                         _A					   =   bbox[i, j, 0]
                         _A					   =   t

       _A					   =   None
       if self.use_input_mask:
             _A					   =   ids_tensor([self.batch_size, self.seq_length]    ,							vocab_size=2		)

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

       _A					   =   None
       _A					   =   None
       if self.use_labels:
             _A					   =   ids_tensor([self.batch_size]    ,							self.type_sequence_label_size		)
             _A					   =   ids_tensor([self.batch_size, self.seq_length]    ,							self.num_labels		)

       _A					   =   self.get_config()

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

 def 			lowerCAmelCase_							( self       :       Dict		):
       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 			lowerCAmelCase_							( self       :       Dict    ,							_UpperCAmelCase       :       List[str]    ,							_UpperCAmelCase       :       Optional[int]    ,							_UpperCAmelCase       :       Optional[int]    ,							_UpperCAmelCase       :       Tuple    ,							_UpperCAmelCase       :       List[str]    ,							_UpperCAmelCase       :       Any    ,							_UpperCAmelCase       :       List[str]    ,							):
       _A					   =   LiltModel(config=UpperCamelCase__		)
       model.to(UpperCamelCase__		)
       model.eval()
       _A					   =   model(UpperCamelCase__    ,							bbox=UpperCamelCase__    ,							attention_mask=UpperCamelCase__    ,							token_type_ids=UpperCamelCase__		)
       _A					   =   model(UpperCamelCase__    ,							bbox=UpperCamelCase__    ,							token_type_ids=UpperCamelCase__		)
       _A					   =   model(UpperCamelCase__    ,							bbox=UpperCamelCase__		)

       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 			lowerCAmelCase_							( self       :       Any    ,							_UpperCAmelCase       :       Dict    ,							_UpperCAmelCase       :       List[Any]    ,							_UpperCAmelCase       :       List[str]    ,							_UpperCAmelCase       :       Tuple    ,							_UpperCAmelCase       :       List[Any]    ,							_UpperCAmelCase       :       Optional[Any]    ,							_UpperCAmelCase       :       Union[str, Any]    ,							):
       _A					   =   self.num_labels
       _A					   =   LiltForTokenClassification(config=UpperCamelCase__		)
       model.to(UpperCamelCase__		)
       model.eval()
       _A					   =   model(
           UpperCamelCase__    ,							bbox=UpperCamelCase__    ,							attention_mask=UpperCamelCase__    ,							token_type_ids=UpperCamelCase__    ,							labels=UpperCamelCase__		)
       self.parent.assertEqual(result.logits.shape    ,							(self.batch_size, self.seq_length, self.num_labels)		)

 def 			lowerCAmelCase_							( self       :       List[str]    ,							_UpperCAmelCase       :       List[str]    ,							_UpperCAmelCase       :       Union[str, Any]    ,							_UpperCAmelCase       :       Union[str, Any]    ,							_UpperCAmelCase       :       Any    ,							_UpperCAmelCase       :       Tuple    ,							_UpperCAmelCase       :       int    ,							_UpperCAmelCase       :       Dict    ,							):
       _A					   =   LiltForQuestionAnswering(config=UpperCamelCase__		)
       model.to(UpperCamelCase__		)
       model.eval()
       _A					   =   model(
           UpperCamelCase__    ,							bbox=UpperCamelCase__    ,							attention_mask=UpperCamelCase__    ,							token_type_ids=UpperCamelCase__    ,							start_positions=UpperCamelCase__    ,							end_positions=UpperCamelCase__    ,							)
       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 			lowerCAmelCase_							( self       :       int		):
       _A					   =   self.prepare_config_and_inputs()
       (
           (
           _A      
       )   ,    (
           _A      
       )   ,    (
           _A      
       )   ,    (
           _A      
       )   ,    (
           _A      
       )   ,    (
           _A      
       )   ,    (
           _A      
       )   ,          
       )					   =   config_and_inputs
       _A					   =   {
           'input_ids': input_ids,
           'bbox': bbox,
           'token_type_ids': token_type_ids,
           'attention_mask': input_mask,
       }
       return config, inputs_dict





@require_torch
class 				lowercase_       (     _a			,	_a			,	_a			,	unittest.TestCase							):



 '''simple docstring'''
 UpperCAmelCase    :     str        =	(
     (
         LiltModel,
         LiltForSequenceClassification,
         LiltForTokenClassification,
         LiltForQuestionAnswering,
     )
     if is_torch_available()
     else ()
 )
 UpperCAmelCase    :     Union[str, Any]        =	(
     {
         '''feature-extraction''': LiltModel,
         '''question-answering''': LiltForQuestionAnswering,
         '''text-classification''': LiltForSequenceClassification,
         '''token-classification''': LiltForTokenClassification,
         '''zero-shot''': LiltForSequenceClassification,
     }
     if is_torch_available()
     else {}
 )
 UpperCAmelCase    :     Tuple        =	False
 UpperCAmelCase    :     Optional[Any]        =	False

 def 			lowerCAmelCase_							( self       :       Union[str, Any]    ,							_UpperCAmelCase       :       Any    ,							_UpperCAmelCase       :       int    ,							_UpperCAmelCase       :       Optional[int]    ,							_UpperCAmelCase       :       List[str]    ,							_UpperCAmelCase       :       List[Any]		):
       return True

 def 			lowerCAmelCase_							( self       :       Union[str, Any]		):
       _A					   =   LiltModelTester(self		)
       _A					   =   ConfigTester(self    ,							config_class=UpperCamelCase__    ,							hidden_size=37		)

 def 			lowerCAmelCase_							( self       :       Union[str, Any]		):
       self.config_tester.run_common_tests()

 def 			lowerCAmelCase_							( self       :       str		):
       _A					   =   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*UpperCamelCase__		)

 def 			lowerCAmelCase_							( self       :       List[str]		):
       _A					   =   self.model_tester.prepare_config_and_inputs()
       for type in ["absolute", "relative_key", "relative_key_query"]:
             _A					   =   type
             self.model_tester.create_and_check_model(*UpperCamelCase__		)

 def 			lowerCAmelCase_							( self       :       Dict		):
       _A					   =   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_token_classification(*UpperCamelCase__		)

 def 			lowerCAmelCase_							( self       :       List[Any]		):
       _A					   =   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_question_answering(*UpperCamelCase__		)



 @slow
 def 			lowerCAmelCase_							( self       :       Union[str, Any]		):







       for model_name in LILT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             _A					   =   LiltModel.from_pretrained(UpperCamelCase__		)
             self.assertIsNotNone(UpperCamelCase__		)





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



 '''simple docstring'''

 def 			lowerCAmelCase_							( self       :       str		):
       _A					   =   LiltModel.from_pretrained('SCUT-DLVCLab/lilt-roberta-en-base'		).to(UpperCamelCase__		)

       _A					   =   torch.tensor([[1, 2]]    ,							device=UpperCamelCase__		)
       _A					   =   torch.tensor([[[1, 2, 3, 4], [5, 6, 7, 8]]]    ,							device=UpperCamelCase__		)

       # forward pass
       with torch.no_grad():
             _A					   =   model(input_ids=UpperCamelCase__    ,							bbox=UpperCamelCase__		)

       _A					   =   torch.Size([1, 2, 768]		)
       _A					   =   torch.tensor(
           [[-0.0653, 0.0950, -0.0061], [-0.0545, 0.0926, -0.0324]]    ,							device=UpperCamelCase__    ,							)

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






'''simple docstring'''


from typing import Optional, Tuple

import jax
import jax.numpy as jnp
from flax import linen as nn
from flax.core.frozen_dict import FrozenDict
from transformers import CLIPConfig, FlaxPreTrainedModel
from transformers.models.clip.modeling_flax_clip import FlaxCLIPVisionModule







def 				__lowerCamelCase				(   A__    ,						A__    ,						A__=1e-1_2    )       ->	Dict:






   """simple docstring"""
   UpperCamelCase					=  jnp.divide(emb_a.T    ,						jnp.clip(jnp.linalg.norm(A__    ,						axis=1    )    ,						a_min=A__    )    ).T
   UpperCamelCase					=  jnp.divide(emb_a.T    ,						jnp.clip(jnp.linalg.norm(A__    ,						axis=1    )    ,						a_min=A__    )    ).T
   return jnp.matmul(A__    ,						norm_emb_a.T    )

class  SCREAMING_SNAKE_CASE (  nn.Module		):






    """simple docstring"""





    _SCREAMING_SNAKE_CASE	  =  42
    _SCREAMING_SNAKE_CASE	  =  jnp.floataa






    def  A       (  self      :					List[Any] ):


       """simple docstring"""







       UpperCamelCase					=  FlaxCLIPVisionModule(self.config.vision_config )
       UpperCamelCase					=  nn.Dense(self.config.projection_dim  ,						use_bias=UpperCamelCase__  ,						dtype=self.dtype )

       UpperCamelCase					=  self.param('concept_embeds'  ,						jax.nn.initializers.ones  ,						(1_7, self.config.projection_dim) )
       UpperCamelCase					=  self.param(
           'special_care_embeds'  ,						jax.nn.initializers.ones  ,						(3, self.config.projection_dim) )

       UpperCamelCase					=  self.param('concept_embeds_weights'  ,						jax.nn.initializers.ones  ,						(1_7,) )
       UpperCamelCase					=  self.param('special_care_embeds_weights'  ,						jax.nn.initializers.ones  ,						(3,) )






    def __call__(  self      :					str  ,						UpperCamelCase__      :					List[str] ):


       """simple docstring"""







       UpperCamelCase					=  self.vision_model(UpperCamelCase__ )[1]
       UpperCamelCase					=  self.visual_projection(UpperCamelCase__ )

       UpperCamelCase					=  jax_cosine_distance(UpperCamelCase__  ,						self.special_care_embeds )
       UpperCamelCase					=  jax_cosine_distance(UpperCamelCase__  ,						self.concept_embeds )

       # increase this value to create a stronger `nfsw` filter
       # at the cost of increasing the possibility of filtering benign image inputs
       UpperCamelCase					=  0.0

       UpperCamelCase					=  special_cos_dist - self.special_care_embeds_weights[None, :] + adjustment
       UpperCamelCase					=  jnp.round(UpperCamelCase__  ,						3 )
       UpperCamelCase					=  jnp.any(special_scores > 0  ,						axis=1  ,						keepdims=UpperCamelCase__ )
       # Use a lower threshold if an image has any special care concept
       UpperCamelCase					=  is_special_care * 0.0_1

       UpperCamelCase					=  cos_dist - self.concept_embeds_weights[None, :] + special_adjustment
       UpperCamelCase					=  jnp.round(UpperCamelCase__  ,						3 )
       UpperCamelCase					=  jnp.any(concept_scores > 0  ,						axis=1 )

       return has_nsfw_concepts

class  SCREAMING_SNAKE_CASE (  _a		):






    """simple docstring"""





    _SCREAMING_SNAKE_CASE	  =  CLIPConfig
    _SCREAMING_SNAKE_CASE	  =  """clip_input"""
    _SCREAMING_SNAKE_CASE	  =  FlaxStableDiffusionSafetyCheckerModule






    def __init__(  self      :					Union[str, Any]  ,						UpperCamelCase__      :					CLIPConfig  ,						UpperCamelCase__      :					Optional[Tuple] = None  ,						UpperCamelCase__      :					int = 0  ,						UpperCamelCase__      :					jnp.dtype = jnp.floataa  ,						UpperCamelCase__      :					bool = True  ,						**UpperCamelCase__      :					List[str]  ,						):


       """simple docstring"""







       if input_shape is None:
          UpperCamelCase					=  (1, 2_2_4, 2_2_4, 3)
       UpperCamelCase					=  self.module_class(config=UpperCamelCase__  ,						dtype=UpperCamelCase__  ,						**UpperCamelCase__ )
       super().__init__(UpperCamelCase__  ,						UpperCamelCase__  ,						input_shape=UpperCamelCase__  ,						seed=UpperCamelCase__  ,						dtype=UpperCamelCase__  ,						_do_init=_do_init )






    def  A       (  self      :					int  ,						UpperCamelCase__      :					jax.random.KeyArray  ,						UpperCamelCase__      :					Tuple  ,						UpperCamelCase__      :					FrozenDict = None ):


       """simple docstring"""







       UpperCamelCase					=  jax.random.normal(UpperCamelCase__  ,						UpperCamelCase__ )

       UpperCamelCase				,      UpperCamelCase					=  jax.random.split(UpperCamelCase__ )
       UpperCamelCase					=  {'params': params_rng, 'dropout': dropout_rng}

       UpperCamelCase					=  self.module.init(UpperCamelCase__  ,						UpperCamelCase__ )['params']

       return random_params






    def __call__(  self      :					List[Any]  ,						UpperCamelCase__      :					Dict  ,						UpperCamelCase__      :					dict = None  ,						):


       """simple docstring"""







       UpperCamelCase					=  jnp.transpose(UpperCamelCase__  ,						(0, 2, 3, 1) )

       return self.module.apply(
           {'params': params or self.params}  ,						jnp.array(UpperCamelCase__  ,						dtype=jnp.floataa )  ,						rngs={}  ,						)

