'''simple docstring'''

import unittest

from parameterized import parameterized

from transformers import OpenLlamaConfig, is_torch_available, set_seed
from transformers.testing_utils import require_torch, 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 OpenLlamaForCausalLM, OpenLlamaForSequenceClassification, OpenLlamaModel







class 				_lowercase     :


 '''simple docstring'''






 def __init__(      self   :	List[Any]      ,       SCREAMING_SNAKE_CASE__   :	Any      ,       SCREAMING_SNAKE_CASE__   :	str=13      ,       SCREAMING_SNAKE_CASE__   :	Optional[Any]=7      ,       SCREAMING_SNAKE_CASE__   :	Any=True      ,       SCREAMING_SNAKE_CASE__   :	List[str]=True      ,       SCREAMING_SNAKE_CASE__   :	Union[str, Any]=False      ,       SCREAMING_SNAKE_CASE__   :	Tuple=True      ,       SCREAMING_SNAKE_CASE__   :	List[Any]=99      ,       SCREAMING_SNAKE_CASE__   :	Tuple=32      ,       SCREAMING_SNAKE_CASE__   :	List[Any]=5      ,       SCREAMING_SNAKE_CASE__   :	int=4      ,       SCREAMING_SNAKE_CASE__   :	int=37      ,       SCREAMING_SNAKE_CASE__   :	Optional[Any]="gelu"      ,       SCREAMING_SNAKE_CASE__   :	int=0.1      ,       SCREAMING_SNAKE_CASE__   :	Union[str, Any]=0.1      ,       SCREAMING_SNAKE_CASE__   :	Dict=5_12      ,       SCREAMING_SNAKE_CASE__   :	Optional[Any]=16      ,       SCREAMING_SNAKE_CASE__   :	Optional[int]=2      ,       SCREAMING_SNAKE_CASE__   :	Tuple=0.0_2      ,       SCREAMING_SNAKE_CASE__   :	Optional[int]=3      ,       SCREAMING_SNAKE_CASE__   :	List[Any]=4      ,       SCREAMING_SNAKE_CASE__   :	Dict=None      ,       )  ->    List[Any]:
  __lowerCAmelCase											=						parent
  __lowerCAmelCase											=						batch_size
  __lowerCAmelCase											=						seq_length
  __lowerCAmelCase											=						is_training
  __lowerCAmelCase											=						use_input_mask
  __lowerCAmelCase											=						use_token_type_ids
  __lowerCAmelCase											=						use_labels
  __lowerCAmelCase											=						vocab_size
  __lowerCAmelCase											=						hidden_size
  __lowerCAmelCase											=						num_hidden_layers
  __lowerCAmelCase											=						num_attention_heads
  __lowerCAmelCase											=						intermediate_size
  __lowerCAmelCase											=						hidden_act
  __lowerCAmelCase											=						hidden_dropout_prob
  __lowerCAmelCase											=						attention_probs_dropout_prob
  __lowerCAmelCase											=						max_position_embeddings
  __lowerCAmelCase											=						type_vocab_size
  __lowerCAmelCase											=						type_sequence_label_size
  __lowerCAmelCase											=						initializer_range
  __lowerCAmelCase											=						num_labels
  __lowerCAmelCase											=						num_choices
  __lowerCAmelCase											=						scope
 def a					(      self   :	int							)  ->    List[str]:
  __lowerCAmelCase											=						ids_tensor([self.batch_size, self.seq_length]      ,       self.vocab_size							)

  __lowerCAmelCase											=						None
  if self.use_input_mask:
   __lowerCAmelCase											=						random_attention_mask([self.batch_size, self.seq_length]							)

  __lowerCAmelCase											=						None
  if self.use_token_type_ids:
   __lowerCAmelCase											=						ids_tensor([self.batch_size, self.seq_length]      ,       self.type_vocab_size							)

  __lowerCAmelCase											=						None
  __lowerCAmelCase											=						None
  __lowerCAmelCase											=						None
  if self.use_labels:
   __lowerCAmelCase											=						ids_tensor([self.batch_size]      ,       self.type_sequence_label_size							)
   __lowerCAmelCase											=						ids_tensor([self.batch_size, self.seq_length]      ,       self.num_labels							)
   __lowerCAmelCase											=						ids_tensor([self.batch_size]      ,       self.num_choices							)

  __lowerCAmelCase											=						self.get_config()

  return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
 def a					(      self   :	List[Any]							)  ->    List[str]:
  return OpenLlamaConfig(
      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=SCREAMING_SNAKE_CASE__      ,       initializer_range=self.initializer_range      ,       use_stable_embedding=SCREAMING_SNAKE_CASE__      ,       )
 def a					(      self   :	List[Any]      ,       SCREAMING_SNAKE_CASE__   :	Any      ,       SCREAMING_SNAKE_CASE__   :	Tuple      ,       SCREAMING_SNAKE_CASE__   :	str      ,       SCREAMING_SNAKE_CASE__   :	Any      ,       SCREAMING_SNAKE_CASE__   :	Optional[Any]      ,       SCREAMING_SNAKE_CASE__   :	int      ,       SCREAMING_SNAKE_CASE__   :	int							)  ->    Any:
  __lowerCAmelCase											=						OpenLlamaModel(config=SCREAMING_SNAKE_CASE__							)
  model.to(SCREAMING_SNAKE_CASE__							)
  model.eval()
  __lowerCAmelCase											=						model(SCREAMING_SNAKE_CASE__      ,       attention_mask=SCREAMING_SNAKE_CASE__							)
  __lowerCAmelCase											=						model(SCREAMING_SNAKE_CASE__							)
  self.parent.assertEqual(result.last_hidden_state.shape      ,       (self.batch_size, self.seq_length, self.hidden_size)							)
 def a					(      self   :	Tuple      ,       SCREAMING_SNAKE_CASE__   :	str      ,       SCREAMING_SNAKE_CASE__   :	Optional[int]      ,       SCREAMING_SNAKE_CASE__   :	Tuple      ,       SCREAMING_SNAKE_CASE__   :	Dict      ,       SCREAMING_SNAKE_CASE__   :	int      ,       SCREAMING_SNAKE_CASE__   :	Tuple      ,       SCREAMING_SNAKE_CASE__   :	Dict      ,       SCREAMING_SNAKE_CASE__   :	str      ,       SCREAMING_SNAKE_CASE__   :	Tuple      ,       )  ->    Any:
  __lowerCAmelCase											=						True
  __lowerCAmelCase											=						OpenLlamaModel(SCREAMING_SNAKE_CASE__							)
  model.to(SCREAMING_SNAKE_CASE__							)
  model.eval()
  __lowerCAmelCase											=						model(
      SCREAMING_SNAKE_CASE__      ,       attention_mask=SCREAMING_SNAKE_CASE__      ,       encoder_hidden_states=SCREAMING_SNAKE_CASE__      ,       encoder_attention_mask=SCREAMING_SNAKE_CASE__      ,       )
  __lowerCAmelCase											=						model(
      SCREAMING_SNAKE_CASE__      ,       attention_mask=SCREAMING_SNAKE_CASE__      ,       encoder_hidden_states=SCREAMING_SNAKE_CASE__      ,       )
  __lowerCAmelCase											=						model(SCREAMING_SNAKE_CASE__      ,       attention_mask=SCREAMING_SNAKE_CASE__							)
  self.parent.assertEqual(result.last_hidden_state.shape      ,       (self.batch_size, self.seq_length, self.hidden_size)							)
 def a					(      self   :	Dict      ,       SCREAMING_SNAKE_CASE__   :	Any      ,       SCREAMING_SNAKE_CASE__   :	List[Any]      ,       SCREAMING_SNAKE_CASE__   :	Tuple      ,       SCREAMING_SNAKE_CASE__   :	List[Any]      ,       SCREAMING_SNAKE_CASE__   :	Optional[Any]      ,       SCREAMING_SNAKE_CASE__   :	Tuple      ,       SCREAMING_SNAKE_CASE__   :	List[Any]      ,       SCREAMING_SNAKE_CASE__   :	List[str]      ,       SCREAMING_SNAKE_CASE__   :	Tuple      ,       )  ->    Any:
  __lowerCAmelCase											=						OpenLlamaForCausalLM(config=SCREAMING_SNAKE_CASE__							)
  model.to(SCREAMING_SNAKE_CASE__							)
  model.eval()
  __lowerCAmelCase											=						model(SCREAMING_SNAKE_CASE__      ,       attention_mask=SCREAMING_SNAKE_CASE__      ,       labels=SCREAMING_SNAKE_CASE__							)
  self.parent.assertEqual(result.logits.shape      ,       (self.batch_size, self.seq_length, self.vocab_size)							)
 def a					(      self   :	Union[str, Any]      ,       SCREAMING_SNAKE_CASE__   :	Any      ,       SCREAMING_SNAKE_CASE__   :	List[str]      ,       SCREAMING_SNAKE_CASE__   :	Dict      ,       SCREAMING_SNAKE_CASE__   :	Union[str, Any]      ,       SCREAMING_SNAKE_CASE__   :	int      ,       SCREAMING_SNAKE_CASE__   :	Tuple      ,       SCREAMING_SNAKE_CASE__   :	Any      ,       SCREAMING_SNAKE_CASE__   :	Tuple      ,       SCREAMING_SNAKE_CASE__   :	int      ,       )  ->    Union[str, Any]:
  __lowerCAmelCase											=						True
  __lowerCAmelCase											=						True
  __lowerCAmelCase											=						OpenLlamaForCausalLM(config=SCREAMING_SNAKE_CASE__							)
  model.to(SCREAMING_SNAKE_CASE__							)
  model.eval()

  # first forward pass
  __lowerCAmelCase											=						model(
      SCREAMING_SNAKE_CASE__      ,       attention_mask=SCREAMING_SNAKE_CASE__      ,       encoder_hidden_states=SCREAMING_SNAKE_CASE__      ,       encoder_attention_mask=SCREAMING_SNAKE_CASE__      ,       use_cache=SCREAMING_SNAKE_CASE__      ,       )
  __lowerCAmelCase											=						outputs.past_key_values

  # create hypothetical multiple next token and extent to next_input_ids
  __lowerCAmelCase											=						ids_tensor((self.batch_size, 3)      ,       config.vocab_size							)
  __lowerCAmelCase											=						ids_tensor((self.batch_size, 3)      ,       vocab_size=2							)

  # append to next input_ids and
  __lowerCAmelCase											=						torch.cat([input_ids, next_tokens]      ,       dim=-1							)
  __lowerCAmelCase											=						torch.cat([input_mask, next_mask]      ,       dim=-1							)

  __lowerCAmelCase											=						model(
      SCREAMING_SNAKE_CASE__      ,       attention_mask=SCREAMING_SNAKE_CASE__      ,       encoder_hidden_states=SCREAMING_SNAKE_CASE__      ,       encoder_attention_mask=SCREAMING_SNAKE_CASE__      ,       output_hidden_states=SCREAMING_SNAKE_CASE__      ,       )["""hidden_states"""][0]
  __lowerCAmelCase											=						model(
      SCREAMING_SNAKE_CASE__      ,       attention_mask=SCREAMING_SNAKE_CASE__      ,       encoder_hidden_states=SCREAMING_SNAKE_CASE__      ,       encoder_attention_mask=SCREAMING_SNAKE_CASE__      ,       past_key_values=SCREAMING_SNAKE_CASE__      ,       output_hidden_states=SCREAMING_SNAKE_CASE__      ,       )["""hidden_states"""][0]

  # select random slice
  __lowerCAmelCase											=						ids_tensor((1,)      ,       output_from_past.shape[-1]							).item()
  __lowerCAmelCase											=						output_from_no_past[:, -3:, random_slice_idx].detach()
  __lowerCAmelCase											=						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(SCREAMING_SNAKE_CASE__      ,       SCREAMING_SNAKE_CASE__      ,       atol=1e-3							)							)





 def a					(      self   :	List[Any]							)  ->    List[str]:
  __lowerCAmelCase											=						self.prepare_config_and_inputs()
  (
      (
      __lowerCAmelCase
  )							,       (
      __lowerCAmelCase
  )							,       (
      __lowerCAmelCase
  )							,       (
      __lowerCAmelCase
  )							,       (
      __lowerCAmelCase
  )							,       (
      __lowerCAmelCase
  )							,       (
      __lowerCAmelCase
  )							,       
  )											=						config_and_inputs
  __lowerCAmelCase											=						{"""input_ids""": input_ids, """attention_mask""": input_mask}
  return config, inputs_dict







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


 '''simple docstring'''






 _SCREAMING_SNAKE_CASE     :       Any      	=					(
     (OpenLlamaModel, OpenLlamaForCausalLM, OpenLlamaForSequenceClassification) if is_torch_available() else ()
 )
 _SCREAMING_SNAKE_CASE     :       Optional[int]      	=					(OpenLlamaForCausalLM,) if is_torch_available() else ()
 _SCREAMING_SNAKE_CASE     :       int      	=					(
     {
         """feature-extraction""": OpenLlamaModel,
         """text-classification""": OpenLlamaForSequenceClassification,
         """text-generation""": OpenLlamaForCausalLM,
         """zero-shot""": OpenLlamaForSequenceClassification,
     }
     if is_torch_available()
     else {}
 )
 _SCREAMING_SNAKE_CASE     :       str      	=					False
 _SCREAMING_SNAKE_CASE     :       List[str]      	=					False
 def a					(      self   :	Any							)  ->    Any:
  __lowerCAmelCase											=						OpenLlamaModelTester(self							)
  __lowerCAmelCase											=						ConfigTester(self      ,       config_class=SCREAMING_SNAKE_CASE__      ,       hidden_size=37							)
 def a					(      self   :	Optional[Any]							)  ->    Dict:
  self.config_tester.run_common_tests()
 def a					(      self   :	Union[str, Any]							)  ->    int:
  __lowerCAmelCase											=						self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE__							)
 def a					(      self   :	Optional[Any]							)  ->    Optional[int]:
  __lowerCAmelCase											=						self.model_tester.prepare_config_and_inputs()
  for type in ["absolute", "relative_key", "relative_key_query"]:
   __lowerCAmelCase											=						type
   self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE__							)
 def a					(      self   :	Optional[Any]							)  ->    Any:
  __lowerCAmelCase							,       __lowerCAmelCase											=						self.model_tester.prepare_config_and_inputs_for_common()
  __lowerCAmelCase											=						3
  __lowerCAmelCase											=						input_dict["""input_ids"""]
  __lowerCAmelCase											=						input_ids.ne(1							).to(SCREAMING_SNAKE_CASE__							)
  __lowerCAmelCase											=						ids_tensor([self.model_tester.batch_size]      ,       self.model_tester.type_sequence_label_size							)
  __lowerCAmelCase											=						OpenLlamaForSequenceClassification(SCREAMING_SNAKE_CASE__							)
  model.to(SCREAMING_SNAKE_CASE__							)
  model.eval()
  __lowerCAmelCase											=						model(SCREAMING_SNAKE_CASE__      ,       attention_mask=SCREAMING_SNAKE_CASE__      ,       labels=SCREAMING_SNAKE_CASE__							)
  self.assertEqual(result.logits.shape      ,       (self.model_tester.batch_size, self.model_tester.num_labels)							)
 def a					(      self   :	Tuple							)  ->    List[str]:
  __lowerCAmelCase							,       __lowerCAmelCase											=						self.model_tester.prepare_config_and_inputs_for_common()
  __lowerCAmelCase											=						3
  __lowerCAmelCase											=						"""single_label_classification"""
  __lowerCAmelCase											=						input_dict["""input_ids"""]
  __lowerCAmelCase											=						input_ids.ne(1							).to(SCREAMING_SNAKE_CASE__							)
  __lowerCAmelCase											=						ids_tensor([self.model_tester.batch_size]      ,       self.model_tester.type_sequence_label_size							)
  __lowerCAmelCase											=						OpenLlamaForSequenceClassification(SCREAMING_SNAKE_CASE__							)
  model.to(SCREAMING_SNAKE_CASE__							)
  model.eval()
  __lowerCAmelCase											=						model(SCREAMING_SNAKE_CASE__      ,       attention_mask=SCREAMING_SNAKE_CASE__      ,       labels=SCREAMING_SNAKE_CASE__							)
  self.assertEqual(result.logits.shape      ,       (self.model_tester.batch_size, self.model_tester.num_labels)							)
 def a					(      self   :	List[Any]							)  ->    Union[str, Any]:
  __lowerCAmelCase							,       __lowerCAmelCase											=						self.model_tester.prepare_config_and_inputs_for_common()
  __lowerCAmelCase											=						3
  __lowerCAmelCase											=						"""multi_label_classification"""
  __lowerCAmelCase											=						input_dict["""input_ids"""]
  __lowerCAmelCase											=						input_ids.ne(1							).to(SCREAMING_SNAKE_CASE__							)
  __lowerCAmelCase											=						ids_tensor(
      [self.model_tester.batch_size, config.num_labels]      ,       self.model_tester.type_sequence_label_size							).to(torch.float							)
  __lowerCAmelCase											=						OpenLlamaForSequenceClassification(SCREAMING_SNAKE_CASE__							)
  model.to(SCREAMING_SNAKE_CASE__							)
  model.eval()
  __lowerCAmelCase											=						model(SCREAMING_SNAKE_CASE__      ,       attention_mask=SCREAMING_SNAKE_CASE__      ,       labels=SCREAMING_SNAKE_CASE__							)
  self.assertEqual(result.logits.shape      ,       (self.model_tester.batch_size, self.model_tester.num_labels)							)
 @unittest.skip("""Open-Llama buffers include complex numbers, which breaks this test"""							)
 def a					(      self   :	Dict							)  ->    str:
  pass





 @parameterized.expand([("""linear""",), ("""dynamic""",)]							)
 def a					(      self   :	Tuple      ,       SCREAMING_SNAKE_CASE__   :	Optional[Any]							)  ->    Tuple:
  __lowerCAmelCase							,       __lowerCAmelCase											=						self.model_tester.prepare_config_and_inputs_for_common()
  __lowerCAmelCase											=						ids_tensor([1, 10]      ,       config.vocab_size							)
  __lowerCAmelCase											=						ids_tensor([1, int(config.max_position_embeddings * 1.5							)]      ,       config.vocab_size							)

  set_seed(42							)  # Fixed seed at init time so the two models get the same random weights
  __lowerCAmelCase											=						OpenLlamaModel(SCREAMING_SNAKE_CASE__							)
  original_model.to(SCREAMING_SNAKE_CASE__							)
  original_model.eval()
  __lowerCAmelCase											=						original_model(SCREAMING_SNAKE_CASE__							).last_hidden_state
  __lowerCAmelCase											=						original_model(SCREAMING_SNAKE_CASE__							).last_hidden_state

  set_seed(42							)  # Fixed seed at init time so the two models get the same random weights
  __lowerCAmelCase											=						{"""type""": scaling_type, """factor""": 1_0.0}
  __lowerCAmelCase											=						OpenLlamaModel(SCREAMING_SNAKE_CASE__							)
  scaled_model.to(SCREAMING_SNAKE_CASE__							)
  scaled_model.eval()
  __lowerCAmelCase											=						scaled_model(SCREAMING_SNAKE_CASE__							).last_hidden_state
  __lowerCAmelCase											=						scaled_model(SCREAMING_SNAKE_CASE__							).last_hidden_state

  # Dynamic scaling does not change the RoPE embeddings until it receives an input longer than the original
  # maximum sequence length, so the outputs for the short input should match.
  if scaling_type == "dynamic":
   self.assertTrue(torch.allclose(SCREAMING_SNAKE_CASE__      ,       SCREAMING_SNAKE_CASE__      ,       atol=1e-5							)							)
  else:
   self.assertFalse(torch.allclose(SCREAMING_SNAKE_CASE__      ,       SCREAMING_SNAKE_CASE__      ,       atol=1e-5							)							)

  # The output should be different for long inputs
  self.assertFalse(torch.allclose(SCREAMING_SNAKE_CASE__      ,       SCREAMING_SNAKE_CASE__      ,       atol=1e-5							)							)







