





"""simple docstring"""



from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available


a_ =			{"""configuration_yolos""": ["""YOLOS_PRETRAINED_CONFIG_ARCHIVE_MAP""", """YolosConfig""", """YolosOnnxConfig"""]}

try:
       if not is_vision_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       a_ =			["""YolosFeatureExtractor"""]
       a_ =			["""YolosImageProcessor"""]

try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       a_ =			[
           """YOLOS_PRETRAINED_MODEL_ARCHIVE_LIST""",
           """YolosForObjectDetection""",
           """YolosModel""",
           """YolosPreTrainedModel""",
       ]


if TYPE_CHECKING:
       from .configuration_yolos import YOLOS_PRETRAINED_CONFIG_ARCHIVE_MAP, YolosConfig, YolosOnnxConfig

       try:
              if not is_vision_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .feature_extraction_yolos import YolosFeatureExtractor
              from .image_processing_yolos import YolosImageProcessor

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_yolos import (
                  YOLOS_PRETRAINED_MODEL_ARCHIVE_LIST,
                  YolosForObjectDetection,
                  YolosModel,
                  YolosPreTrainedModel,
              )


else:
       import sys

       a_ =			_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)











"""simple docstring"""



import random
import unittest

import numpy as np

import transformers
from transformers import is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax


if is_flax_available():
       import os

       import jax.numpy as jnp
       from jax import jit

       from transformers import AutoTokenizer, FlaxAutoModelForCausalLM
       from transformers.modeling_flax_pytorch_utils import load_flax_weights_in_pytorch_model

       a_ =			"""0.12"""  # assumed parallelism: 8


if is_torch_available():
       import torch







def 		__lowercase     (   snake_case_  :  Any				,snake_case_  :  int				,snake_case_  :  List[str]=None						)   ->Union[str, Any]:



    '''simple docstring'''







    if rng is None:
        __A    :						Union[str, Any]						 =     random.Random()

    __A    :						List[str]						 =     1
    for dim in shape:
        total_dims *= dim

    __A    :						List[str]						 =     []
    for _ in range(snake_case_						):
        values.append(rng.randint(0				,vocab_size - 1						)						)

    __A    :						Optional[Any]						 =     np.array(snake_case_				,dtype=jnp.intaa						).reshape(snake_case_						)

    return output







def 		__lowercase     (   snake_case_  :  int				,snake_case_  :  Union[str, Any]=None						)   ->int:



    '''simple docstring'''







    __A    :						int						 =     ids_tensor(snake_case_				,vocab_size=2				,rng=snake_case_						)
    # make sure that at least one token is attended to for each batch
    __A    :						str						 =     1
    return attn_mask






