






"""simple docstring"""






import argparse

import torch

# Step 1. clone https://github.com/microsoft/unilm
# Step 2. git checkout to https://github.com/microsoft/unilm/commit/b94ec76c36f02fb2b0bf0dcb0b8554a2185173cd
# Step 3. cd unilm
# Step 4. ln -s $(realpath wavlm/modules.py) ./  # create simlink
# import classes
from unilm.wavlm.WavLM import WavLM as WavLMOrig
from unilm.wavlm.WavLM import WavLMConfig as WavLMConfigOrig

from transformers import WavLMConfig, WavLMModel, logging


logging.set_verbosity_info()
UpperCAmelCase_     :				Optional[Any]         =				logging.get_logger(__name__)

UpperCAmelCase_     :				Union[str, Any]         =				{
    """post_extract_proj""": """feature_projection.projection""",
    """encoder.pos_conv.0""": """encoder.pos_conv_embed.conv""",
    """self_attn.k_proj""": """encoder.layers.*.attention.k_proj""",
    """self_attn.v_proj""": """encoder.layers.*.attention.v_proj""",
    """self_attn.q_proj""": """encoder.layers.*.attention.q_proj""",
    """self_attn.out_proj""": """encoder.layers.*.attention.out_proj""",
    """self_attn.grep_linear""": """encoder.layers.*.attention.gru_rel_pos_linear""",
    """self_attn.relative_attention_bias""": """encoder.layers.*.attention.rel_attn_embed""",
    """self_attn.grep_a""": """encoder.layers.*.attention.gru_rel_pos_const""",
    """self_attn_layer_norm""": """encoder.layers.*.layer_norm""",
    """fc1""": """encoder.layers.*.feed_forward.intermediate_dense""",
    """fc2""": """encoder.layers.*.feed_forward.output_dense""",
    """final_layer_norm""": """encoder.layers.*.final_layer_norm""",
    """encoder.layer_norm""": """encoder.layer_norm""",
    """w2v_model.layer_norm""": """feature_projection.layer_norm""",
    """quantizer.weight_proj""": """quantizer.weight_proj""",
    """quantizer.vars""": """quantizer.codevectors""",
    """project_q""": """project_q""",
    """final_proj""": """project_hid""",
    """w2v_encoder.proj""": """ctc_proj""",
    """mask_emb""": """masked_spec_embed""",
}
UpperCAmelCase_     :				List[str]         =				[
    """ctc_proj""",
    """quantizer.weight_proj""",
    """quantizer.codevectors""",
    """project_q""",
    """project_hid""",
]



def       _A    (__a					,       __a					,       __a					,       __a					,       __a							)				-> List[str]:
    """simple docstring"""






    for attribute in key.split('''.'''							):
        SCREAMING_SNAKE_CASE_   :						Dict							=	getattr(__a					,       __a							)

    if weight_type is not None:
        SCREAMING_SNAKE_CASE_   :						str							=	getattr(__a					,       __a							).shape
    else:
        SCREAMING_SNAKE_CASE_   :						int							=	hf_pointer.shape

    assert hf_shape == value.shape, (
        f'Shape of hf {key + "." + weight_type if weight_type is not None else ""} is {hf_shape}, but should be'
        f' {value.shape} for {full_name}'
    )

    if weight_type == "weight":
        SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	value
    elif weight_type == "weight_g":
        SCREAMING_SNAKE_CASE_   :						Optional[int]							=	value
    elif weight_type == "weight_v":
        SCREAMING_SNAKE_CASE_   :						int							=	value
    elif weight_type == "bias":
        SCREAMING_SNAKE_CASE_   :						Tuple							=	value
    else:
        SCREAMING_SNAKE_CASE_   :						List[Any]							=	value

    logger.info(f'{key + "." + weight_type if weight_type is not None else ""} was initialized from {full_name}.'							)



