


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

from datasets import Dataset, load_dataset

import transformers
from transformers import (
    CONFIG_MAPPING,
    MODEL_FOR_MASKED_LM_MAPPING,
    AutoConfig,
    AutoModelForMaskedLM,
    AutoTokenizer,
    DataCollatorForWholeWordMask,
    HfArgumentParser,
    Trainer,
    TrainingArguments,
    set_seed,
)
from transformers.trainer_utils import get_last_checkpoint, is_main_process


a					  =     logging.getLogger(__name__)
a					  =     list(MODEL_FOR_MASKED_LM_MAPPING.keys())
a					  =     tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)







@dataclass
class        UpperCamelCase__   :
     __SCREAMING_SNAKE_CASE       :  Optional[str]						=   field(
         default=__lowerCamelCase							,				metadata={
             'help': (
                 'The model checkpoint for weights initialization.Don\'t set if you want to train a model from scratch.'
             )
         }							,				)
     __SCREAMING_SNAKE_CASE       :  Optional[str]						=   field(
         default=__lowerCamelCase							,				metadata={'help': 'If training from scratch, pass a model type from the list: ' + ', '.join(__lowerCamelCase   )}							,				)
     __SCREAMING_SNAKE_CASE       :  Optional[str]						=   field(
         default=__lowerCamelCase							,				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'
             )
         }							,				)
     __SCREAMING_SNAKE_CASE       :  Optional[str]						=   field(
         default=__lowerCamelCase							,				metadata={'help': 'Pretrained config name or path if not the same as model_name'}   )
     __SCREAMING_SNAKE_CASE       :  Optional[str]						=   field(
         default=__lowerCamelCase							,				metadata={'help': 'Pretrained tokenizer name or path if not the same as model_name'}   )
     __SCREAMING_SNAKE_CASE       :  Optional[str]						=   field(
         default=__lowerCamelCase							,				metadata={'help': 'Where do you want to store the pretrained models downloaded from huggingface.co'}							,				)
     __SCREAMING_SNAKE_CASE       :  bool						=   field(
         default=__lowerCamelCase							,				metadata={'help': 'Whether to use one of the fast tokenizer (backed by the tokenizers library) or not.'}							,				)
     __SCREAMING_SNAKE_CASE       :  str						=   field(
         default='main'							,				metadata={'help': 'The specific model version to use (can be a branch name, tag name or commit id).'}							,				)
     __SCREAMING_SNAKE_CASE       :  bool						=   field(
         default=__lowerCamelCase							,				metadata={
             'help': (
                 'Will use the token generated when running `huggingface-cli login` (necessary to use this script '
                 'with private models).'
             )
         }							,				)

     def   UpperCAmelCase__       (     self       :   Dict       ):






          '''simple docstring'''






          if self.config_overrides is not None and (self.config_name is not None or self.model_name_or_path is not None):
               raise ValueError(
                   """--config_overrides can't be used in combination with --config_name or --model_name_or_path"""       )







