
'''simple docstring'''







import json
import os
from dataclasses import dataclass
from functools import partial
from typing import Callable

import flax.linen as nn
import jax
import jax.numpy as jnp
import joblib
import optax
import wandb
from flax import jax_utils, struct, traverse_util
from flax.serialization import from_bytes, to_bytes
from flax.training import train_state
from flax.training.common_utils import shard
from tqdm.auto import tqdm

from transformers import BigBirdConfig, FlaxBigBirdForQuestionAnswering
from transformers.models.big_bird.modeling_flax_big_bird import FlaxBigBirdForQuestionAnsweringModule


class 		__UpperCAmelCase   (      _lowerCamelCase		):
							__lowercase   	= 42
							__lowercase   	= jnp.floataa
							__lowercase   	= True


							def 				lowerCamelCase       ( self						):





								"""simple docstring"""






								super().setup()
								_snake_case     		=	nn.Dense(5					,     dtype=self.dtype						)




							def __call__( self					,     *lowerCAmelCase_					,     **lowerCAmelCase_						):





								"""simple docstring"""






								_snake_case     		=	super().__call__(*lowerCAmelCase_					,     **lowerCAmelCase_						)
								_snake_case     		=	self.cls(outputs[2]						)
								return outputs[:2] + (cls_out,)


class 		__UpperCAmelCase   (      _lowerCamelCase		):
							__lowercase   	= FlaxBigBirdForNaturalQuestionsModule







def        SCREAMING_SNAKE_CASE__			(     __A    ,       __A    ,       __A    ,       __A    ,       __A    ,       __A			) ->			Union[str, Any]:
	def cross_entropy(__A    ,       __A    ,       __A=None			):
		_snake_case     		=	logits.shape[-1]
		_snake_case     		=	(labels[..., None] == jnp.arange(__A			)[None]).astype('f4'			)
		_snake_case     		=	jax.nn.log_softmax(__A    ,       axis=-1			)
		_snake_case     		=	-jnp.sum(labels * logits    ,       axis=-1			)
		if reduction is not None:
			_snake_case     		=	reduction(__A			)
		return loss

	_snake_case     		=	partial(__A    ,       reduction=jnp.mean			)
	_snake_case     		=	cross_entropy(__A    ,       __A			)
	_snake_case     		=	cross_entropy(__A    ,       __A			)
	_snake_case     		=	cross_entropy(__A    ,       __A			)
	return (start_loss + end_loss + pooled_loss) / 3


@dataclass
class 		__UpperCAmelCase   :
							__lowercase   	= "google/bigbird-roberta-base"
							__lowercase   	= 30_00
							__lowercase   	= 1_05_00

							__lowercase   	= 1_28
							__lowercase   	= 3

							__lowercase   	= 1
							__lowercase   	= 5

							# tx_args
							__lowercase   	= 3e-5
							__lowercase   	= 0.0
							__lowercase   	= 2_00_00
							__lowercase   	= 0.0_0_9_5

							__lowercase   	= "bigbird-roberta-natural-questions"
							__lowercase   	= "training-expt"
							__lowercase   	= "data/nq-training.jsonl"
							__lowercase   	= "data/nq-validation.jsonl"


							def 				lowerCamelCase       ( self						):





								"""simple docstring"""






								os.makedirs(self.base_dir					,     exist_ok=lowerCAmelCase_						)
								_snake_case     		=	os.path.join(self.base_dir					,     self.save_dir						)
								_snake_case     		=	self.batch_size_per_device * jax.device_count()


