"""simple docstring"""
import unittest
from unittest.util import safe_repr

from transformers import AutoTokenizer, RwkvConfig, is_torch_available
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, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
					import torch

					from transformers import (
					    RWKV_PRETRAINED_MODEL_ARCHIVE_LIST,
					    RwkvForCausalLM,
					    RwkvModel,
					)
					from transformers.pytorch_utils import is_torch_greater_or_equal_than_2_0
else:
					SCREAMING_SNAKE_CASE__       =	False



class 	lowerCAmelCase_     :



				"""simple docstring"""

				def __init__(					self		,     lowerCAmelCase		,     lowerCAmelCase=14		,     lowerCAmelCase=7		,     lowerCAmelCase=True		,     lowerCAmelCase=False		,     lowerCAmelCase=True		,     lowerCAmelCase=True		,     lowerCAmelCase=True		,     lowerCAmelCase=99		,     lowerCAmelCase=32		,     lowerCAmelCase=5		,     lowerCAmelCase=37		,     lowerCAmelCase="gelu"		,     lowerCAmelCase=0.1		,     lowerCAmelCase=0.1		,     lowerCAmelCase=5_12		,     lowerCAmelCase=16		,     lowerCAmelCase=2		,     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_token_type_ids
										snake_case             =      use_input_mask
										snake_case             =      use_labels
										snake_case             =      use_mc_token_ids
										snake_case             =      vocab_size
										snake_case             =      hidden_size
										snake_case             =      num_hidden_layers
										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             =      num_labels
										snake_case             =      num_choices
										snake_case             =      scope
										snake_case             =      vocab_size - 1
										snake_case             =      vocab_size - 1
										snake_case             =      vocab_size - 1

				def 		snake_case			(					self     ):






										"""simple docstring"""

										return RwkvConfig.from_pretrained('sgugger/rwkv-4-pile-7b'     )

				def 		snake_case			(					self		,     lowerCAmelCase=False		,     lowerCAmelCase=False		,     lowerCAmelCase=False     ):






										"""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
										if self.use_mc_token_ids:
																snake_case             =      ids_tensor([self.batch_size, self.num_choices]		,     self.seq_length     )

										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(
										    gradient_checkpointing=lowerCAmelCase		,     scale_attn_by_inverse_layer_idx=lowerCAmelCase		,     reorder_and_upcast_attn=lowerCAmelCase		,     )

										return (
										    config,
										    input_ids,
										    input_mask,
										    None,
										    token_type_ids,
										    mc_token_ids,
										    sequence_labels,
										    token_labels,
										    choice_labels,
										)

				def 		snake_case			(					self		,     lowerCAmelCase=False		,     lowerCAmelCase=False		,     lowerCAmelCase=False     ):






										"""simple docstring"""

										return RwkvConfig(
										    vocab_size=self.vocab_size		,     hidden_size=self.hidden_size		,     num_hidden_layers=self.num_hidden_layers		,     intermediate_size=self.intermediate_size		,     activation_function=self.hidden_act		,     resid_pdrop=self.hidden_dropout_prob		,     attn_pdrop=self.attention_probs_dropout_prob		,     n_positions=self.max_position_embeddings		,     type_vocab_size=self.type_vocab_size		,     use_cache=lowerCAmelCase		,     bos_token_id=self.bos_token_id		,     eos_token_id=self.eos_token_id		,     pad_token_id=self.pad_token_id		,     gradient_checkpointing=lowerCAmelCase		,     scale_attn_by_inverse_layer_idx=lowerCAmelCase		,     reorder_and_upcast_attn=lowerCAmelCase		,     )

				def 		snake_case			(					self     ):






										"""simple docstring"""

										snake_case             =      self.get_config()
										snake_case             =      3_00
										return config

				def 		snake_case			(					self     ):






										"""simple docstring"""

										(
										    (
										    snake_case
										)       ,(
										    snake_case
										)       ,(
										    snake_case
										)       ,(
										    snake_case
										)       ,(
										    snake_case
										)       ,(
										    snake_case
										)       ,(
										    snake_case
										)       ,(
										    snake_case
										)       ,(
										    snake_case
										)       ,
										)             =      self.prepare_config_and_inputs()

										snake_case             =      floats_tensor([self.batch_size, self.seq_length, self.hidden_size]     )
										snake_case             =      ids_tensor([self.batch_size, self.seq_length]		,     vocab_size=2     )

										return (
										    config,
										    input_ids,
										    input_mask,
										    head_mask,
										    token_type_ids,
										    sequence_labels,
										    token_labels,
										    choice_labels,
										    encoder_hidden_states,
										    encoder_attention_mask,
										)

				def 		snake_case			(					self		,     lowerCAmelCase		,     lowerCAmelCase		,     lowerCAmelCase		,     lowerCAmelCase		,     lowerCAmelCase		,     *lowerCAmelCase     ):






										"""simple docstring"""

										snake_case             =      True
										snake_case             =      RwkvModel(config=lowerCAmelCase     )
										model.to(lowerCAmelCase     )
										model.eval()

										snake_case             =      model(lowerCAmelCase     )

										self.parent.assertEqual(result.last_hidden_state.shape		,     (self.batch_size, self.seq_length, self.hidden_size)     )
										self.parent.assertEqual(len(result.hidden_states     )		,     config.num_hidden_layers + 1     )

				def 		snake_case			(					self		,     lowerCAmelCase		,     lowerCAmelCase		,     lowerCAmelCase		,     lowerCAmelCase		,     lowerCAmelCase		,     *lowerCAmelCase     ):






										"""simple docstring"""

										snake_case             =      RwkvForCausalLM(lowerCAmelCase     )
										model.to(lowerCAmelCase     )
										model.eval()

										snake_case             =      model(lowerCAmelCase		,     labels=lowerCAmelCase     )
										self.parent.assertEqual(result.loss.shape		,     ()     )
										self.parent.assertEqual(result.logits.shape		,     (self.batch_size, self.seq_length, self.vocab_size)     )

				def 		snake_case			(					self		,     lowerCAmelCase		,     lowerCAmelCase		,     lowerCAmelCase		,     lowerCAmelCase		,     lowerCAmelCase		,     *lowerCAmelCase     ):






										"""simple docstring"""

										snake_case             =      RwkvModel(config=lowerCAmelCase     )
										model.to(lowerCAmelCase     )
										model.eval()

										snake_case             =      model(lowerCAmelCase     )
										snake_case             =      outputs.last_hidden_state

										snake_case             =      model(input_ids[:, :2]     )
										snake_case             =      outputs.last_hidden_state

										# Using the state computed on the first inputs, we will get the same output
										snake_case             =      model(input_ids[:, 2:]		,     state=outputs.state     )
										snake_case             =      outputs.last_hidden_state

										self.parent.assertTrue(torch.allclose(torch.cat([output_one, output_two]		,     dim=1     )		,     lowerCAmelCase		,     atol=1E-5     )     )

				def 		snake_case			(					self		,     lowerCAmelCase		,     lowerCAmelCase		,     lowerCAmelCase		,     lowerCAmelCase		,     lowerCAmelCase		,     *lowerCAmelCase		,     lowerCAmelCase=False     ):






										"""simple docstring"""

										snake_case             =      RwkvForCausalLM(lowerCAmelCase     )
										model.to(lowerCAmelCase     )
										if gradient_checkpointing:
																model.gradient_checkpointing_enable()

										snake_case             =      model(lowerCAmelCase		,     labels=lowerCAmelCase     )
										self.parent.assertEqual(result.loss.shape		,     ()     )
										self.parent.assertEqual(result.logits.shape		,     (self.batch_size, self.seq_length, self.vocab_size)     )
										result.loss.backward()

				def 		snake_case			(					self     ):






										"""simple docstring"""

										snake_case             =      self.prepare_config_and_inputs()

										(
										    (
										    snake_case
										)       ,(
										    snake_case
										)       ,(
										    snake_case
										)       ,(
										    snake_case
										)       ,(
										    snake_case
										)       ,(
										    snake_case
										)       ,(
										    snake_case
										)       ,(
										    snake_case
										)       ,(
										    snake_case
										)       ,
										)             =      config_and_inputs

										snake_case             =      {'input_ids': input_ids}

										return config, inputs_dict





@unittest.skipIf(
    not is_torch_greater_or_equal_than_2_0  ,     reason="""See https://github.com/huggingface/transformers/pull/24204"""       )
@require_torch
class 	lowerCAmelCase_     (      lowerCAmelCase  ,     lowerCAmelCase  ,     lowerCAmelCase  ,     unittest.TestCase       ):



				"""simple docstring"""
				_lowerCAmelCase  :     Optional[int]      	=    (RwkvModel, RwkvForCausalLM) if is_torch_available() else ()
				_lowerCAmelCase  :     List[str]      	=    (
				    {"""feature-extraction""": RwkvModel, """text-generation""": RwkvForCausalLM} if is_torch_available() else {}
				)
				# all_generative_model_classes = (RwkvForCausalLM,) if is_torch_available() else ()
				_lowerCAmelCase  :     Union[str, Any]      	=    False
				_lowerCAmelCase  :     Optional[int]      	=    False
				_lowerCAmelCase  :     List[str]      	=    False
				_lowerCAmelCase  :     Optional[int]      	=    False
				_lowerCAmelCase  :     Union[str, Any]      	=    False  # Rwkv does not support head masking

				def 		snake_case			(					self     ):






										"""simple docstring"""

										snake_case             =      RwkvModelTester(self     )
										snake_case             =      ConfigTester(
										    self		,     config_class=lowerCAmelCase		,     n_embd=37		,     common_properties=['hidden_size', 'num_hidden_layers']     )

				def 		snake_case			(					self		,     lowerCAmelCase		,     lowerCAmelCase		,     lowerCAmelCase=None     ):






										"""simple docstring"""

										if isinstance(lowerCAmelCase		,     torch.Tensor     ):
																snake_case       ,snake_case             =      member.max().item(), member.min().item()
										elif isinstance(lowerCAmelCase		,     lowerCAmelCase     ) or isinstance(lowerCAmelCase		,     lowerCAmelCase     ):
																snake_case       ,snake_case             =      max(lowerCAmelCase     ), min(lowerCAmelCase     )

										if not isinstance(lowerCAmelCase		,     lowerCAmelCase     ):
																raise TypeError('container should be a list or tuple'     )
										elif len(lowerCAmelCase     ) != 2:
																raise ValueError('container should have 2 elements'     )

										snake_case       ,snake_case             =      container

										snake_case             =      (min_value >= expected_min) and (max_value <= expected_max)

										if not is_inside_interval:
																snake_case             =      '%s not found in %s' % (safe_repr(lowerCAmelCase     ), safe_repr(lowerCAmelCase     ))
																self.fail(self._formatMessage(lowerCAmelCase		,     lowerCAmelCase     )     )

				def 		snake_case			(					self     ):






										"""simple docstring"""

										self.config_tester.run_common_tests()

				def 		snake_case			(					self     ):






										"""simple docstring"""

										snake_case             =      self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_rwkv_model(*lowerCAmelCase     )

				def 		snake_case			(					self     ):






										"""simple docstring"""

										snake_case             =      self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_causl_lm(*lowerCAmelCase     )

				def 		snake_case			(					self     ):






										"""simple docstring"""

										snake_case             =      self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_state_equivalency(*lowerCAmelCase     )

				def 		snake_case			(					self     ):






										"""simple docstring"""

										snake_case       ,snake_case             =      self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
																snake_case             =      model_class(config=lowerCAmelCase     )
																for name, param in model.named_parameters():
																						if "time_decay" in name:
																												if param.requires_grad:
																																		self.assertTrue(param.data.max().item() == 3.0     )
																																		self.assertTrue(param.data.min().item() == -5.0     )
																						elif "time_first" in name:
																												if param.requires_grad:
																																		# check if it's a ones like
																																		self.assertTrue(torch.allclose(param.data		,     torch.ones_like(param.data     )		,     atol=1E-5		,     rtol=1E-5     )     )
																						elif any(x in name for x in ['time_mix_key', 'time_mix_receptance']     ):
																												if param.requires_grad:
																																		self.assertInterval(
																																		    param.data		,     [0.0, 1.0]		,     msg=F"""Parameter {name} of model {model_class} seems not properly initialized"""		,     )
																						elif "time_mix_value" in name:
																												if param.requires_grad:
																																		self.assertInterval(
																																		    param.data		,     [0.0, 1.3]		,     msg=F"""Parameter {name} of model {model_class} seems not properly initialized"""		,     )

				def 		snake_case			(					self     ):






										"""simple docstring"""

										snake_case       ,snake_case             =      self.model_tester.prepare_config_and_inputs_for_common()
										snake_case             =      True

										snake_case             =      getattr(self.model_tester		,     'seq_length'		,     lowerCAmelCase     )

										for model_class in self.all_model_classes:
																snake_case             =      True
																snake_case             =      False
																snake_case             =      True
																snake_case             =      model_class(lowerCAmelCase     )
																model.to(lowerCAmelCase     )
																model.eval()

																snake_case             =      self._prepare_for_class(lowerCAmelCase		,     lowerCAmelCase     )
																snake_case             =      inputs['input_ids'].shape[0]
																with torch.no_grad():
																						snake_case             =      model(**lowerCAmelCase     )
																snake_case             =      outputs.attentions
																self.assertEqual(len(lowerCAmelCase     )		,     self.model_tester.num_hidden_layers     )

																# check that output_attentions also work using config
																del inputs_dict["output_attentions"]
																snake_case             =      True
																snake_case             =      model_class(lowerCAmelCase     )
																model.to(lowerCAmelCase     )
																model.eval()

																snake_case             =      self._prepare_for_class(lowerCAmelCase		,     lowerCAmelCase     )
																snake_case             =      inputs['input_ids'].shape[0]
																with torch.no_grad():
																						snake_case             =      model(**lowerCAmelCase     )
																snake_case             =      outputs.attentions
																self.assertEqual(len(lowerCAmelCase     )		,     self.model_tester.num_hidden_layers     )

																self.assertListEqual(
																    list(attentions[0].shape[-3:]     )		,     [batch_size, seq_len, config.hidden_size]		,     )
																snake_case             =      len(lowerCAmelCase     )

																# Check attention is always last and order is fine
																snake_case             =      True
																snake_case             =      True
																snake_case             =      model_class(lowerCAmelCase     )
																model.to(lowerCAmelCase     )
																model.eval()

																snake_case             =      self._prepare_for_class(lowerCAmelCase		,     lowerCAmelCase     )
																snake_case             =      inputs['input_ids'].shape[0]
																with torch.no_grad():
																						snake_case             =      model(**lowerCAmelCase     )

																snake_case             =      1
																self.assertEqual(out_len + added_hidden_states		,     len(lowerCAmelCase     )     )

																snake_case             =      outputs.attentions

																self.assertEqual(len(lowerCAmelCase     )		,     self.model_tester.num_hidden_layers     )
																self.assertListEqual(
																    list(self_attentions[0].shape[-3:]     )		,     [batch_size, seq_len, config.hidden_size]		,     )

				@slow
				def 		snake_case			(					self     ):






										"""simple docstring"""

										for model_name in RWKV_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																snake_case             =      RwkvModel.from_pretrained(lowerCAmelCase     )
																self.assertIsNotNone(lowerCAmelCase     )





@unittest.skipIf(
    not is_torch_greater_or_equal_than_2_0  ,     reason="""See https://github.com/huggingface/transformers/pull/24204"""       )
@slow
class 	lowerCAmelCase_     (      unittest.TestCase       ):



				"""simple docstring"""

				def 		snake_case			(					self     ):






										"""simple docstring"""

										snake_case             =      'RWKV/rwkv-4-169m-pile'
										snake_case             =      AutoTokenizer.from_pretrained(self.model_id     )

				def 		snake_case			(					self     ):






										"""simple docstring"""

										snake_case             =      'Hello my name is Jasmine and I am a newbie to the'
										snake_case             =      RwkvForCausalLM.from_pretrained(self.model_id     ).to(lowerCAmelCase     )

										snake_case             =      self.tokenizer('Hello my name is'		,     return_tensors='pt'     ).input_ids.to(lowerCAmelCase     )
										snake_case             =      model.generate(lowerCAmelCase		,     max_new_tokens=10     )
										snake_case             =      self.tokenizer.decode(output[0].tolist()     )

										self.assertEqual(lowerCAmelCase		,     lowerCAmelCase     )

				def 		snake_case			(					self     ):






										"""simple docstring"""

										snake_case             =      'Hello my name is Jasmine and I am a newbie to the'

										snake_case             =      self.tokenizer('Hello my name is'		,     return_tensors='pt'     ).input_ids.to(lowerCAmelCase     )
										snake_case             =      RwkvForCausalLM.from_pretrained(self.model_id		,     torch_dtype=torch.bfloataa     ).to(lowerCAmelCase     )

										snake_case             =      model.generate(lowerCAmelCase		,     max_new_tokens=10     )
										snake_case             =      self.tokenizer.decode(output[0].tolist()     )

										self.assertEqual(lowerCAmelCase		,     lowerCAmelCase     )



"""simple docstring"""
from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_tokenizers_available,
    is_torch_available,
)