def       _A    (__a					,       __a							)				-> Dict:
    """simple docstring"""






    SCREAMING_SNAKE_CASE_   :						Dict							=	[]
    SCREAMING_SNAKE_CASE_   :						int							=	fairseq_model.state_dict()

    SCREAMING_SNAKE_CASE_   :						Optional[int]							=	hf_model.feature_extractor

    for name, value in fairseq_dict.items():
        SCREAMING_SNAKE_CASE_   :						List[Any]							=	False
        if "conv_layers" in name:
            load_conv_layer(
                __a					,       __a					,       __a					,       __a					,       hf_model.config.feat_extract_norm == '''group'''					,       )
            SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	True
        else:
            for key, mapped_key in MAPPING.items():
                if key in name or key.split('''w2v_model.'''							)[-1] == name.split('''.'''							)[0]:
                    SCREAMING_SNAKE_CASE_   :						Tuple							=	True
                    if "*" in mapped_key:
                        SCREAMING_SNAKE_CASE_   :						Optional[int]							=	name.split(__a							)[0].split('''.'''							)[-2]
                        SCREAMING_SNAKE_CASE_   :						Tuple							=	mapped_key.replace('''*'''					,       __a							)
                    if "weight_g" in name:
                        SCREAMING_SNAKE_CASE_   :						Any							=	'''weight_g'''
                    elif "weight_v" in name:
                        SCREAMING_SNAKE_CASE_   :						str							=	'''weight_v'''
                    elif "bias" in name and "relative_attention_bias" not in name:
                        SCREAMING_SNAKE_CASE_   :						Tuple							=	'''bias'''
                    elif "weight" in name:
                        # TODO: don't match quantizer.weight_proj
                        SCREAMING_SNAKE_CASE_   :						str							=	'''weight'''
                    else:
                        SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	None

                    set_recursively(__a					,       __a					,       __a					,       __a					,       __a							)
                continue
        if not is_used:
            unused_weights.append(__a							)

    logger.warning(f'Unused weights: {unused_weights}'							)



def       _A    (__a					,       __a					,       __a					,       __a					,       __a							)				-> Optional[int]:
    """simple docstring"""






    SCREAMING_SNAKE_CASE_   :						List[Any]							=	full_name.split('''conv_layers.'''							)[-1]
    SCREAMING_SNAKE_CASE_   :						int							=	name.split('''.'''							)
    SCREAMING_SNAKE_CASE_   :						Optional[int]							=	int(items[0]							)
    SCREAMING_SNAKE_CASE_   :						str							=	int(items[1]							)

    if type_id == 0:
        if "bias" in name:
            assert value.shape == feature_extractor.conv_layers[layer_id].conv.bias.data.shape, (
                f'{full_name} has size {value.shape}, but'
                f' {feature_extractor.conv_layers[layer_id].conv.bias.data.shape} was found.'
            )
            SCREAMING_SNAKE_CASE_   :						str							=	value
            logger.info(f'Feat extract conv layer {layer_id} was initialized from {full_name}.'							)
        elif "weight" in name:
            assert value.shape == feature_extractor.conv_layers[layer_id].conv.weight.data.shape, (
                f'{full_name} has size {value.shape}, but'
                f' {feature_extractor.conv_layers[layer_id].conv.weight.data.shape} was found.'
            )
            SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	value
            logger.info(f'Feat extract conv layer {layer_id} was initialized from {full_name}.'							)
    elif (type_id == 2 and not use_group_norm) or (type_id == 2 and layer_id == 0 and use_group_norm):
        if "bias" in name:
            assert value.shape == feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape, (
                f'{full_name} has size {value.shape}, but {feature_extractor[layer_id].layer_norm.bias.data.shape} was'
                " found."
            )
            SCREAMING_SNAKE_CASE_   :						Optional[int]							=	value
            logger.info(f'Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'							)
        elif "weight" in name:
            assert value.shape == feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape, (
                f'{full_name} has size {value.shape}, but'
                f' {feature_extractor[layer_id].layer_norm.weight.data.shape} was found.'
            )
            SCREAMING_SNAKE_CASE_   :						int							=	value
            logger.info(f'Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'							)
    else:
        unused_weights.append(__a							)



