
"""simple docstring"""


import os
from pathlib import Path

import numpy as np
import pytest
from pack_dataset import pack_data_dir
from parameterized import parameterized
from save_len_file import save_len_file
from torch.utils.data import DataLoader

from transformers import AutoTokenizer
from transformers.models.mbart.modeling_mbart import shift_tokens_right
from transformers.testing_utils import TestCasePlus, slow
from utils import FAIRSEQ_AVAILABLE, DistributedSortishSampler, LegacySeqaSeqDataset, SeqaSeqDataset


_lowercase				 =					'''bert-base-cased'''
_lowercase				 =					'''google/pegasus-xsum'''
_lowercase				 =					[''' Sam ate lunch today.''', '''Sams lunch ingredients.''']
_lowercase				 =					['''A very interesting story about what I ate for lunch.''', '''Avocado, celery, turkey, coffee''']
_lowercase				 =					'''patrickvonplaten/t5-tiny-random'''
_lowercase				 =					'''sshleifer/bart-tiny-random'''
_lowercase				 =					'''sshleifer/tiny-mbart'''
_lowercase				 =					'''sshleifer/tiny-marian-en-de'''






def    _snake_case					(		snake_case__ :  Path						,   snake_case__ :  list  ):
	A     = '\n'.join(snake_case__  )
	Path(snake_case__  ).open('w'  ).writelines(snake_case__  )






def    _snake_case					(		snake_case__ :  str  ):
	for split in ["train", "val", "test"]:
		_dump_articles(os.path.join(snake_case__						,   F'{split}.source'  )						,   snake_case__  )
		_dump_articles(os.path.join(snake_case__						,   F'{split}.target'  )						,   snake_case__  )
	return tmp_dir







