

import collections
import inspect
import unittest

from transformers import FocalNetConfig
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, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
       import torch
       from torch import nn

       from transformers import (
           FocalNetBackbone,
           FocalNetForImageClassification,
           FocalNetForMaskedImageModeling,
           FocalNetModel,
       )
       from transformers.models.focalnet.modeling_focalnet import FOCALNET_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
       from PIL import Image

       from transformers import AutoImageProcessor


class  SCREAMING_SNAKE_CASE__   :





    def __init__( self		,  a		,  a=13		,  a=32		,  a=2		,  a=3		,  a=16		,  a=[32, 64, 128]		,  a=[1, 2, 1]		,  a=[2, 2, 4]		,  a=2		,  a=2.0		,  a=True		,  a=0.0		,  a=0.0		,  a=0.1		,  a="gelu"		,  a=False		,  a=True		,  a=0.02		,  a=1e-5		,  a=True		,  a=None		,  a=True		,  a=10		,  a=8		,  a=["stage1", "stage2"]		,  a=[1, 2]		,  ):
         lowercase__       :  int =			parent
         lowercase__       :  str =			batch_size
         lowercase__       :  str =			image_size
         lowercase__       :  Tuple =			patch_size
         lowercase__       :  List[Any] =			num_channels
         lowercase__       :  List[Any] =			embed_dim
         lowercase__       :  Any =			hidden_sizes
         lowercase__       :  Optional[Any] =			depths
         lowercase__       :  Union[str, Any] =			num_heads
         lowercase__       :  List[Any] =			window_size
         lowercase__       :  List[Any] =			mlp_ratio
         lowercase__       :  int =			qkv_bias
         lowercase__       :  Any =			hidden_dropout_prob
         lowercase__       :  int =			attention_probs_dropout_prob
         lowercase__       :  str =			drop_path_rate
         lowercase__       :  Dict =			hidden_act
         lowercase__       :  List[str] =			use_absolute_embeddings
         lowercase__       :  Tuple =			patch_norm
         lowercase__       :  int =			layer_norm_eps
         lowercase__       :  Union[str, Any] =			initializer_range
         lowercase__       :  Union[str, Any] =			is_training
         lowercase__       :  Dict =			scope
         lowercase__       :  Optional[Any] =			use_labels
         lowercase__       :  List[Any] =			type_sequence_label_size
         lowercase__       :  int =			encoder_stride
         lowercase__       :  Union[str, Any] =			out_features
         lowercase__       :  int =			out_indices





    def  snake_case_							( self):
         lowercase__       :  Optional[Any] =			floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

         lowercase__       :  Optional[int] =			None
         if self.use_labels:
              lowercase__       :  Any =			ids_tensor([self.batch_size]		,  self.type_sequence_label_size)

         lowercase__       :  str =			self.get_config()

         return config, pixel_values, labels





    def  snake_case_							( self):
         return FocalNetConfig(
             image_size=self.image_size		,  patch_size=self.patch_size		,  num_channels=self.num_channels		,  embed_dim=self.embed_dim		,  hidden_sizes=self.hidden_sizes		,  depths=self.depths		,  num_heads=self.num_heads		,  window_size=self.window_size		,  mlp_ratio=self.mlp_ratio		,  qkv_bias=self.qkv_bias		,  hidden_dropout_prob=self.hidden_dropout_prob		,  attention_probs_dropout_prob=self.attention_probs_dropout_prob		,  drop_path_rate=self.drop_path_rate		,  hidden_act=self.hidden_act		,  use_absolute_embeddings=self.use_absolute_embeddings		,  path_norm=self.patch_norm		,  layer_norm_eps=self.layer_norm_eps		,  initializer_range=self.initializer_range		,  encoder_stride=self.encoder_stride		,  out_features=self.out_features		,  out_indices=self.out_indices		,  )





    def  snake_case_							( self		,  a		,  a		,  a):
         lowercase__       :  Tuple =			FocalNetModel(config=__UpperCamelCase)
         model.to(__UpperCamelCase)
         model.eval()
         lowercase__       :  str =			model(__UpperCamelCase)

         lowercase__       :  Any =			((config.image_size // config.patch_size) ** 2) // (4 ** (len(config.depths) - 1))
         lowercase__       :  List[Any] =			int(config.embed_dim * 2 ** (len(config.depths) - 1))

         self.parent.assertEqual(result.last_hidden_state.shape		,  (self.batch_size, expected_seq_len, expected_dim))





    def  snake_case_							( self		,  a		,  a		,  a):
         lowercase__       :  List[Any] =			FocalNetBackbone(config=__UpperCamelCase)
         model.to(__UpperCamelCase)
         model.eval()
         lowercase__       :  Optional[Any] =			model(__UpperCamelCase)

         # verify feature maps
         self.parent.assertEqual(len(result.feature_maps)		,  len(config.out_features))
         self.parent.assertListEqual(list(result.feature_maps[0].shape)		,  [self.batch_size, self.image_size, 8, 8])

         # 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
         lowercase__       :  Union[str, Any] =			None
         lowercase__       :  Dict =			FocalNetBackbone(config=__UpperCamelCase)
         model.to(__UpperCamelCase)
         model.eval()
         lowercase__       :  int =			model(__UpperCamelCase)

         # verify feature maps
         self.parent.assertEqual(len(result.feature_maps)		,  1)
         self.parent.assertListEqual(list(result.feature_maps[0].shape)		,  [self.batch_size, self.image_size * 2, 4, 4])

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





    def  snake_case_							( self		,  a		,  a		,  a):
         lowercase__       :  Any =			FocalNetForMaskedImageModeling(config=__UpperCamelCase)
         model.to(__UpperCamelCase)
         model.eval()
         lowercase__       :  str =			model(__UpperCamelCase)
         self.parent.assertEqual(
             result.reconstruction.shape		,  (self.batch_size, self.num_channels, self.image_size, self.image_size))

         # test greyscale images
         lowercase__       :  Optional[int] =			1
         lowercase__       :  int =			FocalNetForMaskedImageModeling(__UpperCamelCase)
         model.to(__UpperCamelCase)
         model.eval()

         lowercase__       :  Any =			floats_tensor([self.batch_size, 1, self.image_size, self.image_size])
         lowercase__       :  Union[str, Any] =			model(__UpperCamelCase)
         self.parent.assertEqual(result.reconstruction.shape		,  (self.batch_size, 1, self.image_size, self.image_size))





    def  snake_case_							( self		,  a		,  a		,  a):
         lowercase__       :  Dict =			self.type_sequence_label_size
         lowercase__       :  int =			FocalNetForImageClassification(__UpperCamelCase)
         model.to(__UpperCamelCase)
         model.eval()
         lowercase__       :  Any =			model(__UpperCamelCase		,  labels=__UpperCamelCase)
         self.parent.assertEqual(result.logits.shape		,  (self.batch_size, self.type_sequence_label_size))

         # test greyscale images
         lowercase__       :  int =			1
         lowercase__       :  str =			FocalNetForImageClassification(__UpperCamelCase)
         model.to(__UpperCamelCase)
         model.eval()

         lowercase__       :  Union[str, Any] =			floats_tensor([self.batch_size, 1, self.image_size, self.image_size])
         lowercase__       :  List[Any] =			model(__UpperCamelCase)
         self.parent.assertEqual(result.logits.shape		,  (self.batch_size, self.type_sequence_label_size))





    def  snake_case_							( self):
         lowercase__       :  List[Any] =			self.prepare_config_and_inputs()

         lowercase__ ,					lowercase__ ,					lowercase__       :  Union[str, Any] =			config_and_inputs
         lowercase__       :  Any =			{'pixel_values': pixel_values}
         return config, inputs_dict



@require_torch
class  SCREAMING_SNAKE_CASE__   (_UpperCAmelCase						,			_UpperCAmelCase						,			unittest.TestCase	):
    __lowerCamelCase		:			int								=	(
        (
            FocalNetModel,
            FocalNetForImageClassification,
            FocalNetForMaskedImageModeling,
            FocalNetBackbone,
        )
        if is_torch_available()
        else ()
    )
    __lowerCamelCase		:			Any								=	(
        {'feature-extraction': FocalNetModel, 'image-classification': FocalNetForImageClassification}
        if is_torch_available()
        else {}
    )
    __lowerCamelCase		:			int								=	False

    __lowerCamelCase		:			int								=	False
    __lowerCamelCase		:			Any								=	False
    __lowerCamelCase		:			Any								=	False
    __lowerCamelCase		:			Tuple								=	False





    def  snake_case_							( self):
         lowercase__       :  List[str] =			FocalNetModelTester(self)
         lowercase__       :  str =			ConfigTester(self		,  config_class=__UpperCamelCase		,  embed_dim=37		,  has_text_modality=__UpperCamelCase)





    def  snake_case_							( self):
         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  snake_case_							( self):
         return





    def  snake_case_							( self):
         lowercase__       :  Any =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*__UpperCamelCase)





    def  snake_case_							( self):
         lowercase__       :  int =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_backbone(*__UpperCamelCase)





    def  snake_case_							( self):
         lowercase__       :  List[str] =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_masked_image_modeling(*__UpperCamelCase)





    def  snake_case_							( self):
         lowercase__       :  Union[str, Any] =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_image_classification(*__UpperCamelCase)





    @unittest.skip(reason='FocalNet does not use inputs_embeds')
    def  snake_case_							( self):
         pass





    @unittest.skip(reason='FocalNet does not use feedforward chunking')
    def  snake_case_							( self):
         pass





    def  snake_case_							( self):
         lowercase__ ,					lowercase__       :  Optional[Any] =			self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes[:-1]:
              lowercase__       :  str =			model_class(__UpperCamelCase)
              self.assertIsInstance(model.get_input_embeddings()		,  (nn.Module))
              lowercase__       :  int =			model.get_output_embeddings()
              self.assertTrue(x is None or isinstance(__UpperCamelCase		,  nn.Linear))





    def  snake_case_							( self):
         lowercase__ ,					lowercase__       :  Optional[Any] =			self.model_tester.prepare_config_and_inputs_for_common()

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

              lowercase__       :  Tuple =			['pixel_values']
              self.assertListEqual(arg_names[:1]		,  __UpperCamelCase)





    def  snake_case_							( self		,  a		,  a		,  a		,  a):
         lowercase__       :  Any =			model_class(__UpperCamelCase)
         model.to(__UpperCamelCase)
         model.eval()

         with torch.no_grad():
              lowercase__       :  Tuple =			model(**self._prepare_for_class(__UpperCamelCase		,  __UpperCamelCase))

         lowercase__       :  Tuple =			outputs.hidden_states

         lowercase__       :  Optional[int] =			getattr(
             self.model_tester		,  'expected_num_hidden_layers'		,  len(self.model_tester.depths) + 1)
         self.assertEqual(len(__UpperCamelCase)		,  __UpperCamelCase)

         # FocalNet has a different seq_length
         lowercase__       :  Optional[int] =			(
             config.patch_size
             if isinstance(config.patch_size		,  collections.abc.Iterable)
             else (config.patch_size, config.patch_size)
         )

         lowercase__       :  Any =			(image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])

         self.assertListEqual(
             list(hidden_states[0].shape[-2:])		,  [num_patches, self.model_tester.embed_dim]		,  )

         lowercase__       :  Union[str, Any] =			outputs.reshaped_hidden_states
         self.assertEqual(len(__UpperCamelCase)		,  __UpperCamelCase)

         lowercase__ ,					lowercase__ ,					lowercase__ ,					lowercase__       :  List[Any] =			reshaped_hidden_states[0].shape
         lowercase__       :  str =			(
             reshaped_hidden_states[0].view(__UpperCamelCase		,  __UpperCamelCase		,  height * width).permute(0		,  2		,  1)
         )
         self.assertListEqual(
             list(reshaped_hidden_states.shape[-2:])		,  [num_patches, self.model_tester.embed_dim]		,  )





    def  snake_case_							( self):
         lowercase__ ,					lowercase__       :  Optional[int] =			self.model_tester.prepare_config_and_inputs_for_common()

         lowercase__       :  Union[str, Any] =			(
             self.model_tester.image_size
             if isinstance(self.model_tester.image_size		,  collections.abc.Iterable)
             else (self.model_tester.image_size, self.model_tester.image_size)
         )

         for model_class in self.all_model_classes[:-1]:
              lowercase__       :  Optional[Any] =			True
              self.check_hidden_states_output(__UpperCamelCase		,  __UpperCamelCase		,  __UpperCamelCase		,  __UpperCamelCase)

              # check that output_hidden_states also work using config
              del inputs_dict["output_hidden_states"]
              lowercase__       :  Optional[int] =			True

              self.check_hidden_states_output(__UpperCamelCase		,  __UpperCamelCase		,  __UpperCamelCase		,  __UpperCamelCase)





    def  snake_case_							( self):
         lowercase__ ,					lowercase__       :  str =			self.model_tester.prepare_config_and_inputs_for_common()
         lowercase__       :  Any =			3

         lowercase__       :  Optional[int] =			(
             self.model_tester.image_size
             if isinstance(self.model_tester.image_size		,  collections.abc.Iterable)
             else (self.model_tester.image_size, self.model_tester.image_size)
         )
         lowercase__       :  Tuple =			(
             config.patch_size
             if isinstance(config.patch_size		,  collections.abc.Iterable)
             else (config.patch_size, config.patch_size)
         )

         lowercase__       :  str =			image_size[0] + patch_size[0] - (image_size[0] % patch_size[0])
         lowercase__       :  int =			image_size[1] + patch_size[1] - (image_size[1] % patch_size[1])

         for model_class in self.all_model_classes[:-1]:
              lowercase__       :  Any =			True
              self.check_hidden_states_output(__UpperCamelCase		,  __UpperCamelCase		,  __UpperCamelCase		,  (padded_height, padded_width))

              # check that output_hidden_states also work using config
              del inputs_dict["output_hidden_states"]
              lowercase__       :  Union[str, Any] =			True
              self.check_hidden_states_output(__UpperCamelCase		,  __UpperCamelCase		,  __UpperCamelCase		,  (padded_height, padded_width))





    @slow
    def  snake_case_							( self):






         for model_name in FOCALNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
              lowercase__       :  Dict =			FocalNetModel.from_pretrained(__UpperCamelCase)
              self.assertIsNotNone(__UpperCamelCase)





    def  snake_case_							( self):
         lowercase__ ,					lowercase__       :  Optional[Any] =			self.model_tester.prepare_config_and_inputs_for_common()

         lowercase__       :  str =			_config_zero_init(__UpperCamelCase)
         for model_class in self.all_model_classes:
              lowercase__       :  Optional[Any] =			model_class(config=__UpperCamelCase)
              for name, param in model.named_parameters():
                   if "embeddings" not in name and param.requires_grad:
                        self.assertIn(
                            ((param.data.mean() * 1e9).round() / 1e9).item()		,  [0.0, 1.0]		,  msg=f"""Parameter {name} of model {model_class} seems not properly initialized"""		,  )



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





    @cached_property
    def  snake_case_							( self):
         return AutoImageProcessor.from_pretrained('microsoft/focalnet-tiny') if is_vision_available() else None





    @slow
    def  snake_case_							( self):
         lowercase__       :  Dict =			FocalNetForImageClassification.from_pretrained('microsoft/focalnet-tiny').to(__UpperCamelCase)
         lowercase__       :  List[Any] =			self.default_image_processor

         lowercase__       :  Optional[int] =			Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png')
         lowercase__       :  Optional[Any] =			image_processor(images=__UpperCamelCase		,  return_tensors='pt').to(__UpperCamelCase)

         # forward pass
         with torch.no_grad():
              lowercase__       :  Union[str, Any] =			model(**__UpperCamelCase)

         # verify the logits
         lowercase__       :  Tuple =			torch.Size((1, 1000))
         self.assertEqual(outputs.logits.shape		,  __UpperCamelCase)
         lowercase__       :  str =			torch.tensor([0.2_166, -0.4_368, 0.2_191]).to(__UpperCamelCase)
         self.assertTrue(torch.allclose(outputs.logits[0, :3]		,  __UpperCamelCase		,  atol=1e-4))
         self.assertTrue(outputs.logits.argmax(dim=-1).item()		,  281)



