
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      :



  def __init__(				self       ,   A_       ,   A_=13       ,   A_=7       ,   A_=True       ,   A_=True       ,   A_=False       ,   A_=True       ,   A_=99       ,   A_=32       ,   A_=5       ,   A_=4       ,   A_=37       ,   A_="gelu"       ,   A_=0.1       ,   A_=0.1       ,   A_=512       ,   A_=16       ,   A_=2       ,   A_=0.02       ,   A_=3       ,   A_=4       ,   A_=None       ,   )						->      Union[str, Any]:
    """simple docstring"""
    UpperCamelCase						=  parent
    UpperCamelCase						=  batch_size
    UpperCamelCase						=  seq_length
    UpperCamelCase						=  is_training
    UpperCamelCase						=  use_input_mask
    UpperCamelCase						=  use_token_type_ids
    UpperCamelCase						=  use_labels
    UpperCamelCase						=  vocab_size
    UpperCamelCase						=  hidden_size
    UpperCamelCase						=  num_hidden_layers
    UpperCamelCase						=  num_attention_heads
    UpperCamelCase						=  intermediate_size
    UpperCamelCase						=  hidden_act
    UpperCamelCase						=  hidden_dropout_prob
    UpperCamelCase						=  attention_probs_dropout_prob
    UpperCamelCase						=  max_position_embeddings
    UpperCamelCase						=  type_vocab_size
    UpperCamelCase						=  type_sequence_label_size
    UpperCamelCase						=  initializer_range
    UpperCamelCase						=  num_labels
    UpperCamelCase						=  num_choices
    UpperCamelCase						=  scope



  def __UpperCamelCase				(				self )						->      List[str]:
    """simple docstring"""
    UpperCamelCase						=  ids_tensor([self.batch_size, self.seq_length]       ,   self.vocab_size )

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

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

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

    UpperCamelCase						=  self.get_config()

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



  def __UpperCamelCase				(				self )						->      List[Any]:
    """simple docstring"""
    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=A_       ,   initializer_range=self.initializer_range       ,   use_stable_embedding=A_       ,   )



  def __UpperCamelCase				(				self       ,   A_       ,   A_       ,   A_       ,   A_       ,   A_       ,   A_       ,   A_ )						->      Tuple:
    """simple docstring"""
    UpperCamelCase						=  OpenLlamaModel(config=A_ )
    model.to(A_ )
    model.eval()
    UpperCamelCase						=  model(A_       ,   attention_mask=A_ )
    UpperCamelCase						=  model(A_ )
    self.parent.assertEqual(result.last_hidden_state.shape       ,   (self.batch_size, self.seq_length, self.hidden_size) )



  def __UpperCamelCase				(				self       ,   A_       ,   A_       ,   A_       ,   A_       ,   A_       ,   A_       ,   A_       ,   A_       ,   A_       ,   )						->      Optional[Any]:
    """simple docstring"""
    UpperCamelCase						=  True
    UpperCamelCase						=  OpenLlamaModel(A_ )
    model.to(A_ )
    model.eval()
    UpperCamelCase						=  model(
        A_       ,   attention_mask=A_       ,   encoder_hidden_states=A_       ,   encoder_attention_mask=A_       ,   )
    UpperCamelCase						=  model(
        A_       ,   attention_mask=A_       ,   encoder_hidden_states=A_       ,   )
    UpperCamelCase						=  model(A_       ,   attention_mask=A_ )
    self.parent.assertEqual(result.last_hidden_state.shape       ,   (self.batch_size, self.seq_length, self.hidden_size) )



  def __UpperCamelCase				(				self       ,   A_       ,   A_       ,   A_       ,   A_       ,   A_       ,   A_       ,   A_       ,   A_       ,   A_       ,   )						->      Optional[int]:
    """simple docstring"""
    UpperCamelCase						=  OpenLlamaForCausalLM(config=A_ )
    model.to(A_ )
    model.eval()
    UpperCamelCase						=  model(A_       ,   attention_mask=A_       ,   labels=A_ )
    self.parent.assertEqual(result.logits.shape       ,   (self.batch_size, self.seq_length, self.vocab_size) )



  def __UpperCamelCase				(				self       ,   A_       ,   A_       ,   A_       ,   A_       ,   A_       ,   A_       ,   A_       ,   A_       ,   A_       ,   )						->      Optional[Any]:
    """simple docstring"""
    UpperCamelCase						=  True
    UpperCamelCase						=  True
    UpperCamelCase						=  OpenLlamaForCausalLM(config=A_ )
    model.to(A_ )
    model.eval()

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

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

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

    UpperCamelCase						=  model(
        A_       ,   attention_mask=A_       ,   encoder_hidden_states=A_       ,   encoder_attention_mask=A_       ,   output_hidden_states=A_       ,   )['hidden_states'][0]
    UpperCamelCase						=  model(
        A_       ,   attention_mask=A_       ,   encoder_hidden_states=A_       ,   encoder_attention_mask=A_       ,   past_key_values=A_       ,   output_hidden_states=A_       ,   )['hidden_states'][0]

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



  def __UpperCamelCase				(				self )						->      str:
    """simple docstring"""
    UpperCamelCase						=  self.prepare_config_and_inputs()
    (
        (
        UpperCamelCase						
    )			,						(
        UpperCamelCase						
    )			,						(
        UpperCamelCase						
    )			,						(
        UpperCamelCase						
    )			,						(
        UpperCamelCase						
    )			,						(
        UpperCamelCase						
    )			,						(
        UpperCamelCase						
    )			,												
    )						=  config_and_inputs
    UpperCamelCase						=  {'input_ids': input_ids, 'attention_mask': input_mask}
    return config, inputs_dict






