"""simple docstring"""
import inspect
import unittest

from transformers import SegformerConfig, is_torch_available, is_vision_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, slow, torch_device

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 (
     MODEL_MAPPING,
     SegformerForImageClassification,
     SegformerForSemanticSegmentation,
     SegformerModel,
 )
 from transformers.models.segformer.modeling_segformer import SEGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
 from PIL import Image

 from transformers import SegformerImageProcessor




class    __snake_case       (     __lowerCAmelCase				):



      def   lowerCamelCase_   (						self)	->			Optional[int]:





        '''simple docstring'''





        a__:     List[str]										=	self.config_class(**self.inputs_dict)
        self.parent.assertTrue(hasattr(lowercase    ,			'hidden_sizes'))
        self.parent.assertTrue(hasattr(lowercase    ,			'num_attention_heads'))
        self.parent.assertTrue(hasattr(lowercase    ,			'num_encoder_blocks'))







class    __snake_case       :



      def __init__(						self    ,			lowercase    ,			lowercase=13    ,			lowercase=64    ,			lowercase=3    ,			lowercase=4    ,			lowercase=[2, 2, 2, 2]    ,			lowercase=[8, 4, 2, 1]    ,			lowercase=[16, 32, 64, 1_28]    ,			lowercase=[1, 4, 8, 16]    ,			lowercase=[1, 2, 4, 8]    ,			lowercase=True    ,			lowercase=True    ,			lowercase="gelu"    ,			lowercase=0.1    ,			lowercase=0.1    ,			lowercase=0.02    ,			lowercase=3    ,			lowercase=None    ,			)	->			Any:





        '''simple docstring'''





        a__:     Any										=	parent
        a__:     List[Any]										=	batch_size
        a__:     str										=	image_size
        a__:     int										=	num_channels
        a__:     List[Any]										=	num_encoder_blocks
        a__:     Any										=	sr_ratios
        a__:     Tuple										=	depths
        a__:     Dict										=	hidden_sizes
        a__:     Tuple										=	downsampling_rates
        a__:     Union[str, Any]										=	num_attention_heads
        a__:     str										=	is_training
        a__:     List[Any]										=	use_labels
        a__:     Dict										=	hidden_act
        a__:     Dict										=	hidden_dropout_prob
        a__:     Union[str, Any]										=	attention_probs_dropout_prob
        a__:     Optional[Any]										=	initializer_range
        a__:     Any										=	num_labels
        a__:     Optional[Any]										=	scope



      def   lowerCamelCase_   (						self)	->			Tuple:





        '''simple docstring'''





        a__:     Optional[Any]										=	floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

        a__:     Optional[Any]										=	None
        if self.use_labels:
          a__:     str										=	ids_tensor([self.batch_size, self.image_size, self.image_size]    ,			self.num_labels)

        a__:     Tuple										=	self.get_config()
        return config, pixel_values, labels



      def   lowerCamelCase_   (						self)	->			str:





        '''simple docstring'''





        return SegformerConfig(
            image_size=self.image_size    ,			num_channels=self.num_channels    ,			num_encoder_blocks=self.num_encoder_blocks    ,			depths=self.depths    ,			hidden_sizes=self.hidden_sizes    ,			num_attention_heads=self.num_attention_heads    ,			hidden_act=self.hidden_act    ,			hidden_dropout_prob=self.hidden_dropout_prob    ,			attention_probs_dropout_prob=self.attention_probs_dropout_prob    ,			initializer_range=self.initializer_range    ,			)



      def   lowerCamelCase_   (						self    ,			lowercase    ,			lowercase    ,			lowercase)	->			Dict:





        '''simple docstring'''





        a__:     List[Any]										=	SegformerModel(config=lowercase)
        model.to(lowercase)
        model.eval()
        a__:     Tuple										=	model(lowercase)
        a__:     str										=	self.image_size // (self.downsampling_rates[-1] * 2)
        self.parent.assertEqual(
            result.last_hidden_state.shape    ,			(self.batch_size, self.hidden_sizes[-1], expected_height, expected_width))



      def   lowerCamelCase_   (						self    ,			lowercase    ,			lowercase    ,			lowercase)	->			List[str]:





        '''simple docstring'''





        a__:     Optional[Any]										=	self.num_labels
        a__:     Union[str, Any]										=	SegformerForSemanticSegmentation(lowercase)
        model.to(lowercase)
        model.eval()
        a__:     int										=	model(lowercase)
        self.parent.assertEqual(
            result.logits.shape    ,			(self.batch_size, self.num_labels, self.image_size // 4, self.image_size // 4))
        a__:     int										=	model(lowercase    ,			labels=lowercase)
        self.parent.assertEqual(
            result.logits.shape    ,			(self.batch_size, self.num_labels, self.image_size // 4, self.image_size // 4))
        self.parent.assertGreater(result.loss    ,			0.0)



      def   lowerCamelCase_   (						self    ,			lowercase    ,			lowercase    ,			lowercase)	->			str:





        '''simple docstring'''





        a__:     Tuple										=	1
        a__:     Optional[Any]										=	SegformerForSemanticSegmentation(config=lowercase)
        model.to(lowercase)
        model.eval()
        a__:     List[str]										=	torch.randint(0    ,			1    ,			(self.batch_size, self.image_size, self.image_size)).to(lowercase)
        a__:     Union[str, Any]										=	model(lowercase    ,			labels=lowercase)
        self.parent.assertGreater(result.loss    ,			0.0)



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





        '''simple docstring'''





        a__:     str										=	self.prepare_config_and_inputs()
        a__							,    a__							,    a__:     str										=	config_and_inputs
        a__:     Optional[Any]										=	{'pixel_values': pixel_values}
        return config, inputs_dict







@require_torch
class    __snake_case       (     __lowerCAmelCase       ,      __lowerCAmelCase       ,      unittest.TestCase				):
      a__        =		(
          (
              SegformerModel,
              SegformerForSemanticSegmentation,
              SegformerForImageClassification,
          )
          if is_torch_available()
          else ()
      )
      a__        =		(
          {
              """feature-extraction""": SegformerModel,
              """image-classification""": SegformerForImageClassification,
              """image-segmentation""": SegformerForSemanticSegmentation,
          }
          if is_torch_available()
          else {}
      )

      a__        =		True
      a__        =		False
      a__        =		False
      a__        =		False



      def   lowerCamelCase_   (						self)	->			Optional[int]:





        '''simple docstring'''





        a__:     Optional[int]										=	SegformerModelTester(self)
        a__:     Union[str, Any]										=	SegformerConfigTester(self    ,			config_class=lowercase)



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





        '''simple docstring'''





        self.config_tester.run_common_tests()



      def   lowerCamelCase_   (						self)	->			Any:





        '''simple docstring'''





        a__:     int										=	self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*lowercase)



      def   lowerCamelCase_   (						self)	->			Any:





        '''simple docstring'''





        a__:     Optional[int]										=	self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_binary_image_segmentation(*lowercase)



      def   lowerCamelCase_   (						self)	->			Any:





        '''simple docstring'''





        a__:     List[str]										=	self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_image_segmentation(*lowercase)



      @unittest.skip('SegFormer does not use inputs_embeds')
      def   lowerCamelCase_   (						self)	->			Optional[int]:





        '''simple docstring'''





        pass



      @unittest.skip('SegFormer does not have get_input_embeddings method and get_output_embeddings methods')
      def   lowerCamelCase_   (						self)	->			List[Any]:





        '''simple docstring'''





        pass



      def   lowerCamelCase_   (						self)	->			Optional[int]:





        '''simple docstring'''





        a__							,    a__:     Dict										=	self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
          a__:     Union[str, Any]										=	model_class(lowercase)
          a__:     int										=	inspect.signature(model.forward)
          # signature.parameters is an OrderedDict => so arg_names order is deterministic
          a__:     Union[str, Any]										=	[*signature.parameters.keys()]

          a__:     List[str]										=	['pixel_values']
          self.assertListEqual(arg_names[:1]    ,			lowercase)



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





        '''simple docstring'''





        a__							,    a__:     Optional[Any]										=	self.model_tester.prepare_config_and_inputs_for_common()
        a__:     List[Any]										=	True

        for model_class in self.all_model_classes:
          a__:     Tuple										=	True
          a__:     List[str]										=	False
          a__:     Optional[Any]										=	True
          a__:     Union[str, Any]										=	model_class(lowercase)
          model.to(lowercase)
          model.eval()
          with torch.no_grad():
            a__:     Dict										=	model(**self._prepare_for_class(lowercase    ,			lowercase))
          a__:     Optional[Any]										=	outputs.attentions

          a__:     Dict										=	sum(self.model_tester.depths)
          self.assertEqual(len(lowercase)    ,			lowercase)

          # check that output_attentions also work using config
          del inputs_dict["output_attentions"]
          a__:     List[str]										=	True
          a__:     List[str]										=	model_class(lowercase)
          model.to(lowercase)
          model.eval()
          with torch.no_grad():
            a__:     str										=	model(**self._prepare_for_class(lowercase    ,			lowercase))
          a__:     List[str]										=	outputs.attentions

          self.assertEqual(len(lowercase)    ,			lowercase)

          # verify the first attentions (first block, first layer)
          a__:     Optional[int]										=	(self.model_tester.image_size // 4) ** 2
          a__:     Any										=	(self.model_tester.image_size // (4 * self.model_tester.sr_ratios[0])) ** 2
          self.assertListEqual(
              list(attentions[0].shape[-3:])    ,			[self.model_tester.num_attention_heads[0], expected_seq_len, expected_reduced_seq_len]    ,			)

          # verify the last attentions (last block, last layer)
          a__:     List[Any]										=	(self.model_tester.image_size // 32) ** 2
          a__:     Dict										=	(self.model_tester.image_size // (32 * self.model_tester.sr_ratios[-1])) ** 2
          self.assertListEqual(
              list(attentions[-1].shape[-3:])    ,			[self.model_tester.num_attention_heads[-1], expected_seq_len, expected_reduced_seq_len]    ,			)
          a__:     Tuple										=	len(lowercase)

          # Check attention is always last and order is fine
          a__:     Union[str, Any]										=	True
          a__:     Any										=	True
          a__:     Optional[int]										=	model_class(lowercase)
          model.to(lowercase)
          model.eval()
          with torch.no_grad():
            a__:     int										=	model(**self._prepare_for_class(lowercase    ,			lowercase))

          self.assertEqual(out_len + 1    ,			len(lowercase))

          a__:     Dict										=	outputs.attentions

          self.assertEqual(len(lowercase)    ,			lowercase)
          # verify the first attentions (first block, first layer)
          a__:     Any										=	(self.model_tester.image_size // 4) ** 2
          a__:     Dict										=	(self.model_tester.image_size // (4 * self.model_tester.sr_ratios[0])) ** 2
          self.assertListEqual(
              list(self_attentions[0].shape[-3:])    ,			[self.model_tester.num_attention_heads[0], expected_seq_len, expected_reduced_seq_len]    ,			)



      def   lowerCamelCase_   (						self)	->			Optional[Any]:





        '''simple docstring'''





        def check_hidden_states_output(lowercase    ,			lowercase    ,			lowercase):
          a__:     List[Any]										=	model_class(lowercase)
          model.to(lowercase)
          model.eval()

          with torch.no_grad():
            a__:     Optional[int]										=	model(**self._prepare_for_class(lowercase    ,			lowercase))

          a__:     List[str]										=	outputs.hidden_states

          a__:     Optional[int]										=	self.model_tester.num_encoder_blocks
          self.assertEqual(len(lowercase)    ,			lowercase)

          # verify the first hidden states (first block)
          self.assertListEqual(
              list(hidden_states[0].shape[-3:])    ,			[
                  self.model_tester.hidden_sizes[0],
                  self.model_tester.image_size // 4,
                  self.model_tester.image_size // 4,
              ]    ,			)

        a__							,    a__:     Dict										=	self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
          a__:     List[str]										=	True
          check_hidden_states_output(lowercase    ,			lowercase    ,			lowercase)

          # check that output_hidden_states also work using config
          del inputs_dict["output_hidden_states"]
          a__:     List[Any]										=	True

          check_hidden_states_output(lowercase    ,			lowercase    ,			lowercase)



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





        '''simple docstring'''





        if not self.model_tester.is_training:
          return

        a__							,    a__:     Dict										=	self.model_tester.prepare_config_and_inputs_for_common()
        a__:     Tuple										=	True

        for model_class in self.all_model_classes:
          if model_class in get_values(lowercase):
            continue

          a__:     Union[str, Any]										=	model_class(lowercase)
          model.to(lowercase)
          model.train()
          a__:     List[Any]										=	self._prepare_for_class(lowercase    ,			lowercase    ,			return_labels=lowercase)
          a__:     List[str]										=	model(**lowercase).loss
          loss.backward()



      @unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.')
      def   lowerCamelCase_   (						self)	->			List[Any]:





        '''simple docstring'''





        pass



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





        '''simple docstring'''





        for model_name in SEGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          a__:     Optional[Any]										=	SegformerModel.from_pretrained(lowercase)
          self.assertIsNotNone(lowercase)







def 						__a   (	)	->Optional[Any]:
  a__:     Optional[int]										=	Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'   )
  return image




@require_torch
class    __snake_case       (     unittest.TestCase				):



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





        '''simple docstring'''





        a__:     str										=	SegformerImageProcessor(
            image_scale=(5_12, 5_12)    ,			keep_ratio=lowercase    ,			align=lowercase    ,			do_random_crop=lowercase)
        a__:     List[str]										=	SegformerForSemanticSegmentation.from_pretrained('nvidia/segformer-b0-finetuned-ade-512-512').to(
            lowercase)

        a__:     Optional[int]										=	prepare_img()
        a__:     int										=	image_processor(images=lowercase    ,			return_tensors='pt')
        a__:     Dict										=	encoded_inputs.pixel_values.to(lowercase)

        with torch.no_grad():
          a__:     List[Any]										=	model(lowercase)

        a__:     Tuple										=	torch.Size((1, model.config.num_labels, 1_28, 1_28))
        self.assertEqual(outputs.logits.shape    ,			lowercase)

        a__:     Optional[Any]										=	torch.tensor(
            [
                [[-4.6310, -5.5232, -6.2356], [-5.1921, -6.1444, -6.5996], [-5.4424, -6.2790, -6.7574]],
                [[-12.1391, -13.3122, -13.9554], [-12.8732, -13.9352, -14.3563], [-12.9438, -13.8226, -14.2513]],
                [[-12.5134, -13.4686, -14.4915], [-12.8669, -14.4343, -14.7758], [-13.2523, -14.5819, -15.0694]],
            ]).to(lowercase)
        self.assertTrue(torch.allclose(outputs.logits[0, :3, :3, :3]    ,			lowercase    ,			atol=1e-4))



      @slow
      def   lowerCamelCase_   (						self)	->			Tuple:





        '''simple docstring'''





        a__:     List[str]										=	SegformerImageProcessor(
            image_scale=(5_12, 5_12)    ,			keep_ratio=lowercase    ,			align=lowercase    ,			do_random_crop=lowercase)
        a__:     int										=	SegformerForSemanticSegmentation.from_pretrained(
            'nvidia/segformer-b1-finetuned-cityscapes-1024-1024').to(lowercase)

        a__:     str										=	prepare_img()
        a__:     Optional[Any]										=	image_processor(images=lowercase    ,			return_tensors='pt')
        a__:     Dict										=	encoded_inputs.pixel_values.to(lowercase)

        with torch.no_grad():
          a__:     Optional[Any]										=	model(lowercase)

        a__:     Tuple										=	torch.Size((1, model.config.num_labels, 1_28, 1_28))
        self.assertEqual(outputs.logits.shape    ,			lowercase)

        a__:     Any										=	torch.tensor(
            [
                [[-13.5748, -13.9111, -12.6500], [-14.3500, -15.3683, -14.2328], [-14.7532, -16.0424, -15.6087]],
                [[-17.1651, -15.8725, -12.9653], [-17.2580, -17.3718, -14.8223], [-16.6058, -16.8783, -16.7452]],
                [[-3.6456, -3.0209, -1.4203], [-3.0797, -3.1959, -2.0000], [-1.8757, -1.9217, -1.6997]],
            ]).to(lowercase)
        self.assertTrue(torch.allclose(outputs.logits[0, :3, :3, :3]    ,			lowercase    ,			atol=1e-1))



      @slow
      def   lowerCamelCase_   (						self)	->			List[Any]:





        '''simple docstring'''





        a__:     Optional[Any]										=	SegformerImageProcessor(
            image_scale=(5_12, 5_12)    ,			keep_ratio=lowercase    ,			align=lowercase    ,			do_random_crop=lowercase)
        a__:     Optional[int]										=	SegformerForSemanticSegmentation.from_pretrained('nvidia/segformer-b0-finetuned-ade-512-512').to(
            lowercase)

        a__:     Tuple										=	prepare_img()
        a__:     Optional[Any]										=	image_processor(images=lowercase    ,			return_tensors='pt')
        a__:     Any										=	encoded_inputs.pixel_values.to(lowercase)

        with torch.no_grad():
          a__:     Tuple										=	model(lowercase)

        a__:     Tuple										=	outputs.logits.detach().cpu()

        a__:     Union[str, Any]										=	image_processor.post_process_semantic_segmentation(outputs=lowercase    ,			target_sizes=[(5_00, 3_00)])
        a__:     str										=	torch.Size((5_00, 3_00))
        self.assertEqual(segmentation[0].shape    ,			lowercase)

        a__:     Optional[Any]										=	image_processor.post_process_semantic_segmentation(outputs=lowercase)
        a__:     List[str]										=	torch.Size((1_28, 1_28))
        self.assertEqual(segmentation[0].shape    ,			lowercase)



"""simple docstring"""
import logging
import random

import ray

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


lowercase__					       =	logging.getLogger(__name__)




class    __snake_case       :



      def __init__(						self)	->			Optional[int]:





        '''simple docstring'''





        a__:     Optional[Any]										=	False



      def   lowerCamelCase_   (						self    ,			lowercase    ,			lowercase    ,			lowercase    ,			lowercase)	->			str:





        '''simple docstring'''





        if not self.initialized:
          a__:     Optional[int]										=	RagRetriever(
              lowercase    ,			question_encoder_tokenizer=lowercase    ,			generator_tokenizer=lowercase    ,			index=lowercase    ,			init_retrieval=lowercase    ,			)
          a__:     Optional[int]										=	True



      def   lowerCamelCase_   (						self)	->			int:





        '''simple docstring'''





        self.retriever.index.init_index()



      def   lowerCamelCase_   (						self    ,			lowercase    ,			lowercase)	->			Union[str, Any]:





        '''simple docstring'''





        a__							,    a__:     str										=	self.retriever._main_retrieve(lowercase    ,			lowercase)
        return doc_ids, retrieved_doc_embeds







class    __snake_case       (     __lowerCAmelCase				):



      def __init__(						self    ,			lowercase    ,			lowercase    ,			lowercase    ,			lowercase    ,			lowercase=None)	->			int:





        '''simple docstring'''





        if index is not None and index.is_initialized() and len(lowercase) > 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__(
            lowercase    ,			question_encoder_tokenizer=lowercase    ,			generator_tokenizer=lowercase    ,			index=lowercase    ,			init_retrieval=lowercase    ,			)
        a__:     Any										=	retrieval_workers
        if len(self.retrieval_workers) > 0:
          ray.get(
              [
                  worker.create_rag_retriever.remote(lowercase    ,			lowercase    ,			lowercase    ,			lowercase)
                  for worker in self.retrieval_workers
              ])



      def   lowerCamelCase_   (						self)	->			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   lowerCamelCase_   (						self    ,			lowercase    ,			lowercase)	->			Union[str, Any]:





        '''simple docstring'''





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



      @classmethod
      def   lowerCamelCase_   (						cls    ,			lowercase    ,			lowercase=None    ,			**lowercase)	->			Tuple:





        '''simple docstring'''





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



      @classmethod
      def   lowerCamelCase_   (						cls    ,			lowercase    ,			lowercase    ,			lowercase=None    ,			**lowercase)	->			Union[str, Any]:





        '''simple docstring'''





        a__:     Optional[int]										=	kwargs.pop('config'    ,			lowercase) or RagConfig.from_pretrained(lowercase    ,			**lowercase)
        a__:     Union[str, Any]										=	RagTokenizer.from_pretrained(lowercase    ,			config=lowercase)
        a__:     int										=	rag_tokenizer.question_encoder
        a__:     Any										=	rag_tokenizer.generator
        if indexed_dataset is not None:
          a__:     List[Any]										=	'custom'
          a__:     Optional[Any]										=	CustomHFIndex(config.retrieval_vector_size    ,			lowercase)
        else:
          a__:     Dict										=	cls._build_index(lowercase)
        return cls(
            lowercase    ,			question_encoder_tokenizer=lowercase    ,			generator_tokenizer=lowercase    ,			retrieval_workers=lowercase    ,			index=lowercase    ,			)


