





'''simple docstring'''


import unittest

from parameterized import parameterized

from transformers import LlamaConfig, is_torch_available, set_seed
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 LlamaForCausalLM, LlamaForSequenceClassification, LlamaModel, LlamaTokenizer
class 				a_       :





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





        """simple docstring"""




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





     def 							lowercase__				(			self		: Tuple						):





        """simple docstring"""




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

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

        lowercase_      :Optional[int]         =			None
        if self.use_token_type_ids:
           lowercase_      :Dict         =			ids_tensor([self.batch_size, self.seq_length]     ,				self.type_vocab_size						)

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

        lowercase_      :str         =			self.get_config()

        return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels





     def 							lowercase__				(			self		: List[Any]						):





        """simple docstring"""




        return LlamaConfig(
            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     ,				)





     def 							lowercase__				(			self		: int     ,				lowercase		: Union[str, Any]     ,				lowercase		: List[str]     ,				lowercase		: Optional[int]     ,				lowercase		: List[str]     ,				lowercase		: Tuple     ,				lowercase		: str     ,				lowercase		: Union[str, Any]						):





        """simple docstring"""




        lowercase_      :List[Any]         =			LlamaModel(config=lowercase						)
        model.to(lowercase						)
        model.eval()
        lowercase_      :str         =			model(lowercase     ,				attention_mask=lowercase						)
        lowercase_      :Any         =			model(lowercase						)
        self.parent.assertEqual(result.last_hidden_state.shape     ,				(self.batch_size, self.seq_length, self.hidden_size)						)





     def 							lowercase__				(			self		: Optional[Any]     ,				lowercase		: List[str]     ,				lowercase		: Optional[Any]     ,				lowercase		: Optional[Any]     ,				lowercase		: Tuple     ,				lowercase		: Any     ,				lowercase		: str     ,				lowercase		: Optional[int]     ,				lowercase		: List[Any]     ,				lowercase		: str     ,				):





        """simple docstring"""




        lowercase_      :str         =			True
        lowercase_      :Tuple         =			LlamaModel(lowercase						)
        model.to(lowercase						)
        model.eval()
        lowercase_      :List[str]         =			model(
            lowercase     ,				attention_mask=lowercase     ,				encoder_hidden_states=lowercase     ,				encoder_attention_mask=lowercase     ,				)
        lowercase_      :Optional[int]         =			model(
            lowercase     ,				attention_mask=lowercase     ,				encoder_hidden_states=lowercase     ,				)
        lowercase_      :Optional[int]         =			model(lowercase     ,				attention_mask=lowercase						)
        self.parent.assertEqual(result.last_hidden_state.shape     ,				(self.batch_size, self.seq_length, self.hidden_size)						)





     def 							lowercase__				(			self		: Optional[Any]     ,				lowercase		: List[str]     ,				lowercase		: Any     ,				lowercase		: int     ,				lowercase		: List[str]     ,				lowercase		: Optional[Any]     ,				lowercase		: Optional[int]     ,				lowercase		: Dict     ,				lowercase		: List[str]     ,				lowercase		: str     ,				):





        """simple docstring"""




        lowercase_      :Union[str, Any]         =			LlamaForCausalLM(config=lowercase						)
        model.to(lowercase						)
        model.eval()
        lowercase_      :int         =			model(lowercase     ,				attention_mask=lowercase     ,				labels=lowercase						)
        self.parent.assertEqual(result.logits.shape     ,				(self.batch_size, self.seq_length, self.vocab_size)						)





     def 							lowercase__				(			self		: Optional[int]     ,				lowercase		: int     ,				lowercase		: Union[str, Any]     ,				lowercase		: Dict     ,				lowercase		: Any     ,				lowercase		: Tuple     ,				lowercase		: Any     ,				lowercase		: List[Any]     ,				lowercase		: Optional[Any]     ,				lowercase		: Any     ,				):





        """simple docstring"""




        lowercase_      :Dict         =			True
        lowercase_      :Any         =			True
        lowercase_      :int         =			LlamaForCausalLM(config=lowercase						)
        model.to(lowercase						)
        model.eval()

        # first forward pass
        lowercase_      :Tuple         =			model(
            lowercase     ,				attention_mask=lowercase     ,				encoder_hidden_states=lowercase     ,				encoder_attention_mask=lowercase     ,				use_cache=lowercase     ,				)
        lowercase_      :int         =			outputs.past_key_values

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

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

        lowercase_      :Dict         =			model(
            lowercase     ,				attention_mask=lowercase     ,				encoder_hidden_states=lowercase     ,				encoder_attention_mask=lowercase     ,				output_hidden_states=lowercase     ,				)["hidden_states"][0]
        lowercase_      :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
        lowercase_      :List[str]         =			ids_tensor((1,)     ,				output_from_past.shape[-1]						).item()
        lowercase_      :List[str]         =			output_from_no_past[:, -3:, random_slice_idx].detach()
        lowercase_      :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 							lowercase__				(			self		: Optional[int]						):





        """simple docstring"""




        lowercase_      :Any         =			self.prepare_config_and_inputs()
        (
            (
            lowercase_
        )      ,							(
            lowercase_
        )      ,							(
            lowercase_
        )      ,							(
            lowercase_
        )      ,							(
            lowercase_
        )      ,							(
            lowercase_
        )      ,							(
            lowercase_
        )      ,							
        )      :str         =			config_and_inputs
        lowercase_      :int         =			{"input_ids": input_ids, "attention_mask": input_mask}
        return config, inputs_dict







