





import inspect
import unittest

import numpy as np

from transformers import ViTConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor


if is_flax_available():
  import jax

  from transformers.models.vit.modeling_flax_vit import FlaxViTForImageClassification, FlaxViTModel





class  lowerCAmelCase							(       unittest.TestCase ):







  def __init__(      self  :Tuple   ,							_lowercase  :Any   ,							_lowercase  :Any=13   ,							_lowercase  :Tuple=30   ,							_lowercase  :str=2   ,							_lowercase  :List[Any]=3   ,							_lowercase  :Dict=True   ,							_lowercase  :Dict=True   ,							_lowercase  :Optional[Any]=32   ,							_lowercase  :Union[str, Any]=5   ,							_lowercase  :Any=4   ,							_lowercase  :Union[str, Any]=37   ,							_lowercase  :str="gelu"   ,							_lowercase  :List[str]=0.1   ,							_lowercase  :Optional[int]=0.1   ,							_lowercase  :Any=10   ,							_lowercase  :Any=0.02   ,							):






     '''simple docstring'''





     lowercase__			       =  parent
     lowercase__			       =  batch_size
     lowercase__			       =  image_size
     lowercase__			       =  patch_size
     lowercase__			       =  num_channels
     lowercase__			       =  is_training
     lowercase__			       =  use_labels
     lowercase__			       =  hidden_size
     lowercase__			       =  num_hidden_layers
     lowercase__			       =  num_attention_heads
     lowercase__			       =  intermediate_size
     lowercase__			       =  hidden_act
     lowercase__			       =  hidden_dropout_prob
     lowercase__			       =  attention_probs_dropout_prob
     lowercase__			       =  type_sequence_label_size
     lowercase__			       =  initializer_range

     # in ViT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
     lowercase__			       =  (image_size // patch_size) ** 2
     lowercase__			       =  num_patches + 1







  def   UpperCAmelCase					(      self  :Dict ):






     '''simple docstring'''





     lowercase__			       =  floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size] )

     lowercase__			       =  ViTConfig(
         image_size=self.image_size   ,							patch_size=self.patch_size   ,							num_channels=self.num_channels   ,							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   ,							is_decoder=_lowercase   ,							initializer_range=self.initializer_range   ,							)

     return config, pixel_values







  def   UpperCAmelCase					(      self  :int   ,							_lowercase  :int   ,							_lowercase  :Optional[Any] ):






     '''simple docstring'''





     lowercase__			       =  FlaxViTModel(config=_lowercase )
     lowercase__			       =  model(_lowercase )
     # expected sequence length = num_patches + 1 (we add 1 for the [CLS] token)
     lowercase__			       =  (self.image_size, self.image_size)
     lowercase__			       =  (self.patch_size, self.patch_size)
     lowercase__			       =  (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
     self.parent.assertEqual(result.last_hidden_state.shape   ,							(self.batch_size, num_patches + 1, self.hidden_size) )







  def   UpperCAmelCase					(      self  :int   ,							_lowercase  :Any   ,							_lowercase  :int ):






     '''simple docstring'''





     lowercase__			       =  self.type_sequence_label_size
     lowercase__			       =  FlaxViTForImageClassification(config=_lowercase )
     lowercase__			       =  model(_lowercase )
     self.parent.assertEqual(result.logits.shape   ,							(self.batch_size, self.type_sequence_label_size) )

     # test greyscale images
     lowercase__			       =  1
     lowercase__			       =  FlaxViTForImageClassification(_lowercase )

     lowercase__			       =  floats_tensor([self.batch_size, 1, self.image_size, self.image_size] )
     lowercase__			       =  model(_lowercase )







  def   UpperCAmelCase					(      self  :str ):






     '''simple docstring'''





     lowercase__			       =  self.prepare_config_and_inputs()
     (
         (
         lowercase__
     )       ,							(
         lowercase__
     )       ,							
     )			       =  config_and_inputs
     lowercase__			       =  {"pixel_values": pixel_values}
     return config, inputs_dict





@require_flax
class  lowerCAmelCase							(       lowercase_							,							unittest.TestCase ):
  __lowerCamelCase		=      (FlaxViTModel, FlaxViTForImageClassification) if is_flax_available() else ()







  def   UpperCAmelCase					(      self  :int ):






     '''simple docstring'''





     lowercase__			       =  FlaxViTModelTester(self )
     lowercase__			       =  ConfigTester(self   ,							config_class=_lowercase   ,							has_text_modality=_lowercase   ,							hidden_size=37 )







  def   UpperCAmelCase					(      self  :Any ):






     '''simple docstring'''





     self.config_tester.run_common_tests()







  def   UpperCAmelCase					(      self  :str ):






     '''simple docstring'''





     lowercase__			       =  self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*_lowercase )







  def   UpperCAmelCase					(      self  :Optional[int] ):






     '''simple docstring'''





     lowercase__			       =  self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_image_classification(*_lowercase )







  def   UpperCAmelCase					(      self  :int ):






     '''simple docstring'''





     lowercase__       ,							lowercase__			       =  self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
        lowercase__			       =  model_class(_lowercase )
        lowercase__			       =  inspect.signature(model.__call__ )
        # signature.parameters is an OrderedDict => so arg_names order is deterministic
        lowercase__			       =  [*signature.parameters.keys()]

        lowercase__			       =  ["pixel_values"]
        self.assertListEqual(arg_names[:1]   ,							_lowercase )







  def   UpperCAmelCase					(      self  :Optional[Any] ):






     '''simple docstring'''





     lowercase__       ,							lowercase__			       =  self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
        with self.subTest(model_class.__name__ ):
           lowercase__			       =  self._prepare_for_class(_lowercase   ,							_lowercase )
           lowercase__			       =  model_class(_lowercase )

           @jax.jit
           def model_jitted(_lowercase  :str   ,							**_lowercase  :Optional[int] ):
              return model(pixel_values=_lowercase   ,							**_lowercase )

           with self.subTest("JIT Enabled" ):
              lowercase__			       =  model_jitted(**_lowercase ).to_tuple()

           with self.subTest("JIT Disabled" ):
              with jax.disable_jit():
                 lowercase__			       =  model_jitted(**_lowercase ).to_tuple()

           self.assertEqual(len(_lowercase )   ,							len(_lowercase ) )
           for jitted_output, output in zip(_lowercase   ,							_lowercase ):
              self.assertEqual(jitted_output.shape   ,							output.shape )







  @slow
  def   UpperCAmelCase					(      self  :Any ):






     '''simple docstring'''





     for model_class_name in self.all_model_classes:
        lowercase__			       =  model_class_name.from_pretrained("google/vit-base-patch16-224" )
        lowercase__			       =  model(np.ones((1, 3, 2_24, 2_24) ) )
        self.assertIsNotNone(_lowercase )