@dataclass
class        UpperCamelCase__   :
     __SCREAMING_SNAKE_CASE       :  Optional[str]						=   field(
         default=__lowerCamelCase							,				metadata={'help': 'The name of the dataset to use (via the datasets library).'}   )
     __SCREAMING_SNAKE_CASE       :  Optional[str]						=   field(
         default=__lowerCamelCase							,				metadata={'help': 'The configuration name of the dataset to use (via the datasets library).'}   )
     __SCREAMING_SNAKE_CASE       :  Optional[str]						=   field(default=__lowerCamelCase							,				metadata={'help': 'The input training data file (a text file).'}   )
     __SCREAMING_SNAKE_CASE       :  Optional[str]						=   field(
         default=__lowerCamelCase							,				metadata={'help': 'An optional input evaluation data file to evaluate the perplexity on (a text file).'}							,				)
     __SCREAMING_SNAKE_CASE       :  Optional[str]						=   field(
         default=__lowerCamelCase							,				metadata={'help': 'An optional input train ref data file for whole word masking in Chinese.'}							,				)
     __SCREAMING_SNAKE_CASE       :  Optional[str]						=   field(
         default=__lowerCamelCase							,				metadata={'help': 'An optional input validation ref data file for whole word masking in Chinese.'}							,				)
     __SCREAMING_SNAKE_CASE       :  bool						=   field(
         default=__lowerCamelCase							,				metadata={'help': 'Overwrite the cached training and evaluation sets'}   )
     __SCREAMING_SNAKE_CASE       :  Optional[int]						=   field(
         default=5							,				metadata={
             'help': 'The percentage of the train set used as validation set in case there\'s no validation split'
         }							,				)
     __SCREAMING_SNAKE_CASE       :  Optional[int]						=   field(
         default=__lowerCamelCase							,				metadata={
             'help': (
                 'The maximum total input sequence length after tokenization. Sequences longer '
                 'than this will be truncated. Default to the max input length of the model.'
             )
         }							,				)
     __SCREAMING_SNAKE_CASE       :  Optional[int]						=   field(
         default=__lowerCamelCase							,				metadata={'help': 'The number of processes to use for the preprocessing.'}							,				)
     __SCREAMING_SNAKE_CASE       :  float						=   field(
         default=0.1_5							,				metadata={'help': 'Ratio of tokens to mask for masked language modeling loss'}   )
     __SCREAMING_SNAKE_CASE       :  bool						=   field(
         default=__lowerCamelCase							,				metadata={
             'help': (
                 'Whether to pad all samples to `max_seq_length`. '
                 'If False, will pad the samples dynamically when batching to the maximum length in the batch.'
             )
         }							,				)

     def   UpperCAmelCase__       (     self       :   Tuple       ):






          '''simple docstring'''






          if self.train_file is not None:
               lowercase_  							=			self.train_file.split("""."""       )[-1]
               assert extension in ["csv", "json", "txt"], "`train_file` should be a csv, a json or a txt file."
          if self.validation_file is not None:
               lowercase_  							=			self.validation_file.split("""."""       )[-1]
               assert extension in ["csv", "json", "txt"], "`validation_file` should be a csv, a json or a txt file."



def 			UpperCAmelCase_				(      UpperCAmelCase__						,							UpperCAmelCase__  ):


     with open(lowerCAmelCase_						,							"""r"""						,							encoding="""utf-8"""  ) as f:
          lowercase_  							=			[json.loads(lowerCAmelCase_  ) for line in f.read().splitlines() if (len(lowerCAmelCase_  ) > 0 and not line.isspace())]
     assert len(lowerCAmelCase_  ) == len(lowerCAmelCase_  )

     lowercase_  							=			{c: dataset[c] for c in dataset.column_names}
     lowercase_  							=			refs
     return Dataset.from_dict(lowerCAmelCase_  )


def 			UpperCAmelCase_				(      ):
     lowercase_  							=			HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)  )
     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.
          lowercase_  							=			parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]  )  )
     else:
          lowercase_  							=			parser.parse_args_into_dataclasses()

     # Detecting last checkpoint.
     lowercase_  							=			None
     if os.path.isdir(training_args.output_dir  ) and training_args.do_train and not training_args.overwrite_output_dir:
          lowercase_  							=			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:
               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."""  )

    # 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  )]						,							)
     logger.setLevel(logging.INFO if is_main_process(training_args.local_rank  ) else logging.WARN  )

     # 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}'''  )
     # Set the verbosity to info of the Transformers logger (on main process only):
     if is_main_process(training_args.local_rank  ):
          transformers.utils.logging.set_verbosity_info()
          transformers.utils.logging.enable_default_handler()
          transformers.utils.logging.enable_explicit_format()
     logger.info("""Training/evaluation parameters %s"""						,							lowerCAmelCase_  )

     # Set seed before initializing model.
     set_seed(training_args.seed  )

     # Get the datasets: you can either provide your own CSV/JSON/TXT training and evaluation files (see below)
     # or just provide the name of one of the public datasets available on the hub at https://huggingface.co/datasets/
     # (the dataset will be downloaded automatically from the datasets Hub).
     #
     # For CSV/JSON files, this script will use the column called 'text' or the first column if no column called
     # 'text' is found. You can easily tweak this behavior (see below).
     #
     # In distributed training, the load_dataset function guarantee that only one local process can concurrently
     # download the dataset.
     if data_args.dataset_name is not None:
          # Downloading and loading a dataset from the hub.
          lowercase_  							=			load_dataset(data_args.dataset_name						,							data_args.dataset_config_name  )
          if "validation" not in datasets.keys():
               lowercase_  							=			load_dataset(
                   data_args.dataset_name						,							data_args.dataset_config_name						,							split=F'''train[:{data_args.validation_split_percentage}%]'''						,							)
               lowercase_  							=			load_dataset(
                   data_args.dataset_name						,							data_args.dataset_config_name						,							split=F'''train[{data_args.validation_split_percentage}%:]'''						,							)
     else:
          lowercase_  							=			{}
          if data_args.train_file is not None:
               lowercase_  							=			data_args.train_file
          if data_args.validation_file is not None:
               lowercase_  							=			data_args.validation_file
          lowercase_  							=			data_args.train_file.split("""."""  )[-1]
          if extension == "txt":
               lowercase_  							=			"text"
          lowercase_  							=			load_dataset(lowerCAmelCase_						,							data_files=lowerCAmelCase_  )
     # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at
     # https://huggingface.co/docs/datasets/loading_datasets.html.

     # Load pretrained model and tokenizer
     #
     # Distributed training:
     # The .from_pretrained methods guarantee that only one local process can concurrently
     # download model & vocab.
     lowercase_  							=			{
         "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:
          lowercase_  							=			AutoConfig.from_pretrained(model_args.config_name						,							**lowerCAmelCase_  )
     elif model_args.model_name_or_path:
          lowercase_  							=			AutoConfig.from_pretrained(model_args.model_name_or_path						,							**lowerCAmelCase_  )
     else:
          lowercase_  							=			CONFIG_MAPPING[model_args.model_type]()
          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}'''  )

     lowercase_  							=			{
         "cache_dir": model_args.cache_dir,
         "use_fast": model_args.use_fast_tokenizer,
         "revision": model_args.model_revision,
         "use_auth_token": True if model_args.use_auth_token else None,
     }
     if model_args.tokenizer_name:
          lowercase_  							=			AutoTokenizer.from_pretrained(model_args.tokenizer_name						,							**lowerCAmelCase_  )
     elif model_args.model_name_or_path:
          lowercase_  							=			AutoTokenizer.from_pretrained(model_args.model_name_or_path						,							**lowerCAmelCase_  )
     else:
          raise ValueError(
              """You are instantiating a new tokenizer from scratch. This is not supported by this script."""
              """You can do it from another script, save it, and load it from here, using --tokenizer_name."""  )

     if model_args.model_name_or_path:
          lowercase_  							=			AutoModelForMaskedLM.from_pretrained(
              model_args.model_name_or_path						,							from_tf=bool(""".ckpt""" in model_args.model_name_or_path  )						,							config=lowerCAmelCase_						,							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"""  )
          lowercase_  							=			AutoModelForMaskedLM.from_config(lowerCAmelCase_  )

     model.resize_token_embeddings(len(lowerCAmelCase_  )  )

     # Preprocessing the datasets.
     # First we tokenize all the texts.
     if training_args.do_train:
          lowercase_  							=			datasets["train"].column_names
     else:
          lowercase_  							=			datasets["validation"].column_names
     lowercase_  							=			"text" if "text" in column_names else column_names[0]

     lowercase_  							=			"max_length" if data_args.pad_to_max_length else False

     def tokenize_function(UpperCAmelCase__  ):
          # Remove empty lines
          lowercase_  							=			[line for line in examples["text"] if len(lowerCAmelCase_  ) > 0 and not line.isspace()]
          return tokenizer(examples["""text"""]						,							padding=lowerCAmelCase_						,							truncation=lowerCAmelCase_						,							max_length=data_args.max_seq_length  )

     lowercase_  							=			datasets.map(
         lowerCAmelCase_						,							batched=lowerCAmelCase_						,							num_proc=data_args.preprocessing_num_workers						,							remove_columns=[text_column_name]						,							load_from_cache_file=not data_args.overwrite_cache						,							)

     # Add the chinese references if provided
     if data_args.train_ref_file is not None:
          lowercase_  							=			add_chinese_references(tokenized_datasets["""train"""]						,							data_args.train_ref_file  )
     if data_args.validation_ref_file is not None:
          lowercase_  							=			add_chinese_references(
              tokenized_datasets["""validation"""]						,							data_args.validation_ref_file  )
     # If we have ref files, need to avoid it removed by trainer
     lowercase_  							=			data_args.train_ref_file or data_args.validation_ref_file
     if has_ref:
          lowercase_  							=			False

     # Data collator
     # This one will take care of randomly masking the tokens.
     lowercase_  							=			DataCollatorForWholeWordMask(tokenizer=lowerCAmelCase_						,							mlm_probability=data_args.mlm_probability  )

     # Initialize our Trainer
     lowercase_  							=			Trainer(
         model=lowerCAmelCase_						,							args=lowerCAmelCase_						,							train_dataset=tokenized_datasets["""train"""] if training_args.do_train else None						,							eval_dataset=tokenized_datasets["""validation"""] if training_args.do_eval else None						,							tokenizer=lowerCAmelCase_						,							data_collator=lowerCAmelCase_						,							)

     # Training
     if training_args.do_train:
          if last_checkpoint is not None:
               lowercase_  							=			last_checkpoint
          elif model_args.model_name_or_path is not None and os.path.isdir(model_args.model_name_or_path  ):
               lowercase_  							=			model_args.model_name_or_path
          else:
               lowercase_  							=			None
          lowercase_  							=			trainer.train(resume_from_checkpoint=lowerCAmelCase_  )
          trainer.save_model()  # Saves the tokenizer too for easy upload

          lowercase_  							=			os.path.join(training_args.output_dir						,							"""train_results.txt"""  )
          if trainer.is_world_process_zero():
               with open(lowerCAmelCase_						,							"""w"""  ) as writer:
                    logger.info("""***** Train results *****"""  )
                    for key, value in sorted(train_result.metrics.items()  ):
                         logger.info(F'''  {key} = {value}'''  )
                         writer.write(F'''{key} = {value}\n'''  )

            # Need to save the state, since Trainer.save_model saves only the tokenizer with the model
               trainer.state.save_to_json(os.path.join(training_args.output_dir						,							"""trainer_state.json"""  )  )

    # Evaluation
     lowercase_  							=			{}
     if training_args.do_eval:
          logger.info("""*** Evaluate ***"""  )

          lowercase_  							=			trainer.evaluate()

          lowercase_  							=			math.exp(eval_output["""eval_loss"""]  )
          lowercase_  							=			perplexity

          lowercase_  							=			os.path.join(training_args.output_dir						,							"""eval_results_mlm_wwm.txt"""  )
          if trainer.is_world_process_zero():
               with open(lowerCAmelCase_						,							"""w"""  ) as writer:
                    logger.info("""***** Eval results *****"""  )
                    for key, value in sorted(results.items()  ):
                         logger.info(F'''  {key} = {value}'''  )
                         writer.write(F'''{key} = {value}\n'''  )

     return results


def 			UpperCAmelCase_				(      UpperCAmelCase__  ):
     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       :   List[str]							,					UpperCamelCase__       :   Tuple							,					UpperCamelCase__       :   List[str]=13							,					UpperCamelCase__       :   str=7							,					UpperCamelCase__       :   List[str]=True							,					UpperCamelCase__       :   List[str]=True							,					UpperCamelCase__       :   str=False							,					UpperCamelCase__       :   str=True							,					UpperCamelCase__       :   Union[str, Any]=99							,					UpperCamelCase__       :   Dict=32							,					UpperCamelCase__       :   Union[str, Any]=5							,					UpperCamelCase__       :   Tuple=4							,					UpperCamelCase__       :   Dict=37							,					UpperCamelCase__       :   Any="gelu"							,					UpperCamelCase__       :   Dict=0.1							,					UpperCamelCase__       :   Tuple=0.1							,					UpperCamelCase__       :   Tuple=512							,					UpperCamelCase__       :   Union[str, Any]=16							,					UpperCamelCase__       :   List[str]=2							,					UpperCamelCase__       :   int=0.02							,					UpperCamelCase__       :   int=3							,					UpperCamelCase__       :   Optional[Any]=4							,					UpperCamelCase__       :   int=None							,					):






										'''simple docstring'''






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

					def   UpperCAmelCase__       (     self       :   Optional[Any]       ):






										'''simple docstring'''






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

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

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

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

										lowercase_  							=			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=UpperCamelCase__							,					initializer_range=self.initializer_range							,					use_stable_embedding=UpperCamelCase__							,					)

					def   UpperCAmelCase__       (     self       :   str							,					UpperCamelCase__       :   Optional[Any]							,					UpperCamelCase__       :   Tuple							,					UpperCamelCase__       :   Dict							,					UpperCamelCase__       :   List[str]							,					UpperCamelCase__       :   Dict							,					UpperCamelCase__       :   str							,					UpperCamelCase__       :   Optional[Any]       ):






										'''simple docstring'''






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

					def   UpperCAmelCase__       (     self       :   Union[str, Any]							,					UpperCamelCase__       :   Tuple							,					UpperCamelCase__       :   str							,					UpperCamelCase__       :   Tuple							,					UpperCamelCase__       :   str							,					UpperCamelCase__       :   int							,					UpperCamelCase__       :   str							,					UpperCamelCase__       :   Optional[int]							,					UpperCamelCase__       :   Optional[Any]							,					UpperCamelCase__       :   Optional[Any]							,					):






										'''simple docstring'''






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

					def   UpperCAmelCase__       (     self       :   Any							,					UpperCamelCase__       :   List[str]							,					UpperCamelCase__       :   List[Any]							,					UpperCamelCase__       :   Optional[int]							,					UpperCamelCase__       :   str							,					UpperCamelCase__       :   Any							,					UpperCamelCase__       :   Union[str, Any]							,					UpperCamelCase__       :   Tuple							,					UpperCamelCase__       :   Any							,					UpperCamelCase__       :   List[Any]							,					):






										'''simple docstring'''






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

					def   UpperCAmelCase__       (     self       :   int							,					UpperCamelCase__       :   Optional[int]							,					UpperCamelCase__       :   int							,					UpperCamelCase__       :   Tuple							,					UpperCamelCase__       :   Optional[Any]							,					UpperCamelCase__       :   int							,					UpperCamelCase__       :   List[str]							,					UpperCamelCase__       :   int							,					UpperCamelCase__       :   str							,					UpperCamelCase__       :   Tuple							,					):






										'''simple docstring'''






										lowercase_  							=			True
										lowercase_  							=			True
										lowercase_  							=			OpenLlamaForCausalLM(config=UpperCamelCase__       )
										model.to(UpperCamelCase__       )
										model.eval()

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

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

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

										lowercase_  							=			model(
										    UpperCamelCase__							,					attention_mask=UpperCamelCase__							,					encoder_hidden_states=UpperCamelCase__							,					encoder_attention_mask=UpperCamelCase__							,					output_hidden_states=UpperCamelCase__							,					)["""hidden_states"""][0]
										lowercase_  							=			model(
										    UpperCamelCase__							,					attention_mask=UpperCamelCase__							,					encoder_hidden_states=UpperCamelCase__							,					encoder_attention_mask=UpperCamelCase__							,					past_key_values=UpperCamelCase__							,					output_hidden_states=UpperCamelCase__							,					)["""hidden_states"""][0]

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







					def   UpperCAmelCase__       (     self       :   Optional[Any]       ):






										'''simple docstring'''






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