@require_torch
class 				a_       (    _lowerCAmelCase     ,						_lowerCAmelCase     ,						_lowerCAmelCase     ,						unittest.TestCase							):
     __A								=      (LlamaModel, LlamaForCausalLM, LlamaForSequenceClassification) if is_torch_available() else ()
     __A								=      (LlamaForCausalLM,) if is_torch_available() else ()
     __A								=      (
         {
             "feature-extraction": LlamaModel,
             "text-classification": LlamaForSequenceClassification,
             "text-generation": LlamaForCausalLM,
             "zero-shot": LlamaForSequenceClassification,
         }
         if is_torch_available()
         else {}
     )
     __A								=      False
     __A								=      False





     def 							lowercase__				(			self		: Tuple						):





        """simple docstring"""




        lowercase_      :List[str]         =			LlamaModelTester(self						)
        lowercase_      :List[Any]         =			ConfigTester(self     ,				config_class=lowercase     ,				hidden_size=37						)





     def 							lowercase__				(			self		: Optional[Any]						):





        """simple docstring"""




        self.config_tester.run_common_tests()





     def 							lowercase__				(			self		: Any						):





        """simple docstring"""




        lowercase_      :Optional[int]         =			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*lowercase						)





     def 							lowercase__				(			self		: List[str]						):





        """simple docstring"""




        lowercase_      :Optional[Any]         =			self.model_tester.prepare_config_and_inputs()
        for type in ["absolute", "relative_key", "relative_key_query"]:
           lowercase_      :Dict         =			type
           self.model_tester.create_and_check_model(*lowercase						)





     def 							lowercase__				(			self		: Any						):





        """simple docstring"""




        lowercase_      ,							lowercase_      :List[str]         =			self.model_tester.prepare_config_and_inputs_for_common()
        lowercase_      :str         =			3
        lowercase_      :str         =			input_dict["input_ids"]
        lowercase_      :Optional[Any]         =			input_ids.ne(1						).to(lowercase						)
        lowercase_      :List[str]         =			ids_tensor([self.model_tester.batch_size]     ,				self.model_tester.type_sequence_label_size						)
        lowercase_      :List[Any]         =			LlamaForSequenceClassification(lowercase						)
        model.to(lowercase						)
        model.eval()
        lowercase_      :List[Any]         =			model(lowercase     ,				attention_mask=lowercase     ,				labels=lowercase						)
        self.assertEqual(result.logits.shape     ,				(self.model_tester.batch_size, self.model_tester.num_labels)						)





     def 							lowercase__				(			self		: Optional[Any]						):





        """simple docstring"""




        lowercase_      ,							lowercase_      :Union[str, Any]         =			self.model_tester.prepare_config_and_inputs_for_common()
        lowercase_      :Optional[Any]         =			3
        lowercase_      :List[Any]         =			"single_label_classification"
        lowercase_      :Tuple         =			input_dict["input_ids"]
        lowercase_      :Optional[Any]         =			input_ids.ne(1						).to(lowercase						)
        lowercase_      :Optional[int]         =			ids_tensor([self.model_tester.batch_size]     ,				self.model_tester.type_sequence_label_size						)
        lowercase_      :str         =			LlamaForSequenceClassification(lowercase						)
        model.to(lowercase						)
        model.eval()
        lowercase_      :Any         =			model(lowercase     ,				attention_mask=lowercase     ,				labels=lowercase						)
        self.assertEqual(result.logits.shape     ,				(self.model_tester.batch_size, self.model_tester.num_labels)						)





     def 							lowercase__				(			self		: List[Any]						):





        """simple docstring"""




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





     @unittest.skip("LLaMA buffers include complex numbers, which breaks this test"						)
     def 							lowercase__				(			self		: Any						):





        """simple docstring"""




        pass







     @parameterized.expand([("linear",), ("dynamic",)]						)
     def 							lowercase__				(			self		: List[Any]     ,				lowercase		: Dict						):





        """simple docstring"""




        lowercase_      ,							lowercase_      :str         =			self.model_tester.prepare_config_and_inputs_for_common()
        lowercase_      :List[Any]         =			ids_tensor([1, 10]     ,				config.vocab_size						)
        lowercase_      :str         =			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
        lowercase_      :Any         =			LlamaModel(lowercase						)
        original_model.to(lowercase						)
        original_model.eval()
        lowercase_      :List[str]         =			original_model(lowercase						).last_hidden_state
        lowercase_      :Optional[Any]         =			original_model(lowercase						).last_hidden_state

        set_seed(42						)  # Fixed seed at init time so the two models get the same random weights
        lowercase_      :Optional[Any]         =			{"type": scaling_type, "factor": 10.0}
        lowercase_      :int         =			LlamaModel(lowercase						)
        scaled_model.to(lowercase						)
        scaled_model.eval()
        lowercase_      :Dict         =			scaled_model(lowercase						).last_hidden_state
        lowercase_      :Tuple         =			scaled_model(lowercase						).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(lowercase     ,				lowercase     ,				atol=1e-5						)						)
        else:
           self.assertFalse(torch.allclose(lowercase     ,				lowercase     ,				atol=1e-5						)						)

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