@require_flax
class 		__snake_case					:





 """simple docstring"""


 _lowerCamelCase        =     None
 _lowerCamelCase        =     ()

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A ,       __A    :						Optional[int]						 =     self.model_tester.prepare_config_and_inputs_for_common()

     # cut to half length & take max batch_size 3
     __A    :						Optional[int]						 =     2
     __A    :						Tuple						 =     inputs['''input_ids'''].shape[-1] // 2
     __A    :						List[Any]						 =     inputs['''input_ids'''][:max_batch_size, :sequence_length]

     __A    :						Optional[Any]						 =     jnp.ones_like(__lowerCamelCase  )
     __A    :						List[Any]						 =     attention_mask[:max_batch_size, :sequence_length]

     # generate max 5 tokens
     __A    :						Optional[int]						 =     input_ids.shape[-1] + 5
     if config.eos_token_id is not None and config.pad_token_id is None:
         # hack to allow generate for models such as GPT2 as is done in `generate()`
         __A    :						Optional[int]						 =     config.eos_token_id
     return config, input_ids, attention_mask, max_length

 @is_pt_flax_cross_test
 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A ,       __A ,       __A ,       __A    :						List[Any]						 =     self._get_input_ids_and_config()
     __A    :						List[Any]						 =     False
     __A    :						str						 =     max_length
     __A    :						Tuple						 =     0

     for model_class in self.all_generative_model_classes:
         __A    :						Union[str, Any]						 =     model_class(__lowerCamelCase  )

         __A    :						Any						 =     model_class.__name__[4:]  # Skip the "Flax" at the beginning
         __A    :						Dict						 =     getattr(__lowerCamelCase					,						__lowerCamelCase  )
         __A    :						Tuple						 =     pt_model_class(__lowerCamelCase  ).eval()
         __A    :						Optional[Any]						 =     load_flax_weights_in_pytorch_model(__lowerCamelCase					,						flax_model.params  )

         __A    :						str						 =     flax_model.generate(__lowerCamelCase  ).sequences
         __A    :						List[Any]						 =     pt_model.generate(torch.tensor(__lowerCamelCase					,						dtype=torch.long  )  )

         if flax_generation_outputs.shape[-1] > pt_generation_outputs.shape[-1]:
             __A    :						Any						 =     flax_generation_outputs[:, : pt_generation_outputs.shape[-1]]

         self.assertListEqual(pt_generation_outputs.numpy().tolist()					,						flax_generation_outputs.tolist()  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A ,       __A ,       __A ,       __A    :						Union[str, Any]						 =     self._get_input_ids_and_config()
     __A    :						List[str]						 =     False
     __A    :						Optional[Any]						 =     max_length

     for model_class in self.all_generative_model_classes:
         __A    :						Optional[int]						 =     model_class(__lowerCamelCase  )

         __A    :						Dict						 =     model.generate(__lowerCamelCase  ).sequences
         self.assertEqual(generation_outputs.shape[-1]					,						__lowerCamelCase  )

         __A    :						Union[str, Any]						 =     jit(model.generate  )
         __A    :						str						 =     jit_generate(__lowerCamelCase  ).sequences

         self.assertListEqual(generation_outputs.tolist()					,						jit_generation_outputs.tolist()  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A ,       __A ,       __A ,       __A    :						int						 =     self._get_input_ids_and_config()
     __A    :						List[Any]						 =     True
     __A    :						Union[str, Any]						 =     max_length

     for model_class in self.all_generative_model_classes:
         __A    :						Dict						 =     model_class(__lowerCamelCase  )

         __A    :						List[Any]						 =     model.generate(__lowerCamelCase  ).sequences
         self.assertEqual(generation_outputs.shape[-1]					,						__lowerCamelCase  )

         __A    :						List[Any]						 =     jit(model.generate  )
         __A    :						List[str]						 =     jit_generate(__lowerCamelCase  ).sequences

         self.assertListEqual(generation_outputs.tolist()					,						jit_generation_outputs.tolist()  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A ,       __A ,       __A ,       __A    :						List[Any]						 =     self._get_input_ids_and_config()
     __A    :						List[str]						 =     False
     __A    :						Any						 =     max_length
     __A    :						List[Any]						 =     2

     for model_class in self.all_generative_model_classes:
         __A    :						Tuple						 =     model_class(__lowerCamelCase  )

         __A    :						Optional[Any]						 =     model.generate(__lowerCamelCase  ).sequences
         self.assertEqual(generation_outputs.shape[-1]					,						__lowerCamelCase  )

         __A    :						Union[str, Any]						 =     jit(model.generate  )
         __A    :						Optional[int]						 =     jit_generate(__lowerCamelCase  ).sequences

         self.assertListEqual(generation_outputs.tolist()					,						jit_generation_outputs.tolist()  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A ,       __A ,       __A ,       __A    :						Tuple						 =     self._get_input_ids_and_config()
     __A    :						str						 =     False
     __A    :						Union[str, Any]						 =     max_length
     __A    :						List[Any]						 =     2
     __A    :						Any						 =     2

     for model_class in self.all_generative_model_classes:
         __A    :						List[str]						 =     model_class(__lowerCamelCase  )

         __A    :						Optional[int]						 =     model.generate(__lowerCamelCase  ).sequences
         self.assertEqual(generation_outputs.shape[0]					,						input_ids.shape[0] * config.num_return_sequences  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A ,       __A ,       __A ,       __A    :						Any						 =     self._get_input_ids_and_config()
     __A    :						Optional[Any]						 =     True
     __A    :						Union[str, Any]						 =     max_length
     __A    :						List[str]						 =     0.8
     __A    :						List[str]						 =     10
     __A    :						Union[str, Any]						 =     0.3
     __A    :						Union[str, Any]						 =     1
     __A    :						Optional[Any]						 =     8
     __A    :						Dict						 =     9

     for model_class in self.all_generative_model_classes:
         __A    :						List[Any]						 =     model_class(__lowerCamelCase  )

         __A    :						List[str]						 =     model.generate(__lowerCamelCase  ).sequences
         self.assertEqual(generation_outputs.shape[-1]					,						__lowerCamelCase  )

         __A    :						Any						 =     jit(model.generate  )
         __A    :						str						 =     jit_generate(__lowerCamelCase  ).sequences

         self.assertListEqual(generation_outputs.tolist()					,						jit_generation_outputs.tolist()  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A ,       __A ,       __A ,       __A    :						Union[str, Any]						 =     self._get_input_ids_and_config()
     __A    :						Union[str, Any]						 =     max_length
     __A    :						List[str]						 =     1
     __A    :						str						 =     8
     __A    :						Any						 =     9

     for model_class in self.all_generative_model_classes:
         __A    :						Union[str, Any]						 =     model_class(__lowerCamelCase  )

         __A    :						Tuple						 =     model.generate(__lowerCamelCase  ).sequences
         self.assertEqual(generation_outputs.shape[-1]					,						__lowerCamelCase  )

         __A    :						Optional[Any]						 =     jit(model.generate  )
         __A    :						Optional[Any]						 =     jit_generate(__lowerCamelCase  ).sequences

         self.assertListEqual(generation_outputs.tolist()					,						jit_generation_outputs.tolist()  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A ,       __A ,       __A ,       __A    :						int						 =     self._get_input_ids_and_config()
     __A    :						Optional[int]						 =     max_length
     __A    :						List[str]						 =     2
     __A    :						List[Any]						 =     1
     __A    :						Optional[Any]						 =     8
     __A    :						str						 =     9

     for model_class in self.all_generative_model_classes:
         __A    :						Optional[int]						 =     model_class(__lowerCamelCase  )

         __A    :						int						 =     model.generate(__lowerCamelCase  ).sequences
         self.assertEqual(generation_outputs.shape[-1]					,						__lowerCamelCase  )

         __A    :						str						 =     jit(model.generate  )
         __A    :						Any						 =     jit_generate(__lowerCamelCase  ).sequences

         self.assertListEqual(generation_outputs.tolist()					,						jit_generation_outputs.tolist()  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A ,       __A ,       __A ,       __A    :						int						 =     self._get_input_ids_and_config()

     # pad attention mask on the left
     __A    :						Dict						 =     attention_mask.at[(0, 0)].set(0  )

     __A    :						Optional[Any]						 =     False
     __A    :						Union[str, Any]						 =     max_length

     for model_class in self.all_generative_model_classes:
         __A    :						int						 =     model_class(__lowerCamelCase  )

         __A    :						Union[str, Any]						 =     model.generate(__lowerCamelCase					,						attention_mask=__lowerCamelCase  ).sequences
         self.assertEqual(generation_outputs.shape[-1]					,						__lowerCamelCase  )

         __A    :						Optional[int]						 =     jit(model.generate  )
         __A    :						Any						 =     jit_generate(__lowerCamelCase					,						attention_mask=__lowerCamelCase  ).sequences

         self.assertListEqual(generation_outputs.tolist()					,						jit_generation_outputs.tolist()  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A ,       __A ,       __A ,       __A    :						Optional[int]						 =     self._get_input_ids_and_config()

     # pad attention mask on the left
     __A    :						str						 =     attention_mask.at[(0, 0)].set(0  )

     __A    :						List[Any]						 =     True
     __A    :						Any						 =     max_length

     for model_class in self.all_generative_model_classes:
         __A    :						str						 =     model_class(__lowerCamelCase  )

         __A    :						List[Any]						 =     model.generate(__lowerCamelCase					,						attention_mask=__lowerCamelCase  ).sequences
         self.assertEqual(generation_outputs.shape[-1]					,						__lowerCamelCase  )

         __A    :						List[str]						 =     jit(model.generate  )
         __A    :						Optional[int]						 =     jit_generate(__lowerCamelCase					,						attention_mask=__lowerCamelCase  ).sequences

         self.assertListEqual(generation_outputs.tolist()					,						jit_generation_outputs.tolist()  )







 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A ,       __A ,       __A ,       __A    :						Optional[Any]						 =     self._get_input_ids_and_config()

     # pad attention mask on the left
     __A    :						List[str]						 =     attention_mask.at[(0, 0)].set(0  )

     __A    :						Optional[int]						 =     2
     __A    :						Dict						 =     max_length

     for model_class in self.all_generative_model_classes:
         __A    :						Any						 =     model_class(__lowerCamelCase  )

         __A    :						int						 =     model.generate(__lowerCamelCase					,						attention_mask=__lowerCamelCase  ).sequences
         self.assertEqual(generation_outputs.shape[-1]					,						__lowerCamelCase  )

         __A    :						str						 =     jit(model.generate  )
         __A    :						str						 =     jit_generate(__lowerCamelCase					,						attention_mask=__lowerCamelCase  ).sequences

         self.assertListEqual(generation_outputs.tolist()					,						jit_generation_outputs.tolist()  )






@require_flax
class 		__snake_case					(		unittest.TestCase     ):





 """simple docstring"""


 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A    :						Optional[int]						 =     AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-bert'''  )
     __A    :						Optional[Any]						 =     FlaxAutoModelForCausalLM.from_pretrained('''hf-internal-testing/tiny-bert-flax-only'''  )

     __A    :						Any						 =     '''Hello world'''
     __A    :						Dict						 =     tokenizer(__lowerCamelCase					,						return_tensors='''np'''  ).input_ids

     # typos are quickly detected (the correct argument is `do_sample`)
     with self.assertRaisesRegex(__lowerCamelCase					,						'''do_samples'''  ):
         model.generate(__lowerCamelCase					,						do_samples=__lowerCamelCase  )

     # arbitrary arguments that will not be used anywhere are also not accepted
     with self.assertRaisesRegex(__lowerCamelCase					,						'''foo'''  ):
         __A    :						Any						 =     {'''foo''': '''bar'''}
         model.generate(__lowerCamelCase					,						**__lowerCamelCase  )