'''simple docstring'''

import argparse

import torch
from torch import nn

from transformers import MaMaaaConfig, MaMaaaForConditionalGeneration







def      UpperCamelCase_						(					snake_case_       :       Any      )					->     Optional[Any]:






 '''simple docstring'''

 __lowerCAmelCase											=						[
     """encoder.version""",
     """decoder.version""",
     """model.encoder.version""",
     """model.decoder.version""",
     """decoder.output_projection.weight""",
     """_float_tensor""",
     """encoder.embed_positions._float_tensor""",
     """decoder.embed_positions._float_tensor""",
 ]
 for k in ignore_keys:
  state_dict.pop(snake_case_						,		snake_case_      )







def      UpperCamelCase_						(					snake_case_       :       Optional[Any]      )					->     List[Any]:






 '''simple docstring'''

 __lowerCAmelCase							,       __lowerCAmelCase											=						emb.weight.shape
 __lowerCAmelCase											=						nn.Linear(snake_case_						,		snake_case_						,		bias=snake_case_      )
 __lowerCAmelCase											=						emb.weight.data
 return lin_layer







def      UpperCamelCase_						(					snake_case_       :       Any      )					->     Any:






 '''simple docstring'''

 __lowerCAmelCase											=						torch.load(snake_case_						,		map_location="""cpu"""      )
 __lowerCAmelCase											=						mam_aaa["""args"""] or mam_aaa["""cfg"""]["""model"""]
 __lowerCAmelCase											=						mam_aaa["""model"""]
 remove_ignore_keys_(snake_case_      )
 __lowerCAmelCase											=						state_dict["""encoder.embed_tokens.weight"""].shape[0]

 __lowerCAmelCase											=						MaMaaaConfig(
     vocab_size=snake_case_						,		max_position_embeddings=10_24						,		encoder_layers=args.encoder_layers						,		decoder_layers=args.decoder_layers						,		encoder_attention_heads=args.encoder_attention_heads						,		decoder_attention_heads=args.decoder_attention_heads						,		encoder_ffn_dim=args.encoder_ffn_embed_dim						,		decoder_ffn_dim=args.decoder_ffn_embed_dim						,		d_model=args.encoder_embed_dim						,		encoder_layerdrop=args.encoder_layerdrop						,		decoder_layerdrop=args.decoder_layerdrop						,		dropout=args.dropout						,		attention_dropout=args.attention_dropout						,		activation_dropout=args.activation_dropout						,		activation_function="""relu"""						,		)

 __lowerCAmelCase											=						state_dict["""decoder.embed_tokens.weight"""]
 __lowerCAmelCase											=						MaMaaaForConditionalGeneration(snake_case_      )
 model.model.load_state_dict(snake_case_						,		strict=snake_case_      )
 __lowerCAmelCase											=						make_linear_from_emb(model.model.shared      )

 return model


if __name__ == "__main__":
 _A			:      Union[str, Any]					      =				argparse.ArgumentParser()
 # Required parameters
 parser.add_argument('''fairseq_path''', type=str, help='''path to a model.pt on local filesystem.''')
 parser.add_argument('''pytorch_dump_folder_path''', default=None, type=str, help='''Path to the output PyTorch model.''')
 _A			:      str					      =				parser.parse_args()
 _A			:      Optional[int]					      =				convert_fairseq_mamaaa_checkpoint_from_disk(args.fairseq_pathß)
 model.save_pretrained(args.pytorch_dump_folder_path)






