

import logging
import os
import sys
from dataclasses import dataclass, field
from typing import Optional

import torch
from datasets import load_dataset
from torchvision.transforms import Compose, Lambda, Normalize, RandomHorizontalFlip, RandomResizedCrop, ToTensor
from torchvision.transforms.functional import InterpolationMode

import transformers
from transformers import (
    HfArgumentParser,
    Trainer,
    TrainingArguments,
    ViTImageProcessor,
    ViTMAEConfig,
    ViTMAEForPreTraining,
)
from transformers.trainer_utils import get_last_checkpoint
from transformers.utils import check_min_version, send_example_telemetry
from transformers.utils.versions import require_version
a_    :						List[str]    		=      logging.getLogger(__name__)

# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
check_min_version("4.31.0")

require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/image-pretraining/requirements.txt")




@dataclass
class      UpperCamelCase :
    __UpperCamelCase         =field(
        default="cifar10"	,			metadata={"help": "Name of a dataset from the datasets package"}   )
    __UpperCamelCase         =field(
        default=SCREAMING_SNAKE_CASE	,			metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}   )
    __UpperCamelCase         =field(
        default=SCREAMING_SNAKE_CASE	,			metadata={"help": "The column name of the images in the files."}   )
    __UpperCamelCase         =field(default=SCREAMING_SNAKE_CASE	,			metadata={"help": "A folder containing the training data."}   )
    __UpperCamelCase         =field(default=SCREAMING_SNAKE_CASE	,			metadata={"help": "A folder containing the validation data."}   )
    __UpperCamelCase         =field(
        default=0.15	,			metadata={"help": "Percent to split off of train for validation."}   )
    __UpperCamelCase         =field(
        default=SCREAMING_SNAKE_CASE	,			metadata={
            "help": (
                "For debugging purposes or quicker training, truncate the number of training examples to this "
                "value if set."
            )
        }	,			)
    __UpperCamelCase         =field(
        default=SCREAMING_SNAKE_CASE	,			metadata={
            "help": (
                "For debugging purposes or quicker training, truncate the number of evaluation examples to this "
                "value if set."
            )
        }	,			)

    def UpperCamelCase			(			self :      Optional[Any]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			{}
     if self.train_dir is not None:
      SCREAMING_SNAKE_CASE								=			self.train_dir
     if self.validation_dir is not None:
      SCREAMING_SNAKE_CASE								=			self.validation_dir
     SCREAMING_SNAKE_CASE								=			data_files if data_files else None




@dataclass
class      UpperCamelCase :
    __UpperCamelCase         =field(
        default=SCREAMING_SNAKE_CASE	,			metadata={
            "help": (
                "The model checkpoint for weights initialization.Don't set if you want to train a model from scratch."
            )
        }	,			)
    __UpperCamelCase         =field(
        default=SCREAMING_SNAKE_CASE	,			metadata={"help": "Pretrained config name or path if not the same as model_name_or_path"}   )
    __UpperCamelCase         =field(
        default=SCREAMING_SNAKE_CASE	,			metadata={
            "help": (
                "Override some existing default config settings when a model is trained from scratch. Example: "
                "n_embd=10,resid_pdrop=0.2,scale_attn_weights=false,summary_type=cls_index"
            )
        }	,			)
    __UpperCamelCase         =field(
        default=SCREAMING_SNAKE_CASE	,			metadata={"help": "Where do you want to store the pretrained models downloaded from s3"}   )
    __UpperCamelCase         =field(
        default="main"	,			metadata={"help": "The specific model version to use (can be a branch name, tag name or commit id)."}	,			)
    __UpperCamelCase         =field(default=SCREAMING_SNAKE_CASE	,			metadata={"help": "Name or path of preprocessor config."}   )
    __UpperCamelCase         =field(
        default=SCREAMING_SNAKE_CASE	,			metadata={
            "help": (
                "Will use the token generated when running `huggingface-cli login` (necessary to use this script "
                "with private models)."
            )
        }	,			)
    __UpperCamelCase         =field(
        default=0.75	,			metadata={"help": "The ratio of the number of masked tokens in the input sequence."}   )
    __UpperCamelCase         =field(
        default=SCREAMING_SNAKE_CASE	,			metadata={"help": "Whether or not to train with normalized pixel values as target."}   )




@dataclass
class      UpperCamelCase (     SCREAMING_SNAKE_CASE   ):
    __UpperCamelCase         =field(
        default=1e-3	,			metadata={"help": "Base learning rate: absolute_lr = base_lr * total_batch_size / 256."}   )


def        __lowerCAmelCase    (       _UpperCamelCase		:  Any	)		->							Dict:
 '''simple docstring'''


 SCREAMING_SNAKE_CASE								=			torch.stack([example['pixel_values'] for example in examples]	)
 return {"pixel_values": pixel_values}


def        __lowerCAmelCase    (       )		->							Tuple:
 '''simple docstring'''


 SCREAMING_SNAKE_CASE								=			HfArgumentParser((ModelArguments, DataTrainingArguments, CustomTrainingArguments)	)
 if len(sys.argv	) == 2 and sys.argv[1].endswith('.json'	):
  # If we pass only one argument to the script and it's the path to a json file,
  # let's parse it to get our arguments.
  SCREAMING_SNAKE_CASE						,	SCREAMING_SNAKE_CASE						,	SCREAMING_SNAKE_CASE								=			parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]	)	)
 else:
  SCREAMING_SNAKE_CASE						,	SCREAMING_SNAKE_CASE						,	SCREAMING_SNAKE_CASE								=			parser.parse_args_into_dataclasses()

 # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
 # information sent is the one passed as arguments along with your Python/PyTorch versions.
 send_example_telemetry('run_mae'		,							_UpperCamelCase		,							_UpperCamelCase	)

 # Setup logging
 logging.basicConfig(
     format='%(asctime)s - %(levelname)s - %(name)s - %(message)s'		,							datefmt='%m/%d/%Y %H:%M:%S'		,							handlers=[logging.StreamHandler(sys.stdout	)]		,							)

 if training_args.should_log:
  # The default of training_args.log_level is passive, so we set log level at info here to have that default.
  transformers.utils.logging.set_verbosity_info()

 SCREAMING_SNAKE_CASE								=			training_args.get_process_log_level()
 logger.setLevel(_UpperCamelCase	)
 transformers.utils.logging.set_verbosity(_UpperCamelCase	)
 transformers.utils.logging.enable_default_handler()
 transformers.utils.logging.enable_explicit_format()

 # Log on each process the small summary:
 logger.warning(
     f"""Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"""
     + f"""distributed training: {bool(training_args.local_rank != -1	)}, 16-bits training: {training_args.fpaa}"""	)
 logger.info(f"""Training/evaluation parameters {training_args}"""	)

 # Detecting last checkpoint.
 SCREAMING_SNAKE_CASE								=			None
 if os.path.isdir(training_args.output_dir	) and training_args.do_train and not training_args.overwrite_output_dir:
  SCREAMING_SNAKE_CASE								=			get_last_checkpoint(training_args.output_dir	)
  if last_checkpoint is None and len(os.listdir(training_args.output_dir	)	) > 0:
   raise ValueError(
       f"""Output directory ({training_args.output_dir}) already exists and is not empty. """
       'Use --overwrite_output_dir to overcome.'	)
  elif last_checkpoint is not None and training_args.resume_from_checkpoint is None:
   logger.info(
       f"""Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change """
       'the `--output_dir` or add `--overwrite_output_dir` to train from scratch.'	)

    # Initialize our dataset.
 SCREAMING_SNAKE_CASE								=			load_dataset(
     data_args.dataset_name		,							data_args.dataset_config_name		,							data_files=data_args.data_files		,							cache_dir=model_args.cache_dir		,							use_auth_token=True if model_args.use_auth_token else None		,							)

 # If we don't have a validation split, split off a percentage of train as validation.
 SCREAMING_SNAKE_CASE								=			None if 'validation' in ds.keys() else data_args.train_val_split
 if isinstance(data_args.train_val_split		,							_UpperCamelCase	) and data_args.train_val_split > 0.0:
  SCREAMING_SNAKE_CASE								=			ds['train'].train_test_split(data_args.train_val_split	)
  SCREAMING_SNAKE_CASE								=			split['train']
  SCREAMING_SNAKE_CASE								=			split['test']

 # Load pretrained model and image processor
 #
 # Distributed training:
 # The .from_pretrained methods guarantee that only one local process can concurrently
 # download model & vocab.
 SCREAMING_SNAKE_CASE								=			{
     'cache_dir': model_args.cache_dir,
     'revision': model_args.model_revision,
     'use_auth_token': True if model_args.use_auth_token else None,
 }
 if model_args.config_name:
  SCREAMING_SNAKE_CASE								=			ViTMAEConfig.from_pretrained(model_args.config_name		,							**_UpperCamelCase	)
 elif model_args.model_name_or_path:
  SCREAMING_SNAKE_CASE								=			ViTMAEConfig.from_pretrained(model_args.model_name_or_path		,							**_UpperCamelCase	)
 else:
  SCREAMING_SNAKE_CASE								=			ViTMAEConfig()
  logger.warning('You are instantiating a new config instance from scratch.'	)
  if model_args.config_overrides is not None:
   logger.info(f"""Overriding config: {model_args.config_overrides}"""	)
   config.update_from_string(model_args.config_overrides	)
   logger.info(f"""New config: {config}"""	)

    # adapt config
 config.update(
     {
         'mask_ratio': model_args.mask_ratio,
         'norm_pix_loss': model_args.norm_pix_loss,
     }	)

 # create image processor
 if model_args.image_processor_name:
  SCREAMING_SNAKE_CASE								=			ViTImageProcessor.from_pretrained(model_args.image_processor_name		,							**_UpperCamelCase	)
 elif model_args.model_name_or_path:
  SCREAMING_SNAKE_CASE								=			ViTImageProcessor.from_pretrained(model_args.model_name_or_path		,							**_UpperCamelCase	)
 else:
  SCREAMING_SNAKE_CASE								=			ViTImageProcessor()

 # create model
 if model_args.model_name_or_path:
  SCREAMING_SNAKE_CASE								=			ViTMAEForPreTraining.from_pretrained(
      model_args.model_name_or_path		,							from_tf=bool('.ckpt' in model_args.model_name_or_path	)		,							config=_UpperCamelCase		,							cache_dir=model_args.cache_dir		,							revision=model_args.model_revision		,							use_auth_token=True if model_args.use_auth_token else None		,							)
 else:
  logger.info('Training new model from scratch'	)
  SCREAMING_SNAKE_CASE								=			ViTMAEForPreTraining(_UpperCamelCase	)

 if training_args.do_train:
  SCREAMING_SNAKE_CASE								=			ds['train'].column_names
 else:
  SCREAMING_SNAKE_CASE								=			ds['validation'].column_names

 if data_args.image_column_name is not None:
  SCREAMING_SNAKE_CASE								=			data_args.image_column_name
 elif "image" in column_names:
  SCREAMING_SNAKE_CASE								=			'image'
 elif "img" in column_names:
  SCREAMING_SNAKE_CASE								=			'img'
 else:
  SCREAMING_SNAKE_CASE								=			column_names[0]

 # transformations as done in original MAE paper
 # source: https://github.com/facebookresearch/mae/blob/main/main_pretrain.py
 if "shortest_edge" in image_processor.size:
  SCREAMING_SNAKE_CASE								=			image_processor.size['shortest_edge']
 else:
  SCREAMING_SNAKE_CASE								=			(image_processor.size['height'], image_processor.size['width'])
 SCREAMING_SNAKE_CASE								=			Compose(
     [
         Lambda(lambda _UpperCamelCase	: img.convert('RGB'	) if img.mode != "RGB" else img	),
         RandomResizedCrop(_UpperCamelCase		,							scale=(0.2, 1.0)		,							interpolation=InterpolationMode.BICUBIC	),
         RandomHorizontalFlip(),
         ToTensor(),
         Normalize(mean=image_processor.image_mean		,							std=image_processor.image_std	),
     ]	)

 def preprocess_images(_UpperCamelCase		:  str	):
  SCREAMING_SNAKE_CASE								=			[transforms(_UpperCamelCase	) for image in examples[image_column_name]]
  return examples

 if training_args.do_train:
  if "train" not in ds:
   raise ValueError('--do_train requires a train dataset'	)
  if data_args.max_train_samples is not None:
   SCREAMING_SNAKE_CASE								=			ds['train'].shuffle(seed=training_args.seed	).select(range(data_args.max_train_samples	)	)
  # Set the training transforms
  ds["train"].set_transform(_UpperCamelCase	)

 if training_args.do_eval:
  if "validation" not in ds:
   raise ValueError('--do_eval requires a validation dataset'	)
  if data_args.max_eval_samples is not None:
   SCREAMING_SNAKE_CASE								=			(
       ds['validation'].shuffle(seed=training_args.seed	).select(range(data_args.max_eval_samples	)	)
   )
  # Set the validation transforms
  ds["validation"].set_transform(_UpperCamelCase	)

 # Compute absolute learning rate
 SCREAMING_SNAKE_CASE								=			(
     training_args.train_batch_size * training_args.gradient_accumulation_steps * training_args.world_size
 )
 if training_args.base_learning_rate is not None:
  SCREAMING_SNAKE_CASE								=			training_args.base_learning_rate * total_train_batch_size / 2_56

 # Initialize our trainer
 SCREAMING_SNAKE_CASE								=			Trainer(
     model=_UpperCamelCase		,							args=_UpperCamelCase		,							train_dataset=ds['train'] if training_args.do_train else None		,							eval_dataset=ds['validation'] if training_args.do_eval else None		,							tokenizer=_UpperCamelCase		,							data_collator=_UpperCamelCase		,							)

 # Training
 if training_args.do_train:
  SCREAMING_SNAKE_CASE								=			None
  if training_args.resume_from_checkpoint is not None:
   SCREAMING_SNAKE_CASE								=			training_args.resume_from_checkpoint
  elif last_checkpoint is not None:
   SCREAMING_SNAKE_CASE								=			last_checkpoint
  SCREAMING_SNAKE_CASE								=			trainer.train(resume_from_checkpoint=_UpperCamelCase	)
  trainer.save_model()
  trainer.log_metrics('train'		,							train_result.metrics	)
  trainer.save_metrics('train'		,							train_result.metrics	)
  trainer.save_state()

 # Evaluation
 if training_args.do_eval:
  SCREAMING_SNAKE_CASE								=			trainer.evaluate()
  trainer.log_metrics('eval'		,							_UpperCamelCase	)
  trainer.save_metrics('eval'		,							_UpperCamelCase	)

 # Write model card and (optionally) push to hub
 SCREAMING_SNAKE_CASE								=			{
     'tasks': 'masked-auto-encoding',
     'dataset': data_args.dataset_name,
     'tags': ['masked-auto-encoding'],
 }
 if training_args.push_to_hub:
  trainer.push_to_hub(**_UpperCamelCase	)
 else:
  trainer.create_model_card(**_UpperCamelCase	)