from __future__ import annotations

from collections import deque





class  lowerCAmelCase							:







  def __init__(      self  :List[Any]   ,							_lowercase  :list[str] ):






     '''simple docstring'''





     lowercase__			       =  []
     self.adlist.append(
         {"value": "", "next_states": [], "fail_state": 0, "output": []} )

     for keyword in keywords:
        self.add_keyword(_lowercase )
     self.set_fail_transitions()







  def   UpperCAmelCase					(      self  :str   ,							_lowercase  :int   ,							_lowercase  :str ):






     '''simple docstring'''





     for state in self.adlist[current_state]["next_states"]:
        if char == self.adlist[state]["value"]:
           return state
     return None







  def   UpperCAmelCase					(      self  :List[str]   ,							_lowercase  :str ):






     '''simple docstring'''





     lowercase__			       =  0
     for character in keyword:
        lowercase__			       =  self.find_next_state(_lowercase   ,							_lowercase )
        if next_state is None:
           self.adlist.append(
               {
                   "value": character,
                   "next_states": [],
                   "fail_state": 0,
                   "output": [],
               } )
           self.adlist[current_state]["next_states"].append(len(self.adlist ) - 1 )
           lowercase__			       =  len(self.adlist ) - 1
        else:
           lowercase__			       =  next_state
     self.adlist[current_state]["output"].append(_lowercase )







  def   UpperCAmelCase					(      self  :int ):






     '''simple docstring'''





     lowercase__			       =  deque()
     for node in self.adlist[0]["next_states"]:
        q.append(_lowercase )
        lowercase__			       =  0
     while q:
        lowercase__			       =  q.popleft()
        for child in self.adlist[r]["next_states"]:
           q.append(_lowercase )
           lowercase__			       =  self.adlist[r]["fail_state"]
           while (
               self.find_next_state(_lowercase   ,							self.adlist[child]["value"] ) is None
               and state != 0
           ):
              lowercase__			       =  self.adlist[state]["fail_state"]
           lowercase__			       =  self.find_next_state(
               _lowercase   ,							self.adlist[child]["value"] )
           if self.adlist[child]["fail_state"] is None:
              lowercase__			       =  0
           lowercase__			       =  (
               self.adlist[child]["output"]
               + self.adlist[self.adlist[child]["fail_state"]]["output"]
           )







  def   UpperCAmelCase					(      self  :Optional[Any]   ,							_lowercase  :str ):






     '''simple docstring'''





     lowercase__			       =  {}  # returns a dict with keywords and list of its occurrences
     lowercase__			       =  0
     for i in range(len(_lowercase ) ):
        while (
            self.find_next_state(_lowercase   ,							string[i] ) is None
            and current_state != 0
        ):
           lowercase__			       =  self.adlist[current_state]["fail_state"]
        lowercase__			       =  self.find_next_state(_lowercase   ,							string[i] )
        if next_state is None:
           lowercase__			       =  0
        else:
           lowercase__			       =  next_state
           for key in self.adlist[current_state]["output"]:
              if key not in result:
                 lowercase__			       =  []
              result[key].append(i - len(_lowercase ) + 1 )
     return result


if __name__ == "__main__":
  import doctest

  doctest.testmod()