@require_torch
class 				a_       (    unittest.TestCase							):





     @unittest.skip("Logits are not exactly the same, once we fix the instabalities somehow, will update!"						)
     @slow
     def 							lowercase__				(			self		: Optional[Any]						):





        """simple docstring"""




        lowercase_      :str         =			[1, 306, 4_658, 278, 6_593, 310, 2_834, 338]
        lowercase_      :Union[str, Any]         =			LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf"     ,				device_map="auto"						)
        lowercase_      :Dict         =			model(torch.tensor([input_ids]						)						)
        # Expected mean on dim = -1
        lowercase_      :Optional[int]         =			torch.tensor([[-6.65_50, -4.12_27, -4.98_59, -3.24_06, 0.82_62, -3.00_33, 1.29_64, -3.36_99]]						)
        torch.testing.assert_close(out.mean(-1						)     ,				lowercase     ,				atol=1e-2     ,				rtol=1e-2						)
        # slicing logits[0, 0, 0:30]
        # fmt: off
        lowercase_      :Tuple         =			torch.tensor([-12.82_81, -7.44_53, -0.46_39, -8.06_25, -7.25_00, -8.00_00, -6.48_83, -7.76_95, -7.84_38, -7.03_12, -6.21_88, -7.13_28, -1.84_96, 1.99_61, -8.62_50, -6.72_27, -12.82_81, -6.94_92, -7.07_42, -7.78_52, -7.58_20, -7.90_62, -6.93_75, -7.98_05, -8.34_38, -8.15_62, -8.04_69, -7.62_50, -7.74_22, -7.33_98,]						)
        # fmt: on
        torch.testing.assert_close(out[0, 0, :30]     ,				lowercase     ,				atol=1e-5     ,				rtol=1e-5						)





     @unittest.skip("Logits are not exactly the same, once we fix the instabalities somehow, will update!"						)
     @slow
     def 							lowercase__				(			self		: str						):





        """simple docstring"""




        lowercase_      :Optional[Any]         =			[1, 306, 4_658, 278, 6_593, 310, 2_834, 338]
        lowercase_      :Union[str, Any]         =			LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-13b-hf"     ,				device_map="auto"						)
        lowercase_      :Any         =			model(torch.tensor(lowercase						)						)
        # Expected mean on dim = -1
        lowercase_      :str         =			torch.tensor([[-2.06_22, -1.27_94, -1.16_38, -0.97_88, -1.46_03, -1.02_38, -1.78_93, -1.44_11]]						)
        torch.testing.assert_close(out.mean(-1						)     ,				lowercase     ,				atol=1e-2     ,				rtol=1e-2						)
        # slicing logits[0, 0, 0:30]
        # fmt: off
        lowercase_      :Optional[Any]         =			torch.tensor([-8.14_06, -8.05_47, 2.74_61, -1.23_44, -0.14_48, -1.82_62, -1.00_20, -1.81_54, -1.68_95, -1.85_16, -2.35_74, -0.92_77, 3.75_98, 6.57_42, -1.29_98, -0.11_77, -8.14_06, -2.96_88, -2.91_99, -3.16_99, -3.52_54, -2.35_55, -2.79_88, -3.41_41, -2.82_62, -4.51_95, -3.33_79, -3.31_64, -2.78_32, -3.02_73]						)
        # fmt: on
        torch.testing.assert_close(out[0, 0, :30]     ,				lowercase     ,				atol=1e-5     ,				rtol=1e-5						)





     @unittest.skip("Logits are not exactly the same, once we fix the instabalities somehow, will update!"						)
     @slow
     def 							lowercase__				(			self		: Optional[int]						):





        """simple docstring"""




        lowercase_      :List[str]         =			[1, 306, 4_658, 278, 6_593, 310, 2_834, 338]
        lowercase_      :Dict         =			LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-13b-chat-hf"     ,				device_map="auto"						)
        lowercase_      :str         =			model(torch.tensor(lowercase						)						)
        # Expected mean on dim = -1
        lowercase_      :List[Any]         =			torch.tensor([[-0.85_62, -1.85_20, -0.75_51, -0.41_62, -1.51_61, -1.20_38, -2.48_23, -2.32_54]]						)
        torch.testing.assert_close(out.mean(-1						)     ,				lowercase     ,				atol=1e-2     ,				rtol=1e-2						)
        # slicing logits[0, 0, 0:30]
        # fmt: off
        lowercase_      :Union[str, Any]         =			torch.tensor([-2.22_27, 4.88_28, 0.90_23, -0.45_78, -0.78_71, -0.10_33, -0.62_21, -0.57_86, -0.78_03, -1.06_74, -1.29_20, -0.15_70, 0.80_08, 2.07_23, -0.94_97, 0.27_71, -2.22_27, -0.76_12, -1.43_46, -1.20_61, -1.64_26, -0.30_00, -0.71_39, -1.19_34, -1.86_91, -1.69_73, -1.59_47, -1.27_05, -0.35_23, -0.55_13]						)
        # fmt: on
        torch.testing.assert_close(out.mean(-1						)     ,				lowercase     ,				atol=1e-2     ,				rtol=1e-2						)





     @unittest.skip(
         "Logits are not exactly the same, once we fix the instabalities somehow, will update! Also it is gonna be a `too_slow` test"						)
     @slow
     def 							lowercase__				(			self		: List[str]						):





        """simple docstring"""




        lowercase_      :Dict         =			[1, 306, 4_658, 278, 6_593, 310, 2_834, 338]
        lowercase_      :Optional[int]         =			LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-70b-hf"     ,				device_map="auto"						)
        lowercase_      :Union[str, Any]         =			model(torch.tensor(lowercase						)						)

        lowercase_      :Tuple         =			torch.tensor(
            [[-4.23_27, -3.33_60, -4.66_65, -4.76_31, -1.81_80, -3.41_70, -1.42_11, -3.18_10]]     ,				dtype=torch.floataa						)
        torch.testing.assert_close(out.mean(-1						)     ,				lowercase     ,				atol=1e-2     ,				rtol=1e-2						)
        # fmt: off
        lowercase_      :str         =			torch.tensor([-9.49_22, -3.95_51, 1.79_98, -5.67_58, -5.10_55, -5.89_84, -4.83_20, -6.80_86, -6.53_91, -5.61_72, -5.58_20, -5.53_52, 1.78_81, 3.62_89, -6.51_17, -3.47_85, -9.50_00, -6.03_52, -6.81_25, -6.01_95, -6.68_36, -5.47_27, -6.28_12, -6.03_91, -7.33_98, -7.42_97, -7.48_44, -6.58_20, -5.87_89, -5.53_12]						)
        # fmt: on
        torch.testing.assert_close(out[0, 0, :30]     ,				lowercase     ,				atol=1e-5     ,				rtol=1e-5						)







     @unittest.skip("Model is curently gated"						)
     @slow
     def 							lowercase__				(			self		: Optional[int]						):





        """simple docstring"""




        lowercase_      :Union[str, Any]         =			"Simply put, the theory of relativity states that 1) the laws of physics are the same everywhere in the universe and 2) the passage of time and the length of objects can vary depending on the observer\'s frame of reference.\n\nThe first part of the theory, that the laws of physics are the same everywhere, is known as the \"princi"
        lowercase_      :Union[str, Any]         =			"Simply put, the theory of relativity states that "
        lowercase_      :Any         =			LlamaTokenizer.from_pretrained("meta-llama/Llama-2-13b-chat-hf"						)
        lowercase_      :Dict         =			tokenizer.encode(lowercase     ,				return_tensors="pt"						)
        lowercase_      :Optional[int]         =			LlamaForCausalLM.from_pretrained(
            "meta-llama/Llama-2-13b-chat-hf"     ,				device_map="sequential"     ,				use_safetensors=lowercase						)

        # greedy generation outputs
        lowercase_      :str         =			model.generate(lowercase     ,				max_new_tokens=64     ,				top_p=lowercase     ,				temperature=1     ,				do_sample=lowercase						)
        lowercase_      :List[Any]         =			tokenizer.decode(generated_ids[0]     ,				skip_special_tokens=lowercase						)
        self.assertEqual(lowercase     ,				lowercase						)