@require_torch
class  SCREAMING_SNAKE_CASE__   (_UpperCAmelCase						,			unittest.TestCase	):
    __lowerCamelCase		:			Optional[Any]								=	(FocalNetBackbone,) if is_torch_available() else ()
    __lowerCamelCase		:			Optional[Any]								=	FocalNetConfig

    __lowerCamelCase		:			Optional[int]								=	False





    def  snake_case_							( self):
         lowercase__       :  Dict =			FocalNetModelTester(self)



"""simple docstring"""



import os
import warnings
from typing import List, Optional

from ...tokenization_utils_base import BatchEncoding
from ...utils import logging
from .configuration_rag import RagConfig


a       :						List[str]       		=		logging.get_logger(__name__)




class 		a_    :







    def __init__(      self      :							List[str]  ,       __UpperCamelCase      :							Any  ,       __UpperCamelCase      :							Tuple  )		->Tuple:






        '''simple docstring'''






        _UpperCAmelCase        = question_encoder
        _UpperCAmelCase        = generator
        _UpperCAmelCase        = self.question_encoder







    def   _snake_case       (      self      :							int  ,       __UpperCamelCase      :							str  )		->Dict:






        '''simple docstring'''






        if os.path.isfile(__UpperCamelCase  ):
            raise ValueError(f"""Provided path ({save_directory}) should be a directory, not a file"""  )
        os.makedirs(__UpperCamelCase  ,       exist_ok=__UpperCamelCase  )
        _UpperCAmelCase        = os.path.join(__UpperCamelCase  ,       """question_encoder_tokenizer"""  )
        _UpperCAmelCase        = os.path.join(__UpperCamelCase  ,       """generator_tokenizer"""  )
        self.question_encoder.save_pretrained(__UpperCamelCase  )
        self.generator.save_pretrained(__UpperCamelCase  )







    @classmethod
    def   _snake_case       (      cls      :							Optional[Any]  ,       __UpperCamelCase      :							Tuple  ,       **__UpperCamelCase      :							Optional[Any]  )		->Tuple:






        '''simple docstring'''






        from ..auto.tokenization_auto import AutoTokenizer

        _UpperCAmelCase        = kwargs.pop("""config"""  ,       __UpperCamelCase  )

        if config is None:
            _UpperCAmelCase        = RagConfig.from_pretrained(__UpperCamelCase  )

        _UpperCAmelCase        = AutoTokenizer.from_pretrained(
            __UpperCamelCase  ,       config=config.question_encoder  ,       subfolder="""question_encoder_tokenizer"""  )
        _UpperCAmelCase        = AutoTokenizer.from_pretrained(
            __UpperCamelCase  ,       config=config.generator  ,       subfolder="""generator_tokenizer"""  )
        return cls(question_encoder=__UpperCamelCase  ,       generator=__UpperCamelCase  )







    def __call__(      self      :							Any  ,       *__UpperCamelCase      :							List[Any]  ,       **__UpperCamelCase      :							Optional[Any]  )		->Optional[Any]:






        '''simple docstring'''






        return self.current_tokenizer(*__UpperCamelCase  ,       **__UpperCamelCase  )







    def   _snake_case       (      self      :							Union[str, Any]  ,       *__UpperCamelCase      :							List[Any]  ,       **__UpperCamelCase      :							Tuple  )		->Any:






        '''simple docstring'''






        return self.generator.batch_decode(*__UpperCamelCase  ,       **__UpperCamelCase  )







    def   _snake_case       (      self      :							Optional[int]  ,       *__UpperCamelCase      :							Optional[int]  ,       **__UpperCamelCase      :							List[str]  )		->Union[str, Any]:






        '''simple docstring'''






        return self.generator.decode(*__UpperCamelCase  ,       **__UpperCamelCase  )







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






        '''simple docstring'''






        _UpperCAmelCase        = self.question_encoder







    def   _snake_case       (      self      :							Tuple  )		->Optional[int]:






        '''simple docstring'''






        _UpperCAmelCase        = self.generator







    def   _snake_case       (      self      :							List[Any]  ,       __UpperCamelCase      :							List[str]  ,       __UpperCamelCase      :							Optional[List[str]] = None  ,       __UpperCamelCase      :							Optional[int] = None  ,       __UpperCamelCase      :							Optional[int] = None  ,       __UpperCamelCase      :							str = "longest"  ,       __UpperCamelCase      :							str = None  ,       __UpperCamelCase      :							bool = True  ,       **__UpperCamelCase      :							Tuple  ,       )		->BatchEncoding:






        '''simple docstring'''






        warnings.warn(
            """`prepare_seq2seq_batch` is deprecated and will be removed in version 5 of 🤗 Transformers. Use the """
            """regular `__call__` method to prepare your inputs and the tokenizer under the `with_target_tokenizer` """
            """context manager to prepare your targets. See the documentation of your specific tokenizer for more """
            """details"""  ,       __UpperCamelCase  ,       )
        if max_length is None:
            _UpperCAmelCase        = self.current_tokenizer.model_max_length
        _UpperCAmelCase        = self(
            __UpperCamelCase  ,       add_special_tokens=__UpperCamelCase  ,       return_tensors=__UpperCamelCase  ,       max_length=__UpperCamelCase  ,       padding=__UpperCamelCase  ,       truncation=__UpperCamelCase  ,       **__UpperCamelCase  ,       )
        if tgt_texts is None:
            return model_inputs
        # Process tgt_texts
        if max_target_length is None:
            _UpperCAmelCase        = self.current_tokenizer.model_max_length
        _UpperCAmelCase        = self(
            text_target=__UpperCamelCase  ,       add_special_tokens=__UpperCamelCase  ,       return_tensors=__UpperCamelCase  ,       padding=__UpperCamelCase  ,       max_length=__UpperCamelCase  ,       truncation=__UpperCamelCase  ,       **__UpperCamelCase  ,       )
        _UpperCAmelCase        = labels["""input_ids"""]
        return model_inputs