@torch.no_grad()
def       _A    (__a					,       __a					,       __a=None							)				-> Optional[Any]:
    """simple docstring"""






    SCREAMING_SNAKE_CASE_   :						List[Any]							=	torch.load(__a							)
    SCREAMING_SNAKE_CASE_   :						Any							=	WavLMConfigOrig(checkpoint['''cfg''']							)
    SCREAMING_SNAKE_CASE_   :						int							=	WavLMOrig(__a							)
    model.load_state_dict(checkpoint['''model''']							)
    model.eval()

    if config_path is not None:
        SCREAMING_SNAKE_CASE_   :						Optional[int]							=	WavLMConfig.from_pretrained(__a							)
    else:
        SCREAMING_SNAKE_CASE_   :						List[Any]							=	WavLMConfig()

    SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	WavLMModel(__a							)

    recursively_load_weights(__a					,       __a							)

    hf_wavlm.save_pretrained(__a							)


if __name__ == "__main__":
  UpperCAmelCase_     :				int         =				argparse.ArgumentParser()
  parser.add_argument("""--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the output PyTorch model.""")
  parser.add_argument("""--checkpoint_path""", default=None, type=str, help="""Path to fairseq checkpoint""")
  parser.add_argument("""--config_path""", default=None, type=str, help="""Path to hf config.json of model to convert""")
  UpperCAmelCase_     :				int         =				parser.parse_args()
  convert_wavlm_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path)










"""simple docstring"""






import unittest

from transformers import AutoTokenizer, FalconConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers import (
      FalconForCausalLM,
      FalconForQuestionAnswering,
      FalconForSequenceClassification,
      FalconForTokenClassification,
      FalconModel,
  )