@dataclass
class 		__UpperCAmelCase   :
							__lowercase   	= 42
							__lowercase   	= 40_96  # no dynamic padding on TPUs


							def __call__( self					,     lowerCAmelCase_						):





								"""simple docstring"""






								_snake_case     		=	self.collate_fn(lowerCAmelCase_						)
								_snake_case     		=	jax.tree_util.tree_map(lowerCAmelCase_					,     lowerCAmelCase_						)
								return batch


							def 				lowerCamelCase       ( self					,     lowerCAmelCase_						):





								"""simple docstring"""






								_snake_case		,				_snake_case     		=	self.fetch_inputs(features['input_ids']						)
								_snake_case     		=	{
								    'input_ids': jnp.array(lowerCAmelCase_					,     dtype=jnp.intaa						),
								    'attention_mask': jnp.array(lowerCAmelCase_					,     dtype=jnp.intaa						),
								    'start_labels': jnp.array(features['start_token']					,     dtype=jnp.intaa						),
								    'end_labels': jnp.array(features['end_token']					,     dtype=jnp.intaa						),
								    'pooled_labels': jnp.array(features['category']					,     dtype=jnp.intaa						),
								}
								return batch


							def 				lowerCamelCase       ( self					,     lowerCAmelCase_						):





								"""simple docstring"""






								_snake_case     		=	[self._fetch_inputs(lowerCAmelCase_						) for ids in input_ids]
								return zip(*lowerCAmelCase_						)




							def 				lowerCamelCase       ( self					,     lowerCAmelCase_						):





								"""simple docstring"""






								_snake_case     		=	[1 for _ in range(len(lowerCAmelCase_						)						)]
								while len(lowerCAmelCase_						) < self.max_length:
									input_ids.append(self.pad_id						)
									attention_mask.append(0						)
								return input_ids, attention_mask