'''simple docstring'''


from typing import List, Optional, Union

import numpy as np
import tensorflow as tf

from .utils import logging


lowerCAmelCase :					Tuple 							=logging.get_logger(__name__)
def 							UpperCAmelCase_  (						__lowerCamelCase : Union[tf.Tensor, np.ndarray]			):
   if isinstance(__lowerCamelCase     ,np.ndarray			):
      return list(tensor.shape			)

   lowercase_      :Optional[int]         =			tf.shape(__lowerCamelCase			)

   if tensor.shape == tf.TensorShape(__lowerCamelCase			):
      return dynamic

   lowercase_      :Union[str, Any]         =			tensor.shape.as_list()

   return [dynamic[i] if s is None else s for i, s in enumerate(__lowerCamelCase			)]
def 							UpperCAmelCase_  (						__lowerCamelCase : tf.Tensor     ,__lowerCamelCase : Optional[int] = None     ,__lowerCamelCase : Optional[str] = None			):
   return tf.nn.softmax(logits=logits + 1e-9     ,axis=__lowerCamelCase     ,name=__lowerCamelCase			)
def 							UpperCAmelCase_  (						__lowerCamelCase : str     ,__lowerCamelCase : List[Any]     ,__lowerCamelCase : Any     ,__lowerCamelCase : List[str]=1e-5     ,__lowerCamelCase : List[str]=-1			):
   # This is a very simplified functional layernorm, designed to duplicate
   # the functionality of PyTorch nn.functional.layer_norm when this is needed to port
   # models in Transformers.

   if weight.shape.rank != 1 or bias.shape.rank != 1 or not isinstance(__lowerCamelCase     ,__lowerCamelCase			):
      raise NotImplementedError("Only 1D weight and bias tensors are supported for now, with only a single axis."			)

   # Get mean and variance on the axis to be normalized
   lowercase_      ,							lowercase_      :List[str]         =			tf.nn.moments(__lowerCamelCase     ,axes=[axis]     ,keepdims=__lowerCamelCase			)

   if axis != -1:
      # Reshape scale and weight to have the same rank as inputs, but with 1 dimensions
      # on every dimension except axis
      lowercase_      :Union[str, Any]         =			[1] * inputs.shape.rank
      lowercase_      :Optional[Any]         =			shape_list(__lowerCamelCase			)[axis]
      lowercase_      :List[str]         =			tf.reshape(__lowerCamelCase     ,__lowerCamelCase			)
      lowercase_      :Dict         =			tf.reshape(__lowerCamelCase     ,__lowerCamelCase			)

   # Compute layer normalization using the batch_normalization
   # function.
   lowercase_      :Union[str, Any]         =			tf.nn.batch_normalization(
       __lowerCamelCase     ,__lowerCamelCase     ,__lowerCamelCase     ,offset=__lowerCamelCase     ,scale=__lowerCamelCase     ,variance_epsilon=__lowerCamelCase     ,)
   return outputs