class 							lowerCAmelCase__       :






      '''simple docstring'''

      def __init__(		self  :						Any		,			lowercase_  :						Dict		,			lowercase_  :						str=3		,			lowercase_  :						Dict=7		,			lowercase_  :						Any=True		,			lowercase_  :						List[Any]=True		,			lowercase_  :						Union[str, Any]=False		,			lowercase_  :						Optional[int]=True		,			lowercase_  :						Union[str, Any]=99		,			lowercase_  :						Dict=32		,			lowercase_  :						Union[str, Any]=5		,			lowercase_  :						Optional[Any]=4		,			lowercase_  :						List[str]=37		,			lowercase_  :						int="gelu"		,			lowercase_  :						int=0.1		,			lowercase_  :						List[str]=0.1		,			lowercase_  :						Any=512		,			lowercase_  :						List[Any]=16		,			lowercase_  :						Union[str, Any]=2		,			lowercase_  :						Dict=0.02		,			lowercase_  :						Optional[int]=3		,			lowercase_  :						Dict=4		,			lowercase_  :						Optional[int]=None		,			):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						List[Any]							=	parent
          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	batch_size
          SCREAMING_SNAKE_CASE_   :						List[str]							=	seq_length
          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	is_training
          SCREAMING_SNAKE_CASE_   :						Dict							=	use_input_mask
          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	use_token_type_ids
          SCREAMING_SNAKE_CASE_   :						str							=	use_labels
          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	vocab_size
          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	hidden_size
          SCREAMING_SNAKE_CASE_   :						str							=	num_hidden_layers
          SCREAMING_SNAKE_CASE_   :						List[Any]							=	num_attention_heads
          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	intermediate_size
          SCREAMING_SNAKE_CASE_   :						int							=	hidden_act
          SCREAMING_SNAKE_CASE_   :						Dict							=	hidden_dropout_prob
          SCREAMING_SNAKE_CASE_   :						str							=	attention_probs_dropout_prob
          SCREAMING_SNAKE_CASE_   :						str							=	max_position_embeddings
          SCREAMING_SNAKE_CASE_   :						Dict							=	type_vocab_size
          SCREAMING_SNAKE_CASE_   :						str							=	type_sequence_label_size
          SCREAMING_SNAKE_CASE_   :						List[Any]							=	initializer_range
          SCREAMING_SNAKE_CASE_   :						int							=	num_labels
          SCREAMING_SNAKE_CASE_   :						Dict							=	num_choices
          SCREAMING_SNAKE_CASE_   :						Tuple							=	scope
      def _SCREAMING_SNAKE_CASE				(		self  :						List[Any]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	ids_tensor([self.batch_size, self.seq_length]		,			self.vocab_size)

          SCREAMING_SNAKE_CASE_   :						Tuple							=	None
          if self.use_input_mask:
              SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	random_attention_mask([self.batch_size, self.seq_length])

          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	None

          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	None
          SCREAMING_SNAKE_CASE_   :						Dict							=	None
          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	None
          if self.use_labels:
              SCREAMING_SNAKE_CASE_   :						List[Any]							=	ids_tensor([self.batch_size]		,			self.type_sequence_label_size)
              SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	ids_tensor([self.batch_size, self.seq_length]		,			self.num_labels)
              SCREAMING_SNAKE_CASE_   :						str							=	ids_tensor([self.batch_size]		,			self.num_choices)

          SCREAMING_SNAKE_CASE_   :						Dict							=	self.get_config()

          return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
      def _SCREAMING_SNAKE_CASE				(		self  :						Union[str, Any]):




          '''simple docstring'''


          return FalconConfig(
              vocab_size=self.vocab_size		,			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		,			max_position_embeddings=self.max_position_embeddings		,			type_vocab_size=self.type_vocab_size		,			is_decoder=lowercase_		,			initializer_range=self.initializer_range		,			pad_token_id=1		,			new_decoder_architecture=lowercase_		,			)
      def _SCREAMING_SNAKE_CASE				(		self  :						Tuple		,			lowercase_  :						List[Any]		,			lowercase_  :						int		,			lowercase_  :						List[str]		,			lowercase_  :						Any		,			lowercase_  :						Optional[Any]		,			lowercase_  :						Union[str, Any]		,			lowercase_  :						Any):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Tuple							=	FalconModel(config=lowercase_)
          model.to(lowercase_)
          model.eval()
          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	model(lowercase_		,			attention_mask=lowercase_)
          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	model(lowercase_)
          self.parent.assertEqual(result.last_hidden_state.shape		,			(self.batch_size, self.seq_length, self.hidden_size))
      def _SCREAMING_SNAKE_CASE				(		self  :						Tuple		,			lowercase_  :						int		,			lowercase_  :						Dict		,			lowercase_  :						int		,			lowercase_  :						str		,			lowercase_  :						str		,			lowercase_  :						int		,			lowercase_  :						Tuple		,			lowercase_  :						Dict		,			lowercase_  :						int		,			):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Tuple							=	True
          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	FalconModel(lowercase_)
          model.to(lowercase_)
          model.eval()
          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	model(
              lowercase_		,			attention_mask=lowercase_		,			encoder_hidden_states=lowercase_		,			encoder_attention_mask=lowercase_		,			)
          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	model(
              lowercase_		,			attention_mask=lowercase_		,			encoder_hidden_states=lowercase_		,			)
          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	model(lowercase_		,			attention_mask=lowercase_)
          self.parent.assertEqual(result.last_hidden_state.shape		,			(self.batch_size, self.seq_length, self.hidden_size))
      def _SCREAMING_SNAKE_CASE				(		self  :						Dict		,			lowercase_  :						Union[str, Any]		,			lowercase_  :						Tuple		,			lowercase_  :						int		,			lowercase_  :						Optional[Any]		,			lowercase_  :						Dict		,			lowercase_  :						Dict		,			lowercase_  :						Dict		,			lowercase_  :						List[str]		,			lowercase_  :						List[str]		,			):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						int							=	FalconForCausalLM(config=lowercase_)
          model.to(lowercase_)
          model.eval()
          SCREAMING_SNAKE_CASE_   :						Any							=	model(lowercase_		,			attention_mask=lowercase_		,			labels=lowercase_)
          self.parent.assertEqual(result.logits.shape		,			(self.batch_size, self.seq_length, self.vocab_size))
      def _SCREAMING_SNAKE_CASE				(		self  :						List[str]		,			lowercase_  :						Any		,			lowercase_  :						int		,			lowercase_  :						str		,			lowercase_  :						Tuple		,			lowercase_  :						List[str]		,			lowercase_  :						Optional[int]		,			lowercase_  :						Tuple		,			lowercase_  :						Optional[int]		,			lowercase_  :						str		,			):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Dict							=	True
          SCREAMING_SNAKE_CASE_   :						List[Any]							=	True
          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	FalconForCausalLM(config=lowercase_)
          model.to(lowercase_)
          model.eval()

          # first forward pass
          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	model(
              lowercase_		,			attention_mask=lowercase_		,			encoder_hidden_states=lowercase_		,			encoder_attention_mask=lowercase_		,			use_cache=lowercase_		,			)
          SCREAMING_SNAKE_CASE_   :						List[Any]							=	outputs.past_key_values

          # create hypothetical multiple next token and extent to next_input_ids
          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	ids_tensor((self.batch_size, 3)		,			config.vocab_size)
          SCREAMING_SNAKE_CASE_   :						List[str]							=	ids_tensor((self.batch_size, 3)		,			vocab_size=2)

          # append to next input_ids and
          SCREAMING_SNAKE_CASE_   :						Dict							=	torch.cat([input_ids, next_tokens]		,			dim=-1)
          SCREAMING_SNAKE_CASE_   :						List[Any]							=	torch.cat([input_mask, next_mask]		,			dim=-1)

          SCREAMING_SNAKE_CASE_   :						List[str]							=	model(
              lowercase_		,			attention_mask=lowercase_		,			encoder_hidden_states=lowercase_		,			encoder_attention_mask=lowercase_		,			output_hidden_states=lowercase_		,			)['''hidden_states'''][0]
          SCREAMING_SNAKE_CASE_   :						Any							=	model(
              lowercase_		,			attention_mask=lowercase_		,			encoder_hidden_states=lowercase_		,			encoder_attention_mask=lowercase_		,			past_key_values=lowercase_		,			output_hidden_states=lowercase_		,			)['''hidden_states'''][0]

          # select random slice
          SCREAMING_SNAKE_CASE_   :						int							=	ids_tensor((1,)		,			output_from_past.shape[-1]).item()
          SCREAMING_SNAKE_CASE_   :						List[Any]							=	output_from_no_past[:, -3:, random_slice_idx].detach()
          SCREAMING_SNAKE_CASE_   :						Dict							=	output_from_past[:, :, random_slice_idx].detach()

          self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1])

          # test that outputs are equal for slice
          self.parent.assertTrue(torch.allclose(lowercase_		,			lowercase_		,			atol=1e-3))



      def _SCREAMING_SNAKE_CASE				(		self  :						Union[str, Any]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Tuple							=	self.prepare_config_and_inputs()
          (
              (
              SCREAMING_SNAKE_CASE_
          )		,    (
              SCREAMING_SNAKE_CASE_
          )		,    (
              SCREAMING_SNAKE_CASE_
          )		,    (
              SCREAMING_SNAKE_CASE_
          )		,    (
              SCREAMING_SNAKE_CASE_
          )		,    (
              SCREAMING_SNAKE_CASE_
          )		,    (
              SCREAMING_SNAKE_CASE_
          )		,    
          )   :						Tuple							=	config_and_inputs
          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	{'''input_ids''': input_ids, '''attention_mask''': input_mask}
          return config, inputs_dict




@require_torch
class 							lowerCAmelCase__       (    UpperCAmelCase__       ,	UpperCAmelCase__       ,	UpperCAmelCase__       ,	unittest.TestCase				):






      '''simple docstring'''

      __UpperCamelCase							=						(
          (
              FalconModel,
              FalconForCausalLM,
              FalconForSequenceClassification,
              FalconForTokenClassification,
              FalconForQuestionAnswering,
          )
          if is_torch_available()
          else ()
      )
      __UpperCamelCase							=						(FalconForCausalLM,) if is_torch_available() else ()
      __UpperCamelCase							=						(
          {
              "feature-extraction": FalconModel,
              "text-classification": FalconForSequenceClassification,
              "text-generation": FalconForCausalLM,
              "question-answering": FalconForQuestionAnswering,
              "token-classification": FalconForTokenClassification,
              "zero-shot": FalconForSequenceClassification,
          }
          if is_torch_available()
          else {}
      )
      __UpperCamelCase							=						False
      __UpperCamelCase							=						False
      def _SCREAMING_SNAKE_CASE				(		self  :						Optional[int]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						str							=	FalconModelTester(self)
          SCREAMING_SNAKE_CASE_   :						List[Any]							=	ConfigTester(self		,			config_class=lowercase_		,			hidden_size=37)
      def _SCREAMING_SNAKE_CASE				(		self  :						Tuple):




          '''simple docstring'''


          self.config_tester.run_common_tests()
      def _SCREAMING_SNAKE_CASE				(		self  :						int):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						List[Any]							=	self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*lowercase_)
      def _SCREAMING_SNAKE_CASE				(		self  :						int):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_		,    *SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	self.model_tester.prepare_config_and_inputs()
          for alibi in [True, False]:
              SCREAMING_SNAKE_CASE_   :						Optional[int]							=	alibi
              self.model_tester.create_and_check_model(lowercase_		,			*lowercase_)
      def _SCREAMING_SNAKE_CASE				(		self  :						Tuple):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_		,    SCREAMING_SNAKE_CASE_   :						List[str]							=	self.model_tester.prepare_config_and_inputs_for_common()
          SCREAMING_SNAKE_CASE_   :						Tuple							=	3
          SCREAMING_SNAKE_CASE_   :						Tuple							=	input_dict['''input_ids''']
          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	input_ids.ne(1).to(lowercase_)
          SCREAMING_SNAKE_CASE_   :						Dict							=	ids_tensor([self.model_tester.batch_size]		,			self.model_tester.type_sequence_label_size)
          SCREAMING_SNAKE_CASE_   :						Tuple							=	FalconForSequenceClassification(lowercase_)
          model.to(lowercase_)
          model.eval()
          SCREAMING_SNAKE_CASE_   :						List[str]							=	model(lowercase_		,			attention_mask=lowercase_		,			labels=lowercase_)
          self.assertEqual(result.logits.shape		,			(self.model_tester.batch_size, self.model_tester.num_labels))
      def _SCREAMING_SNAKE_CASE				(		self  :						int):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_		,    SCREAMING_SNAKE_CASE_   :						int							=	self.model_tester.prepare_config_and_inputs_for_common()
          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	3
          SCREAMING_SNAKE_CASE_   :						int							=	'''single_label_classification'''
          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	input_dict['''input_ids''']
          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	input_ids.ne(1).to(lowercase_)
          SCREAMING_SNAKE_CASE_   :						Dict							=	ids_tensor([self.model_tester.batch_size]		,			self.model_tester.type_sequence_label_size)
          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	FalconForSequenceClassification(lowercase_)
          model.to(lowercase_)
          model.eval()
          SCREAMING_SNAKE_CASE_   :						List[str]							=	model(lowercase_		,			attention_mask=lowercase_		,			labels=lowercase_)
          self.assertEqual(result.logits.shape		,			(self.model_tester.batch_size, self.model_tester.num_labels))
      def _SCREAMING_SNAKE_CASE				(		self  :						Dict):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_		,    SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	self.model_tester.prepare_config_and_inputs_for_common()
          SCREAMING_SNAKE_CASE_   :						List[Any]							=	input_dict['''input_ids''']
          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	FalconForCausalLM(lowercase_)
          model.to(lowercase_)
          model.eval()
          SCREAMING_SNAKE_CASE_   :						str							=	model(lowercase_		,			use_cache=lowercase_)
          SCREAMING_SNAKE_CASE_   :						Any							=	input_ids.shape[0]
          SCREAMING_SNAKE_CASE_   :						Dict							=	model._convert_to_rw_cache(result.past_key_values)
          SCREAMING_SNAKE_CASE_   :						int							=	model._convert_cache_to_standard_format(lowercase_		,			lowercase_)
          for layer in range(len(lowercase_)):
              for tensor_idx in range(2):
                  self.assertTrue(rw_cache[layer][tensor_idx].ndim == 3)
                  self.assertTrue(result.past_key_values[layer][tensor_idx].ndim == 4)
                  self.assertTrue(
                      torch.all(result.past_key_values[layer][tensor_idx] == standard_cache[layer][tensor_idx]))
      def _SCREAMING_SNAKE_CASE				(		self  :						Tuple):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_		,    SCREAMING_SNAKE_CASE_   :						Optional[int]							=	self.model_tester.prepare_config_and_inputs_for_common()
          SCREAMING_SNAKE_CASE_   :						int							=	3
          SCREAMING_SNAKE_CASE_   :						Tuple							=	'''multi_label_classification'''
          SCREAMING_SNAKE_CASE_   :						Tuple							=	input_dict['''input_ids''']
          SCREAMING_SNAKE_CASE_   :						List[str]							=	input_ids.ne(1).to(lowercase_)
          SCREAMING_SNAKE_CASE_   :						Tuple							=	ids_tensor(
              [self.model_tester.batch_size, config.num_labels]		,			self.model_tester.type_sequence_label_size).to(torch.float)
          SCREAMING_SNAKE_CASE_   :						int							=	FalconForSequenceClassification(lowercase_)
          model.to(lowercase_)
          model.eval()
          SCREAMING_SNAKE_CASE_   :						Tuple							=	model(lowercase_		,			attention_mask=lowercase_		,			labels=lowercase_)
          self.assertEqual(result.logits.shape		,			(self.model_tester.batch_size, self.model_tester.num_labels))



      def _SCREAMING_SNAKE_CASE				(		self  :						Union[str, Any]):




          '''simple docstring'''


          for model_class in self.all_generative_model_classes:
              SCREAMING_SNAKE_CASE_		,    SCREAMING_SNAKE_CASE_   :						Optional[int]							=	self.model_tester.prepare_config_and_inputs_for_common()

              # If it doesn't support cache, pass the test
              if not hasattr(lowercase_		,			'''use_cache'''):
                  return

              SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	model_class(lowercase_).to(lowercase_)
              if "use_cache" not in inputs:
                  SCREAMING_SNAKE_CASE_   :						Dict							=	True
              SCREAMING_SNAKE_CASE_   :						str							=	model(**lowercase_)

              # If "past_key_values" is not returned, pass the test (e.g. RWKV uses a different cache name and format)
              if "past_key_values" not in outputs:
                  return

              SCREAMING_SNAKE_CASE_   :						Tuple							=	(
                  getattr(lowercase_		,			'''decoder_layers'''		,			lowercase_)
                  or getattr(lowercase_		,			'''num_decoder_layers'''		,			lowercase_)
                  or config.num_hidden_layers
              )
              SCREAMING_SNAKE_CASE_   :						str							=	getattr(lowercase_		,			'''num_kv_heads'''		,			config.num_attention_heads)
              SCREAMING_SNAKE_CASE_   :						Tuple							=	getattr(lowercase_		,			'''d_model'''		,			config.hidden_size)
              SCREAMING_SNAKE_CASE_   :						Dict							=	embed_dim // num_attention_heads

              SCREAMING_SNAKE_CASE_   :						int							=	outputs['''past_key_values''']
              self.assertEqual(len(lowercase_)		,			lowercase_)

              SCREAMING_SNAKE_CASE_		,    SCREAMING_SNAKE_CASE_   :						List[Any]							=	inputs['''input_ids'''].shape
              for i in range(lowercase_):
                  if config.new_decoder_architecture:
                      SCREAMING_SNAKE_CASE_   :						Dict							=	config.num_attention_heads
                  elif config.multi_query:
                      SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	1
                  self.assertEqual(len(past_kv[0])		,			2)  # K V for the decoder = 2
                  self.assertEqual(
                      past_kv[i][0].shape		,			(batch_size, num_attention_heads, seq_length, per_head_embed_dim))
                  self.assertEqual(
                      past_kv[i][1].shape		,			(batch_size, num_attention_heads, seq_length, per_head_embed_dim))




@require_torch
class 							lowerCAmelCase__       (    unittest.TestCase				):






      '''simple docstring'''

      @slow
      def _SCREAMING_SNAKE_CASE				(		self  :						List[str]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Tuple							=	AutoTokenizer.from_pretrained('''Rocketknight1/falcon-rw-1b''')
          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	FalconForCausalLM.from_pretrained('''Rocketknight1/falcon-rw-1b''')
          model.eval()
          model.to(lowercase_)
          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	tokenizer('''My favorite food is'''		,			return_tensors='''pt''').to(lowercase_)

          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	(
              '''My favorite food is pizza. I love it so much that I have a pizza party every year for my birthday.'''
          )

          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	model.generate(**lowercase_		,			do_sample=lowercase_		,			max_new_tokens=19)
          SCREAMING_SNAKE_CASE_   :						List[Any]							=	tokenizer.batch_decode(lowercase_)[0]

          self.assertEqual(lowercase_		,			lowercase_)
      @slow
      def _SCREAMING_SNAKE_CASE				(		self  :						Union[str, Any]):




          '''simple docstring'''


          for repo in ["Rocketknight1/tiny-random-falcon-7b", "Rocketknight1/tiny-random-falcon-40b"]:
              SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	AutoTokenizer.from_pretrained(lowercase_)
              SCREAMING_SNAKE_CASE_   :						str							=	FalconForCausalLM.from_pretrained(lowercase_)
              model.eval()
              model.to(lowercase_)
              SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	tokenizer('''My favorite food is'''		,			return_tensors='''pt''').to(lowercase_)

              # We just test that these run without errors - the models are randomly initialized
              # and so the actual text outputs will be garbage
              model.generate(**lowercase_		,			do_sample=lowercase_		,			max_new_tokens=4)
              model.generate(**lowercase_		,			do_sample=lowercase_		,			max_new_tokens=4)
              model.generate(**lowercase_		,			num_beams=2		,			max_new_tokens=4)



      @slow
      def _SCREAMING_SNAKE_CASE				(		self  :						str):




          '''simple docstring'''


          with torch.no_grad():
              for repo in [
                  "Rocketknight1/falcon-rw-1b",
                  "Rocketknight1/tiny-random-falcon-7b",
                  "Rocketknight1/tiny-random-falcon-40b",
              ]:
                  SCREAMING_SNAKE_CASE_   :						List[str]							=	AutoTokenizer.from_pretrained(lowercase_)
                  SCREAMING_SNAKE_CASE_   :						Tuple							=	FalconForCausalLM.from_pretrained(lowercase_)
                  model.eval()
                  model.to(device=lowercase_)
                  SCREAMING_SNAKE_CASE_   :						int							=	tokenizer('''My favorite food is'''		,			return_tensors='''pt''').to(lowercase_)

                  # Test results are the same with and without cache
                  SCREAMING_SNAKE_CASE_   :						Any							=	model.generate(**lowercase_		,			do_sample=lowercase_		,			max_new_tokens=20		,			use_cache=lowercase_)
                  SCREAMING_SNAKE_CASE_   :						Dict							=	model.generate(**lowercase_		,			do_sample=lowercase_		,			max_new_tokens=20		,			use_cache=lowercase_)
                  self.assertTrue((outputs_cache - outputs_no_cache).sum().item() == 0)