def        __lowerCAmelCase    (       _UpperCamelCase		:  int	)		->							Dict:
 '''simple docstring'''


 main()


if __name__ == "__main__":
      main()




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      UpperCamelCase :

    def __init__(			self :      Dict							,    snake_case__ :      str							,    snake_case__ :      str=1_3							,    snake_case__ :      Tuple=7							,    snake_case__ :      Tuple=True							,    snake_case__ :      Tuple=True							,    snake_case__ :      List[str]=False							,    snake_case__ :      Any=True							,    snake_case__ :      Union[str, Any]=9_9							,    snake_case__ :      Dict=3_2							,    snake_case__ :      Optional[Any]=5							,    snake_case__ :      Optional[Any]=4							,    snake_case__ :      Union[str, Any]=3_7							,    snake_case__ :      Tuple="gelu"							,    snake_case__ :      Dict=0.1							,    snake_case__ :      Any=0.1							,    snake_case__ :      int=5_1_2							,    snake_case__ :      Dict=1_6							,    snake_case__ :      str=2							,    snake_case__ :      Any=0.02							,    snake_case__ :      List[str]=3							,    snake_case__ :      int=4							,    snake_case__ :      List[str]=None							,    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			parent
     SCREAMING_SNAKE_CASE								=			batch_size
     SCREAMING_SNAKE_CASE								=			seq_length
     SCREAMING_SNAKE_CASE								=			is_training
     SCREAMING_SNAKE_CASE								=			use_input_mask
     SCREAMING_SNAKE_CASE								=			use_token_type_ids
     SCREAMING_SNAKE_CASE								=			use_labels
     SCREAMING_SNAKE_CASE								=			vocab_size
     SCREAMING_SNAKE_CASE								=			hidden_size
     SCREAMING_SNAKE_CASE								=			num_hidden_layers
     SCREAMING_SNAKE_CASE								=			num_attention_heads
     SCREAMING_SNAKE_CASE								=			intermediate_size
     SCREAMING_SNAKE_CASE								=			hidden_act
     SCREAMING_SNAKE_CASE								=			hidden_dropout_prob
     SCREAMING_SNAKE_CASE								=			attention_probs_dropout_prob
     SCREAMING_SNAKE_CASE								=			max_position_embeddings
     SCREAMING_SNAKE_CASE								=			type_vocab_size
     SCREAMING_SNAKE_CASE								=			type_sequence_label_size
     SCREAMING_SNAKE_CASE								=			initializer_range
     SCREAMING_SNAKE_CASE								=			num_labels
     SCREAMING_SNAKE_CASE								=			num_choices
     SCREAMING_SNAKE_CASE								=			scope

    def UpperCamelCase			(			self :      List[Any]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			ids_tensor([self.batch_size, self.seq_length]							,    self.vocab_size    )

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

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

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

     SCREAMING_SNAKE_CASE								=			self.get_config()

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

    def UpperCamelCase			(			self :      Dict    ):

     """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=snake_case__							,    initializer_range=self.initializer_range							,    use_stable_embedding=snake_case__							,    )

    def UpperCamelCase			(			self :      int							,    snake_case__ :      Tuple							,    snake_case__ :      str							,    snake_case__ :      Optional[Any]							,    snake_case__ :      Union[str, Any]							,    snake_case__ :      List[str]							,    snake_case__ :      Any							,    snake_case__ :      Optional[int]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			OpenLlamaModel(config=snake_case__    )
     model.to(snake_case__    )
     model.eval()
     SCREAMING_SNAKE_CASE								=			model(snake_case__							,    attention_mask=snake_case__    )
     SCREAMING_SNAKE_CASE								=			model(snake_case__    )
     self.parent.assertEqual(result.last_hidden_state.shape							,    (self.batch_size, self.seq_length, self.hidden_size)    )

    def UpperCamelCase			(			self :      Optional[int]							,    snake_case__ :      Dict							,    snake_case__ :      List[Any]							,    snake_case__ :      str							,    snake_case__ :      str							,    snake_case__ :      int							,    snake_case__ :      Union[str, Any]							,    snake_case__ :      Optional[int]							,    snake_case__ :      List[Any]							,    snake_case__ :      str							,    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			True
     SCREAMING_SNAKE_CASE								=			OpenLlamaModel(snake_case__    )
     model.to(snake_case__    )
     model.eval()
     SCREAMING_SNAKE_CASE								=			model(
         snake_case__							,    attention_mask=snake_case__							,    encoder_hidden_states=snake_case__							,    encoder_attention_mask=snake_case__							,    )
     SCREAMING_SNAKE_CASE								=			model(
         snake_case__							,    attention_mask=snake_case__							,    encoder_hidden_states=snake_case__							,    )
     SCREAMING_SNAKE_CASE								=			model(snake_case__							,    attention_mask=snake_case__    )
     self.parent.assertEqual(result.last_hidden_state.shape							,    (self.batch_size, self.seq_length, self.hidden_size)    )

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

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			OpenLlamaForCausalLM(config=snake_case__    )
     model.to(snake_case__    )
     model.eval()
     SCREAMING_SNAKE_CASE								=			model(snake_case__							,    attention_mask=snake_case__							,    labels=snake_case__    )
     self.parent.assertEqual(result.logits.shape							,    (self.batch_size, self.seq_length, self.vocab_size)    )

    def UpperCamelCase			(			self :      Optional[int]							,    snake_case__ :      Union[str, Any]							,    snake_case__ :      List[Any]							,    snake_case__ :      Any							,    snake_case__ :      int							,    snake_case__ :      Any							,    snake_case__ :      Union[str, Any]							,    snake_case__ :      int							,    snake_case__ :      str							,    snake_case__ :      Union[str, Any]							,    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			True
     SCREAMING_SNAKE_CASE								=			True
     SCREAMING_SNAKE_CASE								=			OpenLlamaForCausalLM(config=snake_case__    )
     model.to(snake_case__    )
     model.eval()

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

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

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

     SCREAMING_SNAKE_CASE								=			model(
         snake_case__							,    attention_mask=snake_case__							,    encoder_hidden_states=snake_case__							,    encoder_attention_mask=snake_case__							,    output_hidden_states=snake_case__							,    )['hidden_states'][0]
     SCREAMING_SNAKE_CASE								=			model(
         snake_case__							,    attention_mask=snake_case__							,    encoder_hidden_states=snake_case__							,    encoder_attention_mask=snake_case__							,    past_key_values=snake_case__							,    output_hidden_states=snake_case__							,    )['hidden_states'][0]

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




    def UpperCamelCase			(			self :      Any    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			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
     )						,	
     )								=			config_and_inputs
     SCREAMING_SNAKE_CASE								=			{'input_ids': input_ids, 'attention_mask': input_mask}
     return config, inputs_dict




@require_torch
class      UpperCamelCase (     SCREAMING_SNAKE_CASE	,			SCREAMING_SNAKE_CASE	,			SCREAMING_SNAKE_CASE	,			unittest.TestCase   ):
    __UpperCamelCase         =(
        (OpenLlamaModel, OpenLlamaForCausalLM, OpenLlamaForSequenceClassification) if is_torch_available() else ()
    )
    __UpperCamelCase         =(OpenLlamaForCausalLM,) if is_torch_available() else ()
    __UpperCamelCase         =(
        {
            "feature-extraction": OpenLlamaModel,
            "text-classification": OpenLlamaForSequenceClassification,
            "text-generation": OpenLlamaForCausalLM,
            "zero-shot": OpenLlamaForSequenceClassification,
        }
        if is_torch_available()
        else {}
    )
    __UpperCamelCase         =False
    __UpperCamelCase         =False

    def UpperCamelCase			(			self :      int    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			OpenLlamaModelTester(self    )
     SCREAMING_SNAKE_CASE								=			ConfigTester(self							,    config_class=snake_case__							,    hidden_size=3_7    )

    def UpperCamelCase			(			self :      Optional[int]    ):

     """simple docstring"""






     self.config_tester.run_common_tests()

    def UpperCamelCase			(			self :      Any    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*snake_case__    )

    def UpperCamelCase			(			self :      List[str]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			self.model_tester.prepare_config_and_inputs()
     for type in ["absolute", "relative_key", "relative_key_query"]:
      SCREAMING_SNAKE_CASE								=			type
      self.model_tester.create_and_check_model(*snake_case__    )

    def UpperCamelCase			(			self :      Any    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE						,	SCREAMING_SNAKE_CASE								=			self.model_tester.prepare_config_and_inputs_for_common()
     SCREAMING_SNAKE_CASE								=			3
     SCREAMING_SNAKE_CASE								=			input_dict['input_ids']
     SCREAMING_SNAKE_CASE								=			input_ids.ne(1    ).to(snake_case__    )
     SCREAMING_SNAKE_CASE								=			ids_tensor([self.model_tester.batch_size]							,    self.model_tester.type_sequence_label_size    )
     SCREAMING_SNAKE_CASE								=			OpenLlamaForSequenceClassification(snake_case__    )
     model.to(snake_case__    )
     model.eval()
     SCREAMING_SNAKE_CASE								=			model(snake_case__							,    attention_mask=snake_case__							,    labels=snake_case__    )
     self.assertEqual(result.logits.shape							,    (self.model_tester.batch_size, self.model_tester.num_labels)    )

    def UpperCamelCase			(			self :      int    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE						,	SCREAMING_SNAKE_CASE								=			self.model_tester.prepare_config_and_inputs_for_common()
     SCREAMING_SNAKE_CASE								=			3
     SCREAMING_SNAKE_CASE								=			'single_label_classification'
     SCREAMING_SNAKE_CASE								=			input_dict['input_ids']
     SCREAMING_SNAKE_CASE								=			input_ids.ne(1    ).to(snake_case__    )
     SCREAMING_SNAKE_CASE								=			ids_tensor([self.model_tester.batch_size]							,    self.model_tester.type_sequence_label_size    )
     SCREAMING_SNAKE_CASE								=			OpenLlamaForSequenceClassification(snake_case__    )
     model.to(snake_case__    )
     model.eval()
     SCREAMING_SNAKE_CASE								=			model(snake_case__							,    attention_mask=snake_case__							,    labels=snake_case__    )
     self.assertEqual(result.logits.shape							,    (self.model_tester.batch_size, self.model_tester.num_labels)    )

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

     """simple docstring"""






     SCREAMING_SNAKE_CASE						,	SCREAMING_SNAKE_CASE								=			self.model_tester.prepare_config_and_inputs_for_common()
     SCREAMING_SNAKE_CASE								=			3
     SCREAMING_SNAKE_CASE								=			'multi_label_classification'
     SCREAMING_SNAKE_CASE								=			input_dict['input_ids']
     SCREAMING_SNAKE_CASE								=			input_ids.ne(1    ).to(snake_case__    )
     SCREAMING_SNAKE_CASE								=			ids_tensor(
         [self.model_tester.batch_size, config.num_labels]							,    self.model_tester.type_sequence_label_size    ).to(torch.float    )
     SCREAMING_SNAKE_CASE								=			OpenLlamaForSequenceClassification(snake_case__    )
     model.to(snake_case__    )
     model.eval()
     SCREAMING_SNAKE_CASE								=			model(snake_case__							,    attention_mask=snake_case__							,    labels=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 UpperCamelCase			(			self :      Optional[int]    ):

     """simple docstring"""






     pass




    @parameterized.expand([('linear',), ('dynamic',)]    )
    def UpperCamelCase			(			self :      str							,    snake_case__ :      Optional[Any]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE						,	SCREAMING_SNAKE_CASE								=			self.model_tester.prepare_config_and_inputs_for_common()
     SCREAMING_SNAKE_CASE								=			ids_tensor([1, 1_0]							,    config.vocab_size    )
     SCREAMING_SNAKE_CASE								=			ids_tensor([1, int(config.max_position_embeddings * 1.5    )]							,    config.vocab_size    )

     set_seed(4_2    )  # Fixed seed at init time so the two models get the same random weights
     SCREAMING_SNAKE_CASE								=			OpenLlamaModel(snake_case__    )
     original_model.to(snake_case__    )
     original_model.eval()
     SCREAMING_SNAKE_CASE								=			original_model(snake_case__    ).last_hidden_state
     SCREAMING_SNAKE_CASE								=			original_model(snake_case__    ).last_hidden_state

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

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