def 							UpperCAmelCase_  (						__lowerCamelCase : Optional[int]     ,__lowerCamelCase : Union[str, Any]=0     ,__lowerCamelCase : Dict=-1			):
   # Replicates the behavior of torch.flatten in TF

   # If end_dim or start_dim is negative, count them from the end
   if end_dim < 0:
      end_dim += input.shape.rank
   if start_dim < 0:
      start_dim += input.shape.rank

   if start_dim == end_dim:
      return input

   lowercase_      :Optional[int]         =			tf.shape(__lowerCamelCase			)
   lowercase_      :Optional[int]         =			tf.math.reduce_prod(in_shape[start_dim : end_dim + 1]			)
   lowercase_      :List[str]         =			tf.concat([in_shape[:start_dim], [flattened_dim], in_shape[end_dim + 1 :]]     ,axis=0			)
   return tf.reshape(__lowerCamelCase     ,__lowerCamelCase			)
def 							UpperCAmelCase_  (						__lowerCamelCase : tf.Tensor			):
   if not isinstance(__lowerCamelCase     ,tf.Tensor			):
      lowercase_      :str         =			tf.convert_to_tensor(__lowerCamelCase			)  # Catches stray NumPy inputs
   if encoder_attention_mask.shape.rank == 3:
      lowercase_      :List[Any]         =			encoder_attention_mask[:, None, :, :]
   if encoder_attention_mask.shape.rank == 2:
      lowercase_      :Optional[int]         =			encoder_attention_mask[:, None, None, :]
   # T5 has a mask that can compare sequence ids, we can simulate this here with this transposition
   # Cf. https://github.com/tensorflow/mesh/blob/8d2465e9bc93129b913b5ccc6a59aa97abd96ec6/mesh_tensorflow
   # /transformer/transformer_layers.py#L270
   # encoder_extended_attention_mask = (encoder_extended_attention_mask ==
   # encoder_extended_attention_mask.transpose(-1, -2))
   lowercase_      :str         =			(
       tf.cast(1     ,encoder_attention_mask.dtype			) - encoder_extended_attention_mask
   ) * encoder_extended_attention_mask.dtype.min

   return encoder_extended_attention_mask