SCREAMING_SNAKE_CASE__       =	{
    "configuration_open_llama": ["OPEN_LLAMA_PRETRAINED_CONFIG_ARCHIVE_MAP", "OpenLlamaConfig"],
}

try:
					if not is_sentencepiece_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					SCREAMING_SNAKE_CASE__       =	["LlamaTokenizer"]

try:
					if not is_tokenizers_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					SCREAMING_SNAKE_CASE__       =	["LlamaTokenizerFast"]

try:
					if not is_torch_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					SCREAMING_SNAKE_CASE__       =	[
					    "OpenLlamaForCausalLM",
					    "OpenLlamaModel",
					    "OpenLlamaPreTrainedModel",
					    "OpenLlamaForSequenceClassification",
					]


if TYPE_CHECKING:
					from .configuration_open_llama import OPEN_LLAMA_PRETRAINED_CONFIG_ARCHIVE_MAP, OpenLlamaConfig

					try:
										if not is_sentencepiece_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from transformers import LlamaTokenizer

					try:
										if not is_tokenizers_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from transformers import LlamaTokenizerFast

					try:
										if not is_torch_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_open_llama import (
										    OpenLlamaForCausalLM,
										    OpenLlamaForSequenceClassification,
										    OpenLlamaModel,
										    OpenLlamaPreTrainedModel,
										)


else:
					import sys

					SCREAMING_SNAKE_CASE__       =	_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)