@require_torch
class        lowercase      (	_SCREAMING_SNAKE_CASE						,     _SCREAMING_SNAKE_CASE						,     _SCREAMING_SNAKE_CASE						,     unittest.TestCase     ):
  __lowercase : Optional[Any]			  =       (
      (OpenLlamaModel, OpenLlamaForCausalLM, OpenLlamaForSequenceClassification) if is_torch_available() else ()
  )
  __lowercase : str			  =       (OpenLlamaForCausalLM,) if is_torch_available() else ()
  __lowercase : List[str]			  =       (
      {
          "feature-extraction": OpenLlamaModel,
          "text-classification": OpenLlamaForSequenceClassification,
          "text-generation": OpenLlamaForCausalLM,
          "zero-shot": OpenLlamaForSequenceClassification,
      }
      if is_torch_available()
      else {}
  )
  __lowercase : Any			  =       False
  __lowercase : Dict			  =       False



  def __UpperCamelCase				(				self )						->      Any:
    """simple docstring"""
    UpperCamelCase						=  OpenLlamaModelTester(self )
    UpperCamelCase						=  ConfigTester(self       ,   config_class=A_       ,   hidden_size=37 )



  def __UpperCamelCase				(				self )						->      Any:
    """simple docstring"""
    self.config_tester.run_common_tests()



  def __UpperCamelCase				(				self )						->      Dict:
    """simple docstring"""
    UpperCamelCase						=  self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_model(*A_ )



  def __UpperCamelCase				(				self )						->      List[str]:
    """simple docstring"""
    UpperCamelCase						=  self.model_tester.prepare_config_and_inputs()
    for type in ["absolute", "relative_key", "relative_key_query"]:
      UpperCamelCase						=  type
      self.model_tester.create_and_check_model(*A_ )



  def __UpperCamelCase				(				self )						->      Optional[int]:
    """simple docstring"""
    UpperCamelCase			,						UpperCamelCase						=  self.model_tester.prepare_config_and_inputs_for_common()
    UpperCamelCase						=  3
    UpperCamelCase						=  input_dict['input_ids']
    UpperCamelCase						=  input_ids.ne(1 ).to(A_ )
    UpperCamelCase						=  ids_tensor([self.model_tester.batch_size]       ,   self.model_tester.type_sequence_label_size )
    UpperCamelCase						=  OpenLlamaForSequenceClassification(A_ )
    model.to(A_ )
    model.eval()
    UpperCamelCase						=  model(A_       ,   attention_mask=A_       ,   labels=A_ )
    self.assertEqual(result.logits.shape       ,   (self.model_tester.batch_size, self.model_tester.num_labels) )



  def __UpperCamelCase				(				self )						->      Optional[Any]:
    """simple docstring"""
    UpperCamelCase			,						UpperCamelCase						=  self.model_tester.prepare_config_and_inputs_for_common()
    UpperCamelCase						=  3
    UpperCamelCase						=  'single_label_classification'
    UpperCamelCase						=  input_dict['input_ids']
    UpperCamelCase						=  input_ids.ne(1 ).to(A_ )
    UpperCamelCase						=  ids_tensor([self.model_tester.batch_size]       ,   self.model_tester.type_sequence_label_size )
    UpperCamelCase						=  OpenLlamaForSequenceClassification(A_ )
    model.to(A_ )
    model.eval()
    UpperCamelCase						=  model(A_       ,   attention_mask=A_       ,   labels=A_ )
    self.assertEqual(result.logits.shape       ,   (self.model_tester.batch_size, self.model_tester.num_labels) )



  def __UpperCamelCase				(				self )						->      Union[str, Any]:
    """simple docstring"""
    UpperCamelCase			,						UpperCamelCase						=  self.model_tester.prepare_config_and_inputs_for_common()
    UpperCamelCase						=  3
    UpperCamelCase						=  'multi_label_classification'
    UpperCamelCase						=  input_dict['input_ids']
    UpperCamelCase						=  input_ids.ne(1 ).to(A_ )
    UpperCamelCase						=  ids_tensor(
        [self.model_tester.batch_size, config.num_labels]       ,   self.model_tester.type_sequence_label_size ).to(torch.float )
    UpperCamelCase						=  OpenLlamaForSequenceClassification(A_ )
    model.to(A_ )
    model.eval()
    UpperCamelCase						=  model(A_       ,   attention_mask=A_       ,   labels=A_ )
    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 __UpperCamelCase				(				self )						->      Any:
    """simple docstring"""
    pass



  @parameterized.expand([('linear',), ('dynamic',)] )
  def __UpperCamelCase				(				self       ,   A_ )						->      Tuple:
    """simple docstring"""
    UpperCamelCase			,						UpperCamelCase						=  self.model_tester.prepare_config_and_inputs_for_common()
    UpperCamelCase						=  ids_tensor([1, 10]       ,   config.vocab_size )
    UpperCamelCase						=  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
    UpperCamelCase						=  OpenLlamaModel(A_ )
    original_model.to(A_ )
    original_model.eval()
    UpperCamelCase						=  original_model(A_ ).last_hidden_state
    UpperCamelCase						=  original_model(A_ ).last_hidden_state

    set_seed(42 )  # Fixed seed at init time so the two models get the same random weights
    UpperCamelCase						=  {'type': scaling_type, 'factor': 10.0}
    UpperCamelCase						=  OpenLlamaModel(A_ )
    scaled_model.to(A_ )
    scaled_model.eval()
    UpperCamelCase						=  scaled_model(A_ ).last_hidden_state
    UpperCamelCase						=  scaled_model(A_ ).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(A_       ,   A_       ,   atol=1e-5 ) )
    else:
      self.assertFalse(torch.allclose(A_       ,   A_       ,   atol=1e-5 ) )

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


import pytest

import datasets.config
from datasets.utils.info_utils import is_small_dataset


@pytest.mark.parametrize('dataset_size'					, [None, 400 * 2**20, 600 * 2**20]				)
@pytest.mark.parametrize('input_in_memory_max_size'					, ['default', 0, 100 * 2**20, 900 * 2**20]				)
def    A (      lowercase					, lowercase					, lowercase				)     ->						Union[str, Any]:
		'''simple docstring'''

		if input_in_memory_max_size != "default":
				monkeypatch.setattr(datasets.config					, 'IN_MEMORY_MAX_SIZE'					, lowercase				)
		UpperCamelCase						=  datasets.config.IN_MEMORY_MAX_SIZE
		if input_in_memory_max_size == "default":
				assert in_memory_max_size == 0
		else:
				assert in_memory_max_size == input_in_memory_max_size
		if dataset_size and in_memory_max_size:
				UpperCamelCase						=  dataset_size < in_memory_max_size
		else:
				UpperCamelCase						=  False
		UpperCamelCase						=  is_small_dataset(lowercase				)
		assert result == expected
