





import inspect
import unittest

from transformers import ConvNextConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
      import torch

      from transformers import ConvNextBackbone, ConvNextForImageClassification, ConvNextModel
      from transformers.models.convnext.modeling_convnext import CONVNEXT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
      from PIL import Image

      from transformers import AutoImageProcessor

class 		UpperCamelCase__			:



     '''simple docstring'''

     def __init__(							self						:   Any						,lowerCamelCase__						:   Dict						,lowerCamelCase__						:   List[Any]=13						,lowerCamelCase__						:   Dict=32						,lowerCamelCase__						:   Optional[Any]=3						,lowerCamelCase__						:   Optional[int]=4						,lowerCamelCase__						:   List[Any]=[10, 20, 30, 40]						,lowerCamelCase__						:   Tuple=[2, 2, 3, 2]						,lowerCamelCase__						:   int=True						,lowerCamelCase__						:   List[str]=True						,lowerCamelCase__						:   Optional[Any]=37						,lowerCamelCase__						:   List[str]="gelu"						,lowerCamelCase__						:   Optional[Any]=10						,lowerCamelCase__						:   int=0.02						,lowerCamelCase__						:   int=["stage2", "stage3", "stage4"]						,lowerCamelCase__						:   Optional[int]=[2, 3, 4]						,lowerCamelCase__						:   int=None						,)					->		Optional[Any]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						parent
           SCREAMING_SNAKE_CASE		=						batch_size
           SCREAMING_SNAKE_CASE		=						image_size
           SCREAMING_SNAKE_CASE		=						num_channels
           SCREAMING_SNAKE_CASE		=						num_stages
           SCREAMING_SNAKE_CASE		=						hidden_sizes
           SCREAMING_SNAKE_CASE		=						depths
           SCREAMING_SNAKE_CASE		=						is_training
           SCREAMING_SNAKE_CASE		=						use_labels
           SCREAMING_SNAKE_CASE		=						intermediate_size
           SCREAMING_SNAKE_CASE		=						hidden_act
           SCREAMING_SNAKE_CASE		=						num_labels
           SCREAMING_SNAKE_CASE		=						initializer_range
           SCREAMING_SNAKE_CASE		=						out_features
           SCREAMING_SNAKE_CASE		=						out_indices
           SCREAMING_SNAKE_CASE		=						scope






     def     SCREAMING_SNAKE_CASE__  (							self						:   Optional[Any]	)					->		str:




           '''simple docstring'''


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

           SCREAMING_SNAKE_CASE		=						None
           if self.use_labels:
                 SCREAMING_SNAKE_CASE		=						ids_tensor([self.batch_size]						,self.num_labels	)

           SCREAMING_SNAKE_CASE		=						self.get_config()
           return config, pixel_values, labels






     def     SCREAMING_SNAKE_CASE__  (							self						:   Optional[Any]	)					->		Dict:




           '''simple docstring'''


           return ConvNextConfig(
               num_channels=self.num_channels						,hidden_sizes=self.hidden_sizes						,depths=self.depths						,num_stages=self.num_stages						,hidden_act=self.hidden_act						,is_decoder=lowerCamelCase__						,initializer_range=self.initializer_range						,out_features=self.out_features						,out_indices=self.out_indices						,num_labels=self.num_labels						,)






     def     SCREAMING_SNAKE_CASE__  (							self						:   Optional[Any]						,lowerCamelCase__						:   Optional[Any]						,lowerCamelCase__						:   Optional[Any]						,lowerCamelCase__						:   List[str]	)					->		Any:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						ConvNextModel(config=lowerCamelCase__	)
           model.to(lowerCamelCase__	)
           model.eval()
           SCREAMING_SNAKE_CASE		=						model(lowerCamelCase__	)
           # expected last hidden states: B, C, H // 32, W // 32
           self.parent.assertEqual(
               result.last_hidden_state.shape						,(self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32)						,)






     def     SCREAMING_SNAKE_CASE__  (							self						:   Any						,lowerCamelCase__						:   str						,lowerCamelCase__						:   List[Any]						,lowerCamelCase__						:   Dict	)					->		Dict:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						ConvNextForImageClassification(lowerCamelCase__	)
           model.to(lowerCamelCase__	)
           model.eval()
           SCREAMING_SNAKE_CASE		=						model(lowerCamelCase__						,labels=lowerCamelCase__	)
           self.parent.assertEqual(result.logits.shape						,(self.batch_size, self.num_labels)	)






     def     SCREAMING_SNAKE_CASE__  (							self						:   Optional[int]						,lowerCamelCase__						:   int						,lowerCamelCase__						:   List[Any]						,lowerCamelCase__						:   Optional[int]	)					->		Optional[int]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						ConvNextBackbone(config=lowerCamelCase__	)
           model.to(lowerCamelCase__	)
           model.eval()
           SCREAMING_SNAKE_CASE		=						model(lowerCamelCase__	)

           # verify hidden states
           self.parent.assertEqual(len(result.feature_maps	)						,len(config.out_features	)	)
           self.parent.assertListEqual(list(result.feature_maps[0].shape	)						,[self.batch_size, self.hidden_sizes[1], 4, 4]	)

           # verify channels
           self.parent.assertEqual(len(model.channels	)						,len(config.out_features	)	)
           self.parent.assertListEqual(model.channels						,config.hidden_sizes[1:]	)

           # verify backbone works with out_features=None
           SCREAMING_SNAKE_CASE		=						None
           SCREAMING_SNAKE_CASE		=						ConvNextBackbone(config=lowerCamelCase__	)
           model.to(lowerCamelCase__	)
           model.eval()
           SCREAMING_SNAKE_CASE		=						model(lowerCamelCase__	)

           # verify feature maps
           self.parent.assertEqual(len(result.feature_maps	)						,1	)
           self.parent.assertListEqual(list(result.feature_maps[0].shape	)						,[self.batch_size, self.hidden_sizes[-1], 1, 1]	)

           # verify channels
           self.parent.assertEqual(len(model.channels	)						,1	)
           self.parent.assertListEqual(model.channels						,[config.hidden_sizes[-1]]	)






     def     SCREAMING_SNAKE_CASE__  (							self						:   Any	)					->		int:




           '''simple docstring'''


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