def        SCREAMING_SNAKE_CASE__			(     __A    ,       __A    ,       __A=None			) ->			List[str]:
	if seed is not None:
		_snake_case     		=	dataset.shuffle(seed=__A			)
	for i in range(len(__A			) // batch_size			):
		_snake_case     		=	dataset[i * batch_size : (i + 1) * batch_size]
		yield dict(__A			)







@partial(jax.pmap    ,       axis_name='batch'			)
def        SCREAMING_SNAKE_CASE__			(     __A    ,       __A    ,       **__A			) ->			Dict:
	def loss_fn(__A			):
		_snake_case     		=	model_inputs.pop('start_labels'			)
		_snake_case     		=	model_inputs.pop('end_labels'			)
		_snake_case     		=	model_inputs.pop('pooled_labels'			)

		_snake_case     		=	state.apply_fn(**__A    ,       params=__A    ,       dropout_rng=__A    ,       train=__A			)
		_snake_case		,				_snake_case		,				_snake_case     		=	outputs

		return state.loss_fn(
		    __A    ,       __A    ,       __A    ,       __A    ,       __A    ,       __A    ,       )

	_snake_case		,				_snake_case     		=	jax.random.split(__A			)
	_snake_case     		=	jax.value_and_grad(__A			)
	_snake_case		,				_snake_case     		=	grad_fn(state.params			)
	_snake_case     		=	jax.lax.pmean({'loss': loss}    ,       axis_name='batch'			)
	_snake_case     		=	jax.lax.pmean(__A    ,       'batch'			)

	_snake_case     		=	state.apply_gradients(grads=__A			)
	return state, metrics, new_drp_rng







@partial(jax.pmap    ,       axis_name='batch'			)
def        SCREAMING_SNAKE_CASE__			(     __A    ,       **__A			) ->			List[str]:
	_snake_case     		=	model_inputs.pop('start_labels'			)
	_snake_case     		=	model_inputs.pop('end_labels'			)
	_snake_case     		=	model_inputs.pop('pooled_labels'			)

	_snake_case     		=	state.apply_fn(**__A    ,       params=state.params    ,       train=__A			)
	_snake_case		,				_snake_case		,				_snake_case     		=	outputs

	_snake_case     		=	state.loss_fn(__A    ,       __A    ,       __A    ,       __A    ,       __A    ,       __A			)
	_snake_case     		=	jax.lax.pmean({'loss': loss}    ,       axis_name='batch'			)
	return metrics


class 		__UpperCAmelCase   (      train_state.TrainState		):
							__lowercase   	= struct.field(pytree_node=_lowerCamelCase		)


@dataclass
class 		__UpperCAmelCase   :
							__lowercase   	= 42
							__lowercase   	= 42
							__lowercase   	= 42
							__lowercase   	= 42
							__lowercase   	= 42
							__lowercase   	= 42
							__lowercase   	= None


							def 				lowerCamelCase       ( self					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_=None						):





								"""simple docstring"""






								_snake_case     		=	model.params
								_snake_case     		=	TrainState.create(
								    apply_fn=model.__call__					,     params=lowerCAmelCase_					,     tx=lowerCAmelCase_					,     loss_fn=lowerCAmelCase_					,     )
								if ckpt_dir is not None:
									_snake_case		,				_snake_case		,				_snake_case		,				_snake_case		,				_snake_case     		=	restore_checkpoint(lowerCAmelCase_					,     lowerCAmelCase_						)
									_snake_case     		=	{
									    'lr': args.lr,
									    'init_lr': args.init_lr,
									    'warmup_steps': args.warmup_steps,
									    'num_train_steps': num_train_steps,
									    'weight_decay': args.weight_decay,
									}
									_snake_case		,				_snake_case     		=	build_tx(**lowerCAmelCase_						)
									_snake_case     		=	train_state.TrainState(
									    step=lowerCAmelCase_					,     apply_fn=model.__call__					,     params=lowerCAmelCase_					,     tx=lowerCAmelCase_					,     opt_state=lowerCAmelCase_					,     )
									_snake_case     		=	args
									_snake_case     		=	data_collator
									_snake_case     		=	lr
									_snake_case     		=	params
								_snake_case     		=	jax_utils.replicate(lowerCAmelCase_						)
								return state


							def 				lowerCamelCase       ( self					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_						):





								"""simple docstring"""






								_snake_case     		=	self.args
								_snake_case     		=	len(lowerCAmelCase_						) // args.batch_size

								_snake_case     		=	jax.random.PRNGKey(0						)
								_snake_case     		=	jax.random.split(lowerCAmelCase_					,     jax.device_count()						)
								for epoch in range(args.max_epochs						):
									_snake_case     		=	jnp.array(0					,     dtype=jnp.floataa						)
									_snake_case     		=	get_batched_dataset(lowerCAmelCase_					,     args.batch_size					,     seed=lowerCAmelCase_						)
									_snake_case     		=	0
									for batch in tqdm(lowerCAmelCase_					,     total=lowerCAmelCase_					,     desc=F'Running EPOCH-{epoch}'						):
										_snake_case     		=	self.data_collator(lowerCAmelCase_						)
										_snake_case		,				_snake_case		,				_snake_case     		=	self.train_step_fn(lowerCAmelCase_					,     lowerCAmelCase_					,     **lowerCAmelCase_						)
										running_loss += jax_utils.unreplicate(metrics['loss']						)
										i += 1
										if i % args.logging_steps == 0:
											_snake_case     		=	jax_utils.unreplicate(state.step						)
											_snake_case     		=	running_loss.item() / i
											_snake_case     		=	self.scheduler_fn(state_step - 1						)

											_snake_case     		=	self.evaluate(lowerCAmelCase_					,     lowerCAmelCase_						)
											_snake_case     		=	{
											    'step': state_step.item(),
											    'eval_loss': eval_loss.item(),
											    'tr_loss': tr_loss,
											    'lr': lr.item(),
											}
											tqdm.write(str(lowerCAmelCase_						)						)
											self.logger.log(lowerCAmelCase_					,     commit=lowerCAmelCase_						)

										if i % args.save_steps == 0:
											self.save_checkpoint(args.save_dir + F'-e{epoch}-s{i}'					,     state=lowerCAmelCase_						)


							def 				lowerCamelCase       ( self					,     lowerCAmelCase_					,     lowerCAmelCase_						):





								"""simple docstring"""






								_snake_case     		=	get_batched_dataset(lowerCAmelCase_					,     self.args.batch_size						)
								_snake_case     		=	len(lowerCAmelCase_						) // self.args.batch_size
								_snake_case     		=	jnp.array(0					,     dtype=jnp.floataa						)
								_snake_case     		=	0
								for batch in tqdm(lowerCAmelCase_					,     total=lowerCAmelCase_					,     desc='Evaluating ... '						):
									_snake_case     		=	self.data_collator(lowerCAmelCase_						)
									_snake_case     		=	self.val_step_fn(lowerCAmelCase_					,     **lowerCAmelCase_						)
									running_loss += jax_utils.unreplicate(metrics['loss']						)
									i += 1
								return running_loss / i




							def 				lowerCamelCase       ( self					,     lowerCAmelCase_					,     lowerCAmelCase_						):





								"""simple docstring"""






								_snake_case     		=	jax_utils.unreplicate(lowerCAmelCase_						)
								print(F'SAVING CHECKPOINT IN {save_dir}'					,     end=' ... '						)
								self.model_save_fn(lowerCAmelCase_					,     params=state.params						)
								with open(os.path.join(lowerCAmelCase_					,     'opt_state.msgpack'						)					,     'wb'						) as f:
									f.write(to_bytes(state.opt_state						)						)
								joblib.dump(self.args					,     os.path.join(lowerCAmelCase_					,     'args.joblib'						)						)
								joblib.dump(self.data_collator					,     os.path.join(lowerCAmelCase_					,     'data_collator.joblib'						)						)
								with open(os.path.join(lowerCAmelCase_					,     'training_state.json'						)					,     'w'						) as f:
									json.dump({'step': state.step.item()}					,     lowerCAmelCase_						)
								print('DONE'						)







def        SCREAMING_SNAKE_CASE__			(     __A    ,       __A			) ->			Tuple:
	print(F'RESTORING CHECKPOINT FROM {save_dir}'    ,       end=' ... '			)
	with open(os.path.join(__A    ,       'flax_model.msgpack'			)    ,       'rb'			) as f:
		_snake_case     		=	from_bytes(state.params    ,       f.read()			)

	with open(os.path.join(__A    ,       'opt_state.msgpack'			)    ,       'rb'			) as f:
		_snake_case     		=	from_bytes(state.opt_state    ,       f.read()			)

	_snake_case     		=	joblib.load(os.path.join(__A    ,       'args.joblib'			)			)
	_snake_case     		=	joblib.load(os.path.join(__A    ,       'data_collator.joblib'			)			)

	with open(os.path.join(__A    ,       'training_state.json'			)    ,       'r'			) as f:
		_snake_case     		=	json.load(__A			)
	_snake_case     		=	training_state['step']

	print('DONE'			)
	return params, opt_state, step, args, data_collator







def        SCREAMING_SNAKE_CASE__			(     __A    ,       __A    ,       __A    ,       __A			) ->			List[Any]:
	_snake_case     		=	num_train_steps - warmup_steps
	_snake_case     		=	optax.linear_schedule(init_value=__A    ,       end_value=__A    ,       transition_steps=__A			)
	_snake_case     		=	optax.linear_schedule(init_value=__A    ,       end_value=1e-7    ,       transition_steps=__A			)
	_snake_case     		=	optax.join_schedules(schedules=[warmup_fn, decay_fn]    ,       boundaries=[warmup_steps]			)
	return lr







def        SCREAMING_SNAKE_CASE__			(     __A    ,       __A    ,       __A    ,       __A    ,       __A			) ->			List[str]:
	def weight_decay_mask(__A			):
		_snake_case     		=	traverse_util.flatten_dict(__A			)
		_snake_case     		=	{k: (v[-1] != 'bias' and v[-2:] != ('LayerNorm', 'scale')) for k, v in params.items()}
		return traverse_util.unflatten_dict(__A			)

	_snake_case     		=	scheduler_fn(__A    ,       __A    ,       __A    ,       __A			)

	_snake_case     		=	optax.adamw(learning_rate=__A    ,       weight_decay=__A    ,       mask=__A			)
	return tx, lr



'''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 		__UpperCAmelCase   :


							def __init__( self					,     lowerCAmelCase_					,     lowerCAmelCase_=13					,     lowerCAmelCase_=7					,     lowerCAmelCase_=True					,     lowerCAmelCase_=True					,     lowerCAmelCase_=False					,     lowerCAmelCase_=True					,     lowerCAmelCase_=99					,     lowerCAmelCase_=32					,     lowerCAmelCase_=5					,     lowerCAmelCase_=4					,     lowerCAmelCase_=37					,     lowerCAmelCase_="gelu"					,     lowerCAmelCase_=0.1					,     lowerCAmelCase_=0.1					,     lowerCAmelCase_=5_12					,     lowerCAmelCase_=16					,     lowerCAmelCase_=2					,     lowerCAmelCase_=0.02					,     lowerCAmelCase_=3					,     lowerCAmelCase_=4					,     lowerCAmelCase_=None					,     ):





								"""simple docstring"""






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


							def 				lowerCamelCase       ( self						):





								"""simple docstring"""






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

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

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

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

								_snake_case     		=	self.get_config()

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


							def 				lowerCamelCase       ( self						):





								"""simple docstring"""






								return LlamaConfig(
								    vocab_size=self.vocab_size					,     hidden_size=self.hidden_size					,     num_hidden_layers=self.num_hidden_layers					,     num_attention_heads=self.num_attention_heads					,     intermediate_size=self.intermediate_size					,     hidden_act=self.hidden_act					,     hidden_dropout_prob=self.hidden_dropout_prob					,     attention_probs_dropout_prob=self.attention_probs_dropout_prob					,     max_position_embeddings=self.max_position_embeddings					,     type_vocab_size=self.type_vocab_size					,     is_decoder=lowerCAmelCase_					,     initializer_range=self.initializer_range					,     )


							def 				lowerCamelCase       ( self					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_						):





								"""simple docstring"""






								_snake_case     		=	LlamaModel(config=lowerCAmelCase_						)
								model.to(lowerCAmelCase_						)
								model.eval()
								_snake_case     		=	model(lowerCAmelCase_					,     attention_mask=lowerCAmelCase_						)
								_snake_case     		=	model(lowerCAmelCase_						)
								self.parent.assertEqual(result.last_hidden_state.shape					,     (self.batch_size, self.seq_length, self.hidden_size)						)


							def 				lowerCamelCase       ( self					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     ):





								"""simple docstring"""






								_snake_case     		=	True
								_snake_case     		=	LlamaModel(lowerCAmelCase_						)
								model.to(lowerCAmelCase_						)
								model.eval()
								_snake_case     		=	model(
								    lowerCAmelCase_					,     attention_mask=lowerCAmelCase_					,     encoder_hidden_states=lowerCAmelCase_					,     encoder_attention_mask=lowerCAmelCase_					,     )
								_snake_case     		=	model(
								    lowerCAmelCase_					,     attention_mask=lowerCAmelCase_					,     encoder_hidden_states=lowerCAmelCase_					,     )
								_snake_case     		=	model(lowerCAmelCase_					,     attention_mask=lowerCAmelCase_						)
								self.parent.assertEqual(result.last_hidden_state.shape					,     (self.batch_size, self.seq_length, self.hidden_size)						)


							def 				lowerCamelCase       ( self					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     ):





								"""simple docstring"""






								_snake_case     		=	LlamaForCausalLM(config=lowerCAmelCase_						)
								model.to(lowerCAmelCase_						)
								model.eval()
								_snake_case     		=	model(lowerCAmelCase_					,     attention_mask=lowerCAmelCase_					,     labels=lowerCAmelCase_						)
								self.parent.assertEqual(result.logits.shape					,     (self.batch_size, self.seq_length, self.vocab_size)						)


							def 				lowerCamelCase       ( self					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     ):





								"""simple docstring"""






								_snake_case     		=	True
								_snake_case     		=	True
								_snake_case     		=	LlamaForCausalLM(config=lowerCAmelCase_						)
								model.to(lowerCAmelCase_						)
								model.eval()

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

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

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

								_snake_case     		=	model(
								    lowerCAmelCase_					,     attention_mask=lowerCAmelCase_					,     encoder_hidden_states=lowerCAmelCase_					,     encoder_attention_mask=lowerCAmelCase_					,     output_hidden_states=lowerCAmelCase_					,     )['hidden_states'][0]
								_snake_case     		=	model(
								    lowerCAmelCase_					,     attention_mask=lowerCAmelCase_					,     encoder_hidden_states=lowerCAmelCase_					,     encoder_attention_mask=lowerCAmelCase_					,     past_key_values=lowerCAmelCase_					,     output_hidden_states=lowerCAmelCase_					,     )['hidden_states'][0]

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




							def 				lowerCamelCase       ( self						):





								"""simple docstring"""






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


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


							def 				lowerCamelCase       ( self						):





								"""simple docstring"""






								_snake_case     		=	LlamaModelTester(self						)
								_snake_case     		=	ConfigTester(self					,     config_class=lowerCAmelCase_					,     hidden_size=37						)


							def 				lowerCamelCase       ( self						):





								"""simple docstring"""






								self.config_tester.run_common_tests()


							def 				lowerCamelCase       ( self						):





								"""simple docstring"""






								_snake_case     		=	self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*lowerCAmelCase_						)


							def 				lowerCamelCase       ( self						):





								"""simple docstring"""






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


							def 				lowerCamelCase       ( self						):





								"""simple docstring"""






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


							def 				lowerCamelCase       ( self						):





								"""simple docstring"""






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


							def 				lowerCamelCase       ( self						):





								"""simple docstring"""






								_snake_case		,				_snake_case     		=	self.model_tester.prepare_config_and_inputs_for_common()
								_snake_case     		=	3
								_snake_case     		=	'multi_label_classification'
								_snake_case     		=	input_dict['input_ids']
								_snake_case     		=	input_ids.ne(1						).to(lowerCAmelCase_						)
								_snake_case     		=	ids_tensor(
								    [self.model_tester.batch_size, config.num_labels]					,     self.model_tester.type_sequence_label_size						).to(torch.float						)
								_snake_case     		=	LlamaForSequenceClassification(lowerCAmelCase_						)
								model.to(lowerCAmelCase_						)
								model.eval()
								_snake_case     		=	model(lowerCAmelCase_					,     attention_mask=lowerCAmelCase_					,     labels=lowerCAmelCase_						)
								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 				lowerCamelCase       ( self						):





								"""simple docstring"""






								pass




							@parameterized.expand([('linear',), ('dynamic',)]						)
							def 				lowerCamelCase       ( self					,     lowerCAmelCase_						):





								"""simple docstring"""






								_snake_case		,				_snake_case     		=	self.model_tester.prepare_config_and_inputs_for_common()
								_snake_case     		=	ids_tensor([1, 10]					,     config.vocab_size						)
								_snake_case     		=	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
								_snake_case     		=	LlamaModel(lowerCAmelCase_						)
								original_model.to(lowerCAmelCase_						)
								original_model.eval()
								_snake_case     		=	original_model(lowerCAmelCase_						).last_hidden_state
								_snake_case     		=	original_model(lowerCAmelCase_						).last_hidden_state

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

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


@require_torch
class 		__UpperCAmelCase   (      unittest.TestCase		):


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





								"""simple docstring"""






								_snake_case     		=	[1, 3_06, 46_58, 2_78, 65_93, 3_10, 28_34, 3_38]
								_snake_case     		=	LlamaForCausalLM.from_pretrained('meta-llama/Llama-2-7b-hf'					,     device_map='auto'						)
								_snake_case     		=	model(torch.tensor([input_ids]						)						)
								# Expected mean on dim = -1
								_snake_case     		=	torch.tensor([[-6.6550, -4.1227, -4.9859, -3.2406, 0.8262, -3.0033, 1.2964, -3.3699]]						)
								torch.testing.assert_close(out.mean(-1						)					,     lowerCAmelCase_					,     atol=1E-2					,     rtol=1E-2						)
								# slicing logits[0, 0, 0:30]
								# fmt: off
								_snake_case     		=	torch.tensor([-12.8281, -7.4453, -0.4639, -8.0625, -7.2500, -8.0000, -6.4883, -7.7695, -7.8438, -7.0312, -6.2188, -7.1328, -1.8496, 1.9961, -8.6250, -6.7227, -12.8281, -6.9492, -7.0742, -7.7852, -7.5820, -7.9062, -6.9375, -7.9805, -8.3438, -8.1562, -8.0469, -7.6250, -7.7422, -7.3398,]						)
								# fmt: on
								torch.testing.assert_close(out[0, 0, :30]					,     lowerCAmelCase_					,     atol=1E-5					,     rtol=1E-5						)


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





								"""simple docstring"""






								_snake_case     		=	[1, 3_06, 46_58, 2_78, 65_93, 3_10, 28_34, 3_38]
								_snake_case     		=	LlamaForCausalLM.from_pretrained('meta-llama/Llama-2-13b-hf'					,     device_map='auto'						)
								_snake_case     		=	model(torch.tensor(lowerCAmelCase_						)						)
								# Expected mean on dim = -1
								_snake_case     		=	torch.tensor([[-2.0622, -1.2794, -1.1638, -0.9788, -1.4603, -1.0238, -1.7893, -1.4411]]						)
								torch.testing.assert_close(out.mean(-1						)					,     lowerCAmelCase_					,     atol=1E-2					,     rtol=1E-2						)
								# slicing logits[0, 0, 0:30]
								# fmt: off
								_snake_case     		=	torch.tensor([-8.1406, -8.0547, 2.7461, -1.2344, -0.1448, -1.8262, -1.0020, -1.8154, -1.6895, -1.8516, -2.3574, -0.9277, 3.7598, 6.5742, -1.2998, -0.1177, -8.1406, -2.9688, -2.9199, -3.1699, -3.5254, -2.3555, -2.7988, -3.4141, -2.8262, -4.5195, -3.3379, -3.3164, -2.7832, -3.0273]						)
								# fmt: on
								torch.testing.assert_close(out[0, 0, :30]					,     lowerCAmelCase_					,     atol=1E-5					,     rtol=1E-5						)


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





								"""simple docstring"""






								_snake_case     		=	[1, 3_06, 46_58, 2_78, 65_93, 3_10, 28_34, 3_38]
								_snake_case     		=	LlamaForCausalLM.from_pretrained('meta-llama/Llama-2-13b-chat-hf'					,     device_map='auto'						)
								_snake_case     		=	model(torch.tensor(lowerCAmelCase_						)						)
								# Expected mean on dim = -1
								_snake_case     		=	torch.tensor([[-0.8562, -1.8520, -0.7551, -0.4162, -1.5161, -1.2038, -2.4823, -2.3254]]						)
								torch.testing.assert_close(out.mean(-1						)					,     lowerCAmelCase_					,     atol=1E-2					,     rtol=1E-2						)
								# slicing logits[0, 0, 0:30]
								# fmt: off
								_snake_case     		=	torch.tensor([-2.2227, 4.8828, 0.9023, -0.4578, -0.7871, -0.1033, -0.6221, -0.5786, -0.7803, -1.0674, -1.2920, -0.1570, 0.8008, 2.0723, -0.9497, 0.2771, -2.2227, -0.7612, -1.4346, -1.2061, -1.6426, -0.3000, -0.7139, -1.1934, -1.8691, -1.6973, -1.5947, -1.2705, -0.3523, -0.5513]						)
								# fmt: on
								torch.testing.assert_close(out.mean(-1						)					,     lowerCAmelCase_					,     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 				lowerCamelCase       ( self						):





								"""simple docstring"""






								_snake_case     		=	[1, 3_06, 46_58, 2_78, 65_93, 3_10, 28_34, 3_38]
								_snake_case     		=	LlamaForCausalLM.from_pretrained('meta-llama/Llama-2-70b-hf'					,     device_map='auto'						)
								_snake_case     		=	model(torch.tensor(lowerCAmelCase_						)						)

								_snake_case     		=	torch.tensor(
								    [[-4.2327, -3.3360, -4.6665, -4.7631, -1.8180, -3.4170, -1.4211, -3.1810]]					,     dtype=torch.floataa						)
								torch.testing.assert_close(out.mean(-1						)					,     lowerCAmelCase_					,     atol=1E-2					,     rtol=1E-2						)
								# fmt: off
								_snake_case     		=	torch.tensor([-9.4922, -3.9551, 1.7998, -5.6758, -5.1055, -5.8984, -4.8320, -6.8086, -6.5391, -5.6172, -5.5820, -5.5352, 1.7881, 3.6289, -6.5117, -3.4785, -9.5000, -6.0352, -6.8125, -6.0195, -6.6836, -5.4727, -6.2812, -6.0391, -7.3398, -7.4297, -7.4844, -6.5820, -5.8789, -5.5312]						)
								# fmt: on
								torch.testing.assert_close(out[0, 0, :30]					,     lowerCAmelCase_					,     atol=1E-5					,     rtol=1E-5						)




							@unittest.skip('Model is curently gated'						)
							@slow
							def 				lowerCamelCase       ( self						):





								"""simple docstring"""






								_snake_case     		=	'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'
								_snake_case     		=	'Simply put, the theory of relativity states that '
								_snake_case     		=	LlamaTokenizer.from_pretrained('meta-llama/Llama-2-13b-chat-hf'						)
								_snake_case     		=	tokenizer.encode(lowerCAmelCase_					,     return_tensors='pt'						)
								_snake_case     		=	LlamaForCausalLM.from_pretrained(
								    'meta-llama/Llama-2-13b-chat-hf'					,     device_map='sequential'					,     use_safetensors=lowerCAmelCase_						)

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