def 							UpperCAmelCase_  (						__lowerCamelCase : tf.Tensor     ,__lowerCamelCase : int     ,__lowerCamelCase : str = "input_ids"			):
   tf.debugging.assert_less(
       __lowerCamelCase     ,tf.cast(__lowerCamelCase     ,dtype=tensor.dtype			)     ,message=(
           F'The maximum value of {tensor_name} ({tf.math.reduce_max(__lowerCamelCase			)}) must be smaller than the embedding '
           F'layer\'s input dimension ({embed_dim}). The likely cause is some problem at tokenization time.'
       )     ,)
def 							UpperCAmelCase_  (						__lowerCamelCase : List[str]     ,__lowerCamelCase : Tuple     ,__lowerCamelCase : Dict			):
   lowercase_      :int         =			6_45_12
   # Check that no item in `data` is larger than `HDF5_OBJECT_HEADER_LIMIT`
   # because in that case even chunking the array would not make the saving
   # possible.
   lowercase_      :Union[str, Any]         =			[x for x in data if len(__lowerCamelCase			) > HDF5_OBJECT_HEADER_LIMIT]

   # Expecting this to never be true.
   if bad_attributes:
      raise RuntimeError(
          "The following attributes cannot be saved to HDF5 file because "
          F'they are larger than {HDF5_OBJECT_HEADER_LIMIT} '
          F'bytes: {bad_attributes}'			)

   lowercase_      :Union[str, Any]         =			np.asarray(__lowerCamelCase			)

   lowercase_      :Optional[int]         =			1
   lowercase_      :int         =			np.array_split(__lowerCamelCase     ,__lowerCamelCase			)

   # This will never loop forever thanks to the test above.
   while any(x.nbytes > HDF5_OBJECT_HEADER_LIMIT for x in chunked_data			):
      num_chunks += 1
      lowercase_      :List[Any]         =			np.array_split(__lowerCamelCase     ,__lowerCamelCase			)

   if num_chunks > 1:
      for chunk_id, chunk_data in enumerate(__lowerCamelCase			):
         lowercase_      :int         =			chunk_data
   else:
      lowercase_      :Tuple         =			data
def 							UpperCAmelCase_  (						__lowerCamelCase : str     ,__lowerCamelCase : Tuple			):
   if name in group.attrs:
      lowercase_      :Optional[Any]         =			[n.decode("utf8"			) if hasattr(__lowerCamelCase     ,"decode"			) else n for n in group.attrs[name]]
   else:
      lowercase_      :List[str]         =			[]
      lowercase_      :str         =			0
      while "%s%d" % (name, chunk_id) in group.attrs:
         data.extend(
             [n.decode("utf8"			) if hasattr(__lowerCamelCase     ,"decode"			) else n for n in group.attrs["%s%d" % (name, chunk_id)]]			)
         chunk_id += 1
   return data


def 							UpperCAmelCase_  (						__lowerCamelCase : str			):

   def _expand_single_ad_tensor(__lowerCamelCase : Tuple			):
      if isinstance(__lowerCamelCase     ,tf.Tensor			) and t.shape.rank == 1:
         return tf.expand_dims(__lowerCamelCase     ,axis=-1			)
      return t

   return tf.nest.map_structure(_expand_single_ad_tensor     ,__lowerCamelCase			)