@require_torch
class 		UpperCamelCase__			(				lowerCAmelCase_							,				lowerCAmelCase_							,				unittest.TestCase       ):



     '''simple docstring'''

     __snake_case   :     Optional[int]		     =      (
         (
             ConvNextModel,
             ConvNextForImageClassification,
             ConvNextBackbone,
         )
         if is_torch_available()
         else ()
     )
     __snake_case   :     Optional[Any]		     =      (
         {"feature-extraction": ConvNextModel, "image-classification": ConvNextForImageClassification}
         if is_torch_available()
         else {}
     )

     __snake_case   :     Any		     =      True
     __snake_case   :     Optional[Any]		     =      False
     __snake_case   :     Union[str, Any]		     =      False
     __snake_case   :     Optional[int]		     =      False
     __snake_case   :     Union[str, Any]		     =      False






     def     SCREAMING_SNAKE_CASE__  (							self						:   str	)					->		Tuple:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						ConvNextModelTester(self	)
           SCREAMING_SNAKE_CASE		=						ConfigTester(self						,config_class=lowerCamelCase__						,has_text_modality=lowerCamelCase__						,hidden_size=37	)






     def     SCREAMING_SNAKE_CASE__  (							self						:   List[Any]	)					->		Tuple:




           '''simple docstring'''


           self.create_and_test_config_common_properties()
           self.config_tester.create_and_test_config_to_json_string()
           self.config_tester.create_and_test_config_to_json_file()
           self.config_tester.create_and_test_config_from_and_save_pretrained()
           self.config_tester.create_and_test_config_with_num_labels()
           self.config_tester.check_config_can_be_init_without_params()
           self.config_tester.check_config_arguments_init()






     def     SCREAMING_SNAKE_CASE__  (							self						:   Dict	)					->		List[Any]:




           '''simple docstring'''


           return






     @unittest.skip(reason="""ConvNext does not use inputs_embeds"""	)
     def     SCREAMING_SNAKE_CASE__  (							self						:   Optional[Any]	)					->		Optional[int]:




           '''simple docstring'''


           pass






     @unittest.skip(reason="""ConvNext does not support input and output embeddings"""	)
     def     SCREAMING_SNAKE_CASE__  (							self						:   Dict	)					->		Optional[Any]:




           '''simple docstring'''


           pass






     @unittest.skip(reason="""ConvNext does not use feedforward chunking"""	)
     def     SCREAMING_SNAKE_CASE__  (							self						:   Tuple	)					->		List[Any]:




           '''simple docstring'''


           pass






     def     SCREAMING_SNAKE_CASE__  (							self						:   Any	)					->		List[Any]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE, SCREAMING_SNAKE_CASE		=						self.model_tester.prepare_config_and_inputs_for_common()

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

                 SCREAMING_SNAKE_CASE		=						["""pixel_values"""]
                 self.assertListEqual(arg_names[:1]						,lowerCamelCase__	)






     def     SCREAMING_SNAKE_CASE__  (							self						:   List[Any]	)					->		Dict:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*lowerCamelCase__	)






     def     SCREAMING_SNAKE_CASE__  (							self						:   Optional[Any]	)					->		Optional[int]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_backbone(*lowerCamelCase__	)






     def     SCREAMING_SNAKE_CASE__  (							self						:   Dict	)					->		Optional[Any]:




           '''simple docstring'''


           def check_hidden_states_output(lowerCamelCase__						:   Optional[Any]						,lowerCamelCase__						:   Tuple						,lowerCamelCase__						:   int	):
                 SCREAMING_SNAKE_CASE		=						model_class(lowerCamelCase__	)
                 model.to(lowerCamelCase__	)
                 model.eval()

                 with torch.no_grad():
                       SCREAMING_SNAKE_CASE		=						model(**self._prepare_for_class(lowerCamelCase__						,lowerCamelCase__	)	)

                 SCREAMING_SNAKE_CASE		=						outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

                 SCREAMING_SNAKE_CASE		=						self.model_tester.num_stages
                 self.assertEqual(len(lowerCamelCase__	)						,expected_num_stages + 1	)

                 # ConvNext's feature maps are of shape (batch_size, num_channels, height, width)
                 self.assertListEqual(
                     list(hidden_states[0].shape[-2:]	)						,[self.model_tester.image_size // 4, self.model_tester.image_size // 4]						,)

           SCREAMING_SNAKE_CASE, SCREAMING_SNAKE_CASE		=						self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
                 SCREAMING_SNAKE_CASE		=						True
                 check_hidden_states_output(lowerCamelCase__						,lowerCamelCase__						,lowerCamelCase__	)

                 # check that output_hidden_states also work using config
                 del inputs_dict["output_hidden_states"]
                 SCREAMING_SNAKE_CASE		=						True

                 check_hidden_states_output(lowerCamelCase__						,lowerCamelCase__						,lowerCamelCase__	)






     def     SCREAMING_SNAKE_CASE__  (							self						:   Tuple	)					->		Union[str, Any]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_image_classification(*lowerCamelCase__	)






     @slow
     def     SCREAMING_SNAKE_CASE__  (							self						:   Any	)					->		Dict:




           '''simple docstring'''


           for model_name in CONVNEXT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 SCREAMING_SNAKE_CASE		=						ConvNextModel.from_pretrained(lowerCamelCase__	)
                 self.assertIsNotNone(lowerCamelCase__	)



def        __lowercase       (				)			->       List[str]:


      '''simple docstring'''


      SCREAMING_SNAKE_CASE		=						Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""     )
      return image

@require_torch
@require_vision
class 		UpperCamelCase__			(				unittest.TestCase       ):



     '''simple docstring'''

     @cached_property
     def     SCREAMING_SNAKE_CASE__  (							self						:   int	)					->		str:




           '''simple docstring'''


           return AutoImageProcessor.from_pretrained("""facebook/convnext-tiny-224"""	) if is_vision_available() else None






     @slow
     def     SCREAMING_SNAKE_CASE__  (							self						:   Union[str, Any]	)					->		str:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						ConvNextForImageClassification.from_pretrained("""facebook/convnext-tiny-224"""	).to(lowerCamelCase__	)

           SCREAMING_SNAKE_CASE		=						self.default_image_processor
           SCREAMING_SNAKE_CASE		=						prepare_img()
           SCREAMING_SNAKE_CASE		=						image_processor(images=lowerCamelCase__						,return_tensors="""pt"""	).to(lowerCamelCase__	)

           # forward pass
           with torch.no_grad():
                 SCREAMING_SNAKE_CASE		=						model(**lowerCamelCase__	)

           # verify the logits
           SCREAMING_SNAKE_CASE		=						torch.Size((1, 1000)	)
           self.assertEqual(outputs.logits.shape						,lowerCamelCase__	)

           SCREAMING_SNAKE_CASE		=						torch.tensor([-0.0260, -0.4739, 0.1911]	).to(lowerCamelCase__	)

           self.assertTrue(torch.allclose(outputs.logits[0, :3]						,lowerCamelCase__						,atol=1e-4	)	)

@require_torch
class 		UpperCamelCase__			(				unittest.TestCase							,				lowerCAmelCase_       ):



     '''simple docstring'''

     __snake_case   :     Union[str, Any]		     =      (ConvNextBackbone,) if is_torch_available() else ()
     __snake_case   :     str		     =      ConvNextConfig

     __snake_case   :     int		     =      False






     def     SCREAMING_SNAKE_CASE__  (							self						:   List[str]	)					->		Union[str, Any]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						ConvNextModelTester(self	)









import logging
import random

import ray

from transformers import RagConfig, RagRetriever, RagTokenizer
from transformers.models.rag.retrieval_rag import CustomHFIndex


SCREAMING_SNAKE_CASE_						       =	logging.getLogger(__name__)

class 		UpperCamelCase__			:



     '''simple docstring'''

     def __init__(							self						:   str	)					->		List[str]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						False






     def     SCREAMING_SNAKE_CASE__  (							self						:   Any						,lowerCamelCase__						:   int						,lowerCamelCase__						:   Dict						,lowerCamelCase__						:   List[Any]						,lowerCamelCase__						:   Tuple	)					->		Optional[Any]:




           '''simple docstring'''


           if not self.initialized:
                 SCREAMING_SNAKE_CASE		=						RagRetriever(
                     lowerCamelCase__						,question_encoder_tokenizer=lowerCamelCase__						,generator_tokenizer=lowerCamelCase__						,index=lowerCamelCase__						,init_retrieval=lowerCamelCase__						,)
                 SCREAMING_SNAKE_CASE		=						True






     def     SCREAMING_SNAKE_CASE__  (							self						:   Dict	)					->		Any:




           '''simple docstring'''


           self.retriever.index.init_index()






     def     SCREAMING_SNAKE_CASE__  (							self						:   Any						,lowerCamelCase__						:   List[Any]						,lowerCamelCase__						:   List[str]	)					->		List[Any]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE, SCREAMING_SNAKE_CASE		=						self.retriever._main_retrieve(lowerCamelCase__						,lowerCamelCase__	)
           return doc_ids, retrieved_doc_embeds

class 		UpperCamelCase__			(				lowerCAmelCase_       ):



     '''simple docstring'''

     def __init__(							self						:   int						,lowerCamelCase__						:   Dict						,lowerCamelCase__						:   Optional[Any]						,lowerCamelCase__						:   List[str]						,lowerCamelCase__						:   Optional[Any]						,lowerCamelCase__						:   Dict=None	)					->		Any:




           '''simple docstring'''


           if index is not None and index.is_initialized() and len(lowerCamelCase__	) > 0:
                 raise ValueError(
                     """When using Ray for distributed fine-tuning, """
                     """you'll need to provide the paths instead, """
                     """as the dataset and the index are loaded """
                     """separately. More info in examples/rag/use_own_knowledge_dataset.py """	)
           super().__init__(
               lowerCamelCase__						,question_encoder_tokenizer=lowerCamelCase__						,generator_tokenizer=lowerCamelCase__						,index=lowerCamelCase__						,init_retrieval=lowerCamelCase__						,)
           SCREAMING_SNAKE_CASE		=						retrieval_workers
           if len(self.retrieval_workers	) > 0:
                 ray.get(
                     [
                         worker.create_rag_retriever.remote(lowerCamelCase__						,lowerCamelCase__						,lowerCamelCase__						,lowerCamelCase__	)
                         for worker in self.retrieval_workers
                     ]	)






     def     SCREAMING_SNAKE_CASE__  (							self						:   List[str]	)					->		Union[str, Any]:




           '''simple docstring'''


           logger.info("""initializing retrieval"""	)

           if len(self.retrieval_workers	) > 0:
                 ray.get([worker.init_retrieval.remote() for worker in self.retrieval_workers]	)
           else:
                 # Non-distributed training. Load index into this same process.
                 self.index.init_index()






     def     SCREAMING_SNAKE_CASE__  (							self						:   Any						,lowerCamelCase__						:   Any						,lowerCamelCase__						:   int	)					->		Dict:




           '''simple docstring'''


           if len(self.retrieval_workers	) > 0:
                 # Select a random retrieval actor.
                 SCREAMING_SNAKE_CASE		=						self.retrieval_workers[random.randint(0						,len(self.retrieval_workers	) - 1	)]
                 SCREAMING_SNAKE_CASE, SCREAMING_SNAKE_CASE		=						ray.get(random_worker.retrieve.remote(lowerCamelCase__						,lowerCamelCase__	)	)
           else:
                 SCREAMING_SNAKE_CASE, SCREAMING_SNAKE_CASE		=						self._main_retrieve(lowerCamelCase__						,lowerCamelCase__	)
           return retrieved_doc_embeds, doc_ids, self.index.get_doc_dicts(lowerCamelCase__	)






     @classmethod
     def     SCREAMING_SNAKE_CASE__  (							cls						:   Optional[Any]						,lowerCamelCase__						:   List[Any]						,lowerCamelCase__						:   Union[str, Any]=None						,**lowerCamelCase__						:   Optional[Any]	)					->		Any:




           '''simple docstring'''


           return super(lowerCamelCase__						,cls	).get_tokenizers(lowerCamelCase__						,lowerCamelCase__						,**lowerCamelCase__	)






     @classmethod
     def     SCREAMING_SNAKE_CASE__  (							cls						:   Any						,lowerCamelCase__						:   Tuple						,lowerCamelCase__						:   List[Any]						,lowerCamelCase__						:   List[Any]=None						,**lowerCamelCase__						:   Any	)					->		Tuple:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						kwargs.pop("""config"""						,lowerCamelCase__	) or RagConfig.from_pretrained(lowerCamelCase__						,**lowerCamelCase__	)
           SCREAMING_SNAKE_CASE		=						RagTokenizer.from_pretrained(lowerCamelCase__						,config=lowerCamelCase__	)
           SCREAMING_SNAKE_CASE		=						rag_tokenizer.question_encoder
           SCREAMING_SNAKE_CASE		=						rag_tokenizer.generator
           if indexed_dataset is not None:
                 SCREAMING_SNAKE_CASE		=						"""custom"""
                 SCREAMING_SNAKE_CASE		=						CustomHFIndex(config.retrieval_vector_size						,lowerCamelCase__	)
           else:
                 SCREAMING_SNAKE_CASE		=						cls._build_index(lowerCamelCase__	)
           return cls(
               lowerCamelCase__						,question_encoder_tokenizer=lowerCamelCase__						,generator_tokenizer=lowerCamelCase__						,retrieval_workers=lowerCamelCase__						,index=lowerCamelCase__						,)