class 					lowerCAmelCase_					(   _lowercase     ):



						'''simple docstring'''


						@parameterized.expand(
						    [
						        MBART_TINY,
						        MARIAN_TINY,
						        T5_TINY,
						        BART_TINY,
						        PEGASUS_XSUM,
						    ]	,)
						@slow
						def     _SCREAMING_SNAKE_CASE   (			self  :   int	,A_  :   Any   )  -> Union[str, Any]:
							A     = AutoTokenizer.from_pretrained(A_   )
							A     = make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()   )
							A     = max(len(tokenizer.encode(A_   )   ) for a in ARTICLES   )
							A     = max(len(tokenizer.encode(A_   )   ) for a in SUMMARIES   )
							A     = 4
							A     = 8
							assert max_len_target > max_src_len  # Will be truncated
							assert max_len_source > max_src_len  # Will be truncated
							A ,   A     = 'ro_RO', 'de_DE'  # ignored for all but mbart, but never causes error.
							A     = SeqaSeqDataset(
							    A_	,data_dir=A_	,type_path='train'	,max_source_length=A_	,max_target_length=A_	,src_lang=A_	,tgt_lang=A_	,)
							A     = DataLoader(A_	,batch_size=2	,collate_fn=train_dataset.collate_fn   )
							for batch in dataloader:
								assert isinstance(A_	,A_   )
								assert batch["attention_mask"].shape == batch["input_ids"].shape
								# show that articles were trimmed.
								assert batch["input_ids"].shape[1] == max_src_len
								# show that targets are the same len
								assert batch["labels"].shape[1] == max_tgt_len
								if tok_name != MBART_TINY:
									continue
								# check language codes in correct place
								A     = shift_tokens_right(batch['labels']	,tokenizer.pad_token_id   )
								assert batch["decoder_input_ids"][0, 0].item() == tokenizer.lang_code_to_id[tgt_lang]
								assert batch["decoder_input_ids"][0, -1].item() == tokenizer.eos_token_id
								assert batch["input_ids"][0, -2].item() == tokenizer.eos_token_id
								assert batch["input_ids"][0, -1].item() == tokenizer.lang_code_to_id[src_lang]

								break  # No need to test every batch


						@parameterized.expand([BART_TINY, BERT_BASE_CASED]   )
						def     _SCREAMING_SNAKE_CASE   (			self  :   Union[str, Any]	,A_  :   Any   )  -> Dict:
							A     = AutoTokenizer.from_pretrained(A_   )
							A     = make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()   )
							A     = max(len(tokenizer.encode(A_   )   ) for a in ARTICLES   )
							A     = max(len(tokenizer.encode(A_   )   ) for a in SUMMARIES   )
							A     = 4
							A     = LegacySeqaSeqDataset(
							    A_	,data_dir=A_	,type_path='train'	,max_source_length=20	,max_target_length=A_	,)
							A     = DataLoader(A_	,batch_size=2	,collate_fn=train_dataset.collate_fn   )
							for batch in dataloader:
								assert batch["attention_mask"].shape == batch["input_ids"].shape
								# show that articles were trimmed.
								assert batch["input_ids"].shape[1] == max_len_source
								assert 20 >= batch["input_ids"].shape[1]  # trimmed significantly
								# show that targets were truncated
								assert batch["labels"].shape[1] == trunc_target  # Truncated
								assert max_len_target > trunc_target  # Truncated
								break  # No need to test every batch


						def     _SCREAMING_SNAKE_CASE   (			self  :   List[str]   )  -> Optional[int]:
							A     = AutoTokenizer.from_pretrained('facebook/mbart-large-cc25'   )

							A     = Path(make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()   )   )
							A     = tmp_dir.joinpath('train.source'   ).open().readlines()
							A     = Path(make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()   )   )
							pack_data_dir(A_	,A_	,128	,A_   )
							A     = {x.name for x in tmp_dir.iterdir()}
							A     = {x.name for x in save_dir.iterdir()}
							A     = save_dir.joinpath('train.source'   ).open().readlines()
							# orig: [' Sam ate lunch today.\n', 'Sams lunch ingredients.']
							# desired_packed: [' Sam ate lunch today.\n Sams lunch ingredients.']
							assert len(A_   ) < len(A_   )
							assert len(A_   ) == 1
							assert len(packed_examples[0]   ) == sum(len(A_   ) for x in orig_examples   )
							assert orig_paths == new_paths


						@pytest.mark.skipif(not FAIRSEQ_AVAILABLE	,reason='This test requires fairseq'   )
						def     _SCREAMING_SNAKE_CASE   (			self  :   Tuple   )  -> List[Any]:
							if not FAIRSEQ_AVAILABLE:
								return
							A ,   A ,   A     = self._get_dataset(max_len=64   )
							A     = 64
							A     = ds.make_dynamic_sampler(A_	,required_batch_size_multiple=A_   )
							A     = [len(A_   ) for x in batch_sampler]
							assert len(set(A_   )   ) > 1  # it's not dynamic batch size if every batch is the same length
							assert sum(A_   ) == len(A_   )  # no dropped or added examples
							A     = DataLoader(A_	,batch_sampler=A_	,collate_fn=ds.collate_fn	,num_workers=2   )
							A     = []
							A     = []
							for batch in data_loader:
								A     = batch['input_ids'].shape
								A     = src_shape[0]
								assert bs % required_batch_size_multiple == 0 or bs < required_batch_size_multiple
								A     = np.product(batch['input_ids'].shape   )
								num_src_per_batch.append(A_   )
								if num_src_tokens > (max_tokens * 1.1):
									failures.append(A_   )
							assert num_src_per_batch[0] == max(A_   )
							if failures:
								raise AssertionError(F'too many tokens in {len(A_   )} batches'   )


						def     _SCREAMING_SNAKE_CASE   (			self  :   List[str]   )  -> Tuple:
							A ,   A ,   A     = self._get_dataset(max_len=512   )
							A     = 2
							A     = ds.make_sortish_sampler(A_	,shuffle=A_   )

							A     = DataLoader(A_	,batch_size=A_	,collate_fn=ds.collate_fn	,num_workers=2   )
							A     = DataLoader(A_	,batch_size=A_	,collate_fn=ds.collate_fn	,num_workers=2	,sampler=A_   )

							A     = tokenizer.pad_token_id

							def count_pad_tokens(A_  :   Tuple	,A_  :   Union[str, Any]="input_ids"   ):
								return [batch[k].eq(A_   ).sum().item() for batch in data_loader]

							assert sum(count_pad_tokens(A_	,k='labels'   )   ) < sum(count_pad_tokens(A_	,k='labels'   )   )
							assert sum(count_pad_tokens(A_   )   ) < sum(count_pad_tokens(A_   )   )
							assert len(A_   ) == len(A_   )


						def     _SCREAMING_SNAKE_CASE   (			self  :   List[Any]	,A_  :   Optional[Any]=1000	,A_  :   Union[str, Any]=128   )  -> Union[str, Any]:
							if os.getenv('USE_REAL_DATA'	,A_   ):
								A     = 'examples/seq2seq/wmt_en_ro'
								A     = max_len * 2 * 64
								if not Path(A_   ).joinpath('train.len'   ).exists():
									save_len_file(A_	,A_   )
							else:
								A     = 'examples/seq2seq/test_data/wmt_en_ro'
								A     = max_len * 4
								save_len_file(A_	,A_   )

							A     = AutoTokenizer.from_pretrained(A_   )
							A     = SeqaSeqDataset(
							    A_	,data_dir=A_	,type_path='train'	,max_source_length=A_	,max_target_length=A_	,n_obs=A_	,)
							return ds, max_tokens, tokenizer


						def     _SCREAMING_SNAKE_CASE   (			self  :   Optional[int]   )  -> Tuple:
							A ,   A ,   A     = self._get_dataset()
							A     = set(DistributedSortishSampler(A_	,256	,num_replicas=2	,rank=0	,add_extra_examples=A_   )   )
							A     = set(DistributedSortishSampler(A_	,256	,num_replicas=2	,rank=1	,add_extra_examples=A_   )   )
							assert idsa.intersection(A_   ) == set()




						@parameterized.expand(
						    [
						        MBART_TINY,
						        MARIAN_TINY,
						        T5_TINY,
						        BART_TINY,
						        PEGASUS_XSUM,
						    ]	,)
						def     _SCREAMING_SNAKE_CASE   (			self  :   str	,A_  :   Any   )  -> Optional[int]:
							A     = AutoTokenizer.from_pretrained(A_	,use_fast=A_   )
							if tok_name == MBART_TINY:
								A     = SeqaSeqDataset(
								    A_	,data_dir=make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()   )	,type_path='train'	,max_source_length=4	,max_target_length=8	,src_lang='EN'	,tgt_lang='FR'	,)
								A     = train_dataset.dataset_kwargs
								assert "src_lang" in kwargs and "tgt_lang" in kwargs
							else:
								A     = SeqaSeqDataset(
								    A_	,data_dir=make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()   )	,type_path='train'	,max_source_length=4	,max_target_length=8	,)
								A     = train_dataset.dataset_kwargs
								assert "add_prefix_space" not in kwargs if tok_name != BART_TINY else "add_prefix_space" in kwargs
								assert len(A_   ) == 1 if tok_name == BART_TINY else len(A_   ) == 0