@require_torch
class        UpperCamelCase__   (	__magic_name__							,				__magic_name__							,				__magic_name__							,				unittest.TestCase   ):
					__SCREAMING_SNAKE_CASE       :  Union[str, Any]						=   (
					    (OpenLlamaModel, OpenLlamaForCausalLM, OpenLlamaForSequenceClassification) if is_torch_available() else ()
					)
					__SCREAMING_SNAKE_CASE       :  List[str]						=   (OpenLlamaForCausalLM,) if is_torch_available() else ()
					__SCREAMING_SNAKE_CASE       :  List[Any]						=   (
					    {
					        'feature-extraction': OpenLlamaModel,
					        'text-classification': OpenLlamaForSequenceClassification,
					        'text-generation': OpenLlamaForCausalLM,
					        'zero-shot': OpenLlamaForSequenceClassification,
					    }
					    if is_torch_available()
					    else {}
					)
					__SCREAMING_SNAKE_CASE       :  Tuple						=   False
					__SCREAMING_SNAKE_CASE       :  Optional[int]						=   False

					def   UpperCAmelCase__       (     self       :   List[Any]       ):






										'''simple docstring'''






										lowercase_  							=			OpenLlamaModelTester(self       )
										lowercase_  							=			ConfigTester(self							,					config_class=UpperCamelCase__							,					hidden_size=37       )

					def   UpperCAmelCase__       (     self       :   Optional[Any]       ):






										'''simple docstring'''






										self.config_tester.run_common_tests()

					def   UpperCAmelCase__       (     self       :   str       ):






										'''simple docstring'''






										lowercase_  							=			self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*UpperCamelCase__       )

					def   UpperCAmelCase__       (     self       :   Dict       ):






										'''simple docstring'''






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

					def   UpperCAmelCase__       (     self       :   List[Any]       ):






										'''simple docstring'''






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

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






										'''simple docstring'''






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

					def   UpperCAmelCase__       (     self       :   Dict       ):






										'''simple docstring'''






										lowercase_							,					lowercase_  							=			self.model_tester.prepare_config_and_inputs_for_common()
										lowercase_  							=			3
										lowercase_  							=			"""multi_label_classification"""
										lowercase_  							=			input_dict["""input_ids"""]
										lowercase_  							=			input_ids.ne(1       ).to(UpperCamelCase__       )
										lowercase_  							=			ids_tensor(
										    [self.model_tester.batch_size, config.num_labels]							,					self.model_tester.type_sequence_label_size       ).to(torch.float       )
										lowercase_  							=			OpenLlamaForSequenceClassification(UpperCamelCase__       )
										model.to(UpperCamelCase__       )
										model.eval()
										lowercase_  							=			model(UpperCamelCase__							,					attention_mask=UpperCamelCase__							,					labels=UpperCamelCase__       )
										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       :   Dict       ):






										'''simple docstring'''






										pass







					@parameterized.expand([("""linear""",), ("""dynamic""",)]       )
					def   UpperCAmelCase__       (     self       :   Tuple							,					UpperCamelCase__       :   Union[str, Any]       ):






										'''simple docstring'''






										lowercase_							,					lowercase_  							=			self.model_tester.prepare_config_and_inputs_for_common()
										lowercase_  							=			ids_tensor([1, 10]							,					config.vocab_size       )
										lowercase_  							=			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_  							=			OpenLlamaModel(UpperCamelCase__       )
										original_model.to(UpperCamelCase__       )
										original_model.eval()
										lowercase_  							=			original_model(UpperCamelCase__       ).last_hidden_state
										lowercase_  							=			original_model(UpperCamelCase__       ).last_hidden_state

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

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