"""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 					lowerCAmelCase_					:



						'''simple docstring'''


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


						def     _SCREAMING_SNAKE_CASE   (			self  :   List[Any]   )  -> Any:
							A     = ids_tensor([self.batch_size, self.seq_length]	,self.vocab_size   )

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

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

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

							A     = self.get_config()

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


						def     _SCREAMING_SNAKE_CASE   (			self  :   Optional[Any]   )  -> Any:
							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=A_	,initializer_range=self.initializer_range	,)


						def     _SCREAMING_SNAKE_CASE   (			self  :   Any	,A_  :   Dict	,A_  :   Optional[int]	,A_  :   Any	,A_  :   Optional[Any]	,A_  :   Any	,A_  :   Union[str, Any]	,A_  :   Tuple   )  -> List[Any]:
							A     = LlamaModel(config=A_   )
							model.to(A_   )
							model.eval()
							A     = model(A_	,attention_mask=A_   )
							A     = model(A_   )
							self.parent.assertEqual(result.last_hidden_state.shape	,(self.batch_size, self.seq_length, self.hidden_size)   )


						def     _SCREAMING_SNAKE_CASE   (			self  :   Tuple	,A_  :   int	,A_  :   Optional[Any]	,A_  :   Optional[int]	,A_  :   Any	,A_  :   Union[str, Any]	,A_  :   Any	,A_  :   Tuple	,A_  :   Union[str, Any]	,A_  :   Dict	,)  -> List[str]:
							A     = True
							A     = LlamaModel(A_   )
							model.to(A_   )
							model.eval()
							A     = model(
							    A_	,attention_mask=A_	,encoder_hidden_states=A_	,encoder_attention_mask=A_	,)
							A     = model(
							    A_	,attention_mask=A_	,encoder_hidden_states=A_	,)
							A     = model(A_	,attention_mask=A_   )
							self.parent.assertEqual(result.last_hidden_state.shape	,(self.batch_size, self.seq_length, self.hidden_size)   )


						def     _SCREAMING_SNAKE_CASE   (			self  :   int	,A_  :   List[str]	,A_  :   Optional[int]	,A_  :   Any	,A_  :   str	,A_  :   Dict	,A_  :   Dict	,A_  :   Tuple	,A_  :   Tuple	,A_  :   Dict	,)  -> Union[str, Any]:
							A     = LlamaForCausalLM(config=A_   )
							model.to(A_   )
							model.eval()
							A     = model(A_	,attention_mask=A_	,labels=A_   )
							self.parent.assertEqual(result.logits.shape	,(self.batch_size, self.seq_length, self.vocab_size)   )


						def     _SCREAMING_SNAKE_CASE   (			self  :   int	,A_  :   Dict	,A_  :   Any	,A_  :   int	,A_  :   List[str]	,A_  :   Tuple	,A_  :   Any	,A_  :   Union[str, Any]	,A_  :   Any	,A_  :   int	,)  -> List[Any]:
							A     = True
							A     = True
							A     = LlamaForCausalLM(config=A_   )
							model.to(A_   )
							model.eval()

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

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

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

							A     = model(
							    A_	,attention_mask=A_	,encoder_hidden_states=A_	,encoder_attention_mask=A_	,output_hidden_states=A_	,)['hidden_states'][0]
							A     = 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
							A     = ids_tensor((1,)	,output_from_past.shape[-1]   ).item()
							A     = output_from_no_past[:, -3:, random_slice_idx].detach()
							A     = 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     _SCREAMING_SNAKE_CASE   (			self  :   Any   )  -> Optional[int]:
							A     = self.prepare_config_and_inputs()
							(
							    (
							    A
							) ,   (
							    A
							) ,   (
							    A
							) ,   (
							    A
							) ,   (
							    A
							) ,   (
							    A
							) ,   (
							    A
							) ,   
							)     = config_and_inputs
							A     = {'input_ids': input_ids, 'attention_mask': input_mask}
							return config, inputs_dict







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



						'''simple docstring'''


						_lowerCamelCase:     Dict   =     (LlamaModel, LlamaForCausalLM, LlamaForSequenceClassification) if is_torch_available() else ()
						_lowerCamelCase:     List[Any]   =     (LlamaForCausalLM,) if is_torch_available() else ()
						_lowerCamelCase:     Any   =     (
						    {
						        '''feature-extraction''': LlamaModel,
						        '''text-classification''': LlamaForSequenceClassification,
						        '''text-generation''': LlamaForCausalLM,
						        '''zero-shot''': LlamaForSequenceClassification,
						    }
						    if is_torch_available()
						    else {}
						)
						_lowerCamelCase:     int   =     False
						_lowerCamelCase:     List[str]   =     False


						def     _SCREAMING_SNAKE_CASE   (			self  :   Any   )  -> Dict:
							A     = LlamaModelTester(self   )
							A     = ConfigTester(self	,config_class=A_	,hidden_size=37   )


						def     _SCREAMING_SNAKE_CASE   (			self  :   Union[str, Any]   )  -> Any:
							self.config_tester.run_common_tests()


						def     _SCREAMING_SNAKE_CASE   (			self  :   List[Any]   )  -> Tuple:
							A     = self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*A_   )


						def     _SCREAMING_SNAKE_CASE   (			self  :   str   )  -> Optional[int]:
							A     = self.model_tester.prepare_config_and_inputs()
							for type in ["absolute", "relative_key", "relative_key_query"]:
								A     = type
								self.model_tester.create_and_check_model(*A_   )


						def     _SCREAMING_SNAKE_CASE   (			self  :   List[Any]   )  -> List[Any]:
							A ,   A     = self.model_tester.prepare_config_and_inputs_for_common()
							A     = 3
							A     = input_dict['input_ids']
							A     = input_ids.ne(1   ).to(A_   )
							A     = ids_tensor([self.model_tester.batch_size]	,self.model_tester.type_sequence_label_size   )
							A     = LlamaForSequenceClassification(A_   )
							model.to(A_   )
							model.eval()
							A     = model(A_	,attention_mask=A_	,labels=A_   )
							self.assertEqual(result.logits.shape	,(self.model_tester.batch_size, self.model_tester.num_labels)   )


						def     _SCREAMING_SNAKE_CASE   (			self  :   int   )  -> Optional[Any]:
							A ,   A     = self.model_tester.prepare_config_and_inputs_for_common()
							A     = 3
							A     = 'single_label_classification'
							A     = input_dict['input_ids']
							A     = input_ids.ne(1   ).to(A_   )
							A     = ids_tensor([self.model_tester.batch_size]	,self.model_tester.type_sequence_label_size   )
							A     = LlamaForSequenceClassification(A_   )
							model.to(A_   )
							model.eval()
							A     = model(A_	,attention_mask=A_	,labels=A_   )
							self.assertEqual(result.logits.shape	,(self.model_tester.batch_size, self.model_tester.num_labels)   )


						def     _SCREAMING_SNAKE_CASE   (			self  :   str   )  -> Dict:
							A ,   A     = self.model_tester.prepare_config_and_inputs_for_common()
							A     = 3
							A     = 'multi_label_classification'
							A     = input_dict['input_ids']
							A     = input_ids.ne(1   ).to(A_   )
							A     = ids_tensor(
							    [self.model_tester.batch_size, config.num_labels]	,self.model_tester.type_sequence_label_size   ).to(torch.float   )
							A     = LlamaForSequenceClassification(A_   )
							model.to(A_   )
							model.eval()
							A     = model(A_	,attention_mask=A_	,labels=A_   )
							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     _SCREAMING_SNAKE_CASE   (			self  :   Dict   )  -> Optional[Any]:
							pass




						@parameterized.expand([('linear',), ('dynamic',)]   )
						def     _SCREAMING_SNAKE_CASE   (			self  :   List[Any]	,A_  :   Any   )  -> str:
							A ,   A     = self.model_tester.prepare_config_and_inputs_for_common()
							A     = ids_tensor([1, 10]	,config.vocab_size   )
							A     = 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
							A     = LlamaModel(A_   )
							original_model.to(A_   )
							original_model.eval()
							A     = original_model(A_   ).last_hidden_state
							A     = original_model(A_   ).last_hidden_state

							set_seed(42   )  # Fixed seed at init time so the two models get the same random weights
							A     = {'type': scaling_type, 'factor': 10.0}
							A     = LlamaModel(A_   )
							scaled_model.to(A_   )
							scaled_model.eval()
							A     = scaled_model(A_   ).last_hidden_state
							A     = 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   )   )







@require_torch
class 					lowerCAmelCase_					(   unittest.TestCase     ):



						'''simple docstring'''


						@unittest.skip('Logits are not exactly the same, once we fix the instabalities somehow, will update!'   )
						@slow
						def     _SCREAMING_SNAKE_CASE   (			self  :   List[str]   )  -> Optional[Any]:
							A     = [1, 306, 4658, 278, 6593, 310, 2834, 338]
							A     = LlamaForCausalLM.from_pretrained('meta-llama/Llama-2-7b-hf'	,device_map='auto'   )
							A     = model(torch.tensor([input_ids]   )   )
							# Expected mean on dim = -1
							A     = 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   )	,A_	,atol=1e-2	,rtol=1e-2   )
							# slicing logits[0, 0, 0:30]
							# fmt: off
							A     = 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]	,A_	,atol=1e-5	,rtol=1e-5   )


						@unittest.skip('Logits are not exactly the same, once we fix the instabalities somehow, will update!'   )
						@slow
						def     _SCREAMING_SNAKE_CASE   (			self  :   Any   )  -> int:
							A     = [1, 306, 4658, 278, 6593, 310, 2834, 338]
							A     = LlamaForCausalLM.from_pretrained('meta-llama/Llama-2-13b-hf'	,device_map='auto'   )
							A     = model(torch.tensor(A_   )   )
							# Expected mean on dim = -1
							A     = 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   )	,A_	,atol=1e-2	,rtol=1e-2   )
							# slicing logits[0, 0, 0:30]
							# fmt: off
							A     = 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]	,A_	,atol=1e-5	,rtol=1e-5   )


						@unittest.skip('Logits are not exactly the same, once we fix the instabalities somehow, will update!'   )
						@slow
						def     _SCREAMING_SNAKE_CASE   (			self  :   int   )  -> str:
							A     = [1, 306, 4658, 278, 6593, 310, 2834, 338]
							A     = LlamaForCausalLM.from_pretrained('meta-llama/Llama-2-13b-chat-hf'	,device_map='auto'   )
							A     = model(torch.tensor(A_   )   )
							# Expected mean on dim = -1
							A     = 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   )	,A_	,atol=1e-2	,rtol=1e-2   )
							# slicing logits[0, 0, 0:30]
							# fmt: off
							A     = 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   )	,A_	,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     _SCREAMING_SNAKE_CASE   (			self  :   Any   )  -> str:
							A     = [1, 306, 4658, 278, 6593, 310, 2834, 338]
							A     = LlamaForCausalLM.from_pretrained('meta-llama/Llama-2-70b-hf'	,device_map='auto'   )
							A     = model(torch.tensor(A_   )   )

							A     = 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   )	,A_	,atol=1e-2	,rtol=1e-2   )
							# fmt: off
							A     = 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]	,A_	,atol=1e-5	,rtol=1e-5   )




						@unittest.skip('Model is curently gated'   )
						@slow
						def     _SCREAMING_SNAKE_CASE   (			self  :   int   )  -> Tuple:
							A     = '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'
							A     = 'Simply put, the theory of relativity states that '
							A     = LlamaTokenizer.from_pretrained('meta-llama/Llama-2-13b-chat-hf'   )
							A     = tokenizer.encode(A_	,return_tensors='pt'   )
							A     = LlamaForCausalLM.from_pretrained(
							    'meta-llama/Llama-2-13b-chat-hf'	,device_map='sequential'	,use_safetensors=A_   )

							# greedy generation outputs
							A     = model.generate(A_	,max_new_tokens=64	,top_p=A_	,temperature=1	,do_sample=A_   )
							A     = tokenizer.decode(generated_ids[0]	,skip_special_tokens=A_   )
							self.assertEqual(A_	,A_   )