




import csv
import json
import logging
import os
import re
import subprocess
import tempfile
import unittest
import zipfile
from pathlib import Path
from typing import Optional
from unittest import mock

import torch

# We use TF to parse the logs
from accelerate import Accelerator
from accelerate.test_utils.testing import (
    MockingTestCase,
    TempDirTestCase,
    require_comet_ml,
    require_tensorboard,
    require_wandb,
    skip,
)
from accelerate.tracking import CometMLTracker, GeneralTracker
from accelerate.utils import ProjectConfiguration, is_comet_ml_available, is_tensorboard_available


if is_comet_ml_available():
				from comet_ml import OfflineExperiment

if is_tensorboard_available():
				import struct

				import tensorboard.compat.proto.event_pba as event_pba

lowercase_			:       List[Any]						       =				logging.getLogger(__name__)






@require_tensorboard
class 					__lowerCAmelCase				( unittest.TestCase				):







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




									"""simple docstring"""


									_UpperCAmelCase									=  "test_project_with_config"
									with tempfile.TemporaryDirectory() as dirpath:
															_UpperCAmelCase									=  Accelerator(log_with="tensorboard"					,					project_dir=snake_case__   )
															_UpperCAmelCase									=  {"num_iterations": 12, "learning_rate": 1e-2, "some_boolean": False, "some_string": "some_value"}
															accelerator.init_trackers(snake_case__					,					snake_case__   )
															accelerator.end_training()
															for child in Path(F"""{dirpath}/{project_name}"""   ).glob("*/**"   ):
																					_UpperCAmelCase									=  list(filter(lambda snake_case__   : x.is_file()					,					child.iterdir()   )   )[0]
															self.assertNotEqual(str(snake_case__   )					,					""   )







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




									"""simple docstring"""


									_UpperCAmelCase									=  "test_project_with_log"
									with tempfile.TemporaryDirectory() as dirpath:
															_UpperCAmelCase									=  Accelerator(log_with="tensorboard"					,					project_dir=snake_case__   )
															accelerator.init_trackers(snake_case__   )
															_UpperCAmelCase									=  {"total_loss": 0.1, "iteration": 1, "my_text": "some_value"}
															accelerator.log(snake_case__					,					step=0   )
															accelerator.end_training()
															# Logged values are stored in the outermost-tfevents file and can be read in as a TFRecord
															# Names are randomly generated each time
															_UpperCAmelCase									=  list(filter(lambda snake_case__   : x.is_file()					,					Path(F"""{dirpath}/{project_name}"""   ).iterdir()   )   )[0]
															self.assertNotEqual(str(snake_case__   )					,					""   )







			def UpperCamelCase    ( self					:    Optional[Any]   ):




									"""simple docstring"""


									_UpperCAmelCase									=  "test_project_with_log"
									with tempfile.TemporaryDirectory() as dirpath:
															_UpperCAmelCase									=  Accelerator(log_with="tensorboard"					,					project_dir=snake_case__   )
															accelerator.init_trackers(snake_case__   )
															_UpperCAmelCase									=  {"tensor": torch.tensor(1   )}
															accelerator.log(snake_case__					,					step=0   )
															accelerator.end_training()
															# Logged values are stored in the outermost-tfevents file and can be read in as a TFRecord
															# Names are randomly generated each time
															_UpperCAmelCase									=  list(filter(lambda snake_case__   : x.is_file()					,					Path(F"""{dirpath}/{project_name}"""   ).iterdir()   )   )[0]
															# Reading implementation based on https://github.com/pytorch/pytorch/issues/45327#issuecomment-703757685
															with open(snake_case__					,					"rb"   ) as f:
																					_UpperCAmelCase									=  f.read()
															_UpperCAmelCase									=  False
															while data:
																					_UpperCAmelCase									=  struct.unpack("Q"					,					data[:8]   )

																					_UpperCAmelCase									=  data[12 : 12 + int(header[0]   )]  # 8+4
																					_UpperCAmelCase									=  data[12 + int(header[0]   ) + 4 :]
																					_UpperCAmelCase									=  event_pba.Event()

																					event.ParseFromString(snake_case__   )
																					if event.HasField("summary"   ):
																											for value in event.summary.value:
																																	if value.simple_value == 1.0 and value.tag == "tensor":
																																							_UpperCAmelCase									=  True
															self.assertTrue(snake_case__					,					"Converted tensor was not found in the log file!"   )







			def UpperCamelCase    ( self					:    Optional[int]   ):




									"""simple docstring"""


									with self.assertRaisesRegex(snake_case__					,					"Logging with `tensorboard` requires a `logging_dir`"   ):
															_UpperCAmelCase									=  Accelerator(log_with="tensorboard"   )
									with tempfile.TemporaryDirectory() as dirpath:
															_UpperCAmelCase									=  Accelerator(log_with="tensorboard"					,					project_dir=snake_case__   )







			def UpperCamelCase    ( self					:    Any   ):




									"""simple docstring"""


									_UpperCAmelCase									=  ProjectConfiguration(total_limit=30   )
									with tempfile.TemporaryDirectory() as dirpath:
															_UpperCAmelCase									=  Accelerator(log_with="tensorboard"					,					project_dir=snake_case__					,					project_config=snake_case__   )






@require_wandb
@mock.patch.dict(os.environ      ,	{"WANDB_MODE": "offline"}				)
class 					__lowerCAmelCase				( UpperCAmelCase__      ,	UpperCAmelCase__				):







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




									"""simple docstring"""


									super().setUp()
									# wandb let's us override where logs are stored to via the WANDB_DIR env var
									self.add_mocks(mock.patch.dict(os.environ					,					{"WANDB_DIR": self.tmpdir}   )   )







			@staticmethod
			def UpperCamelCase    ( snake_case__					:    str					,					snake_case__					:    str					,					snake_case__					:    bool = True   ):




									"""simple docstring"""


									_UpperCAmelCase									=  RF"""{section} ([\S\s]*?)\n\n"""
									if record:
															_UpperCAmelCase									=  RF"""Record: {pattern}"""
									_UpperCAmelCase									=  re.findall(snake_case__					,					snake_case__   )[0]
									# A config
									if section == "config" or section == "history":
															_UpperCAmelCase									=  re.findall(R"\"([a-zA-Z0-9_.,]+)"					,					snake_case__   )
															return {key: val for key, val in zip(cleaned_record[0::2]					,					cleaned_record[1::2]   )}
									# Everything else
									else:
															return dict(re.findall(R"(\w+): \"([^\s]+)\""					,					snake_case__   )   )







			@skip
			def UpperCamelCase    ( self					:    str   ):




									"""simple docstring"""


									_UpperCAmelCase									=  "test_project_with_config"
									_UpperCAmelCase									=  Accelerator(log_with="wandb"   )
									_UpperCAmelCase									=  {"num_iterations": 12, "learning_rate": 1e-2, "some_boolean": False, "some_string": "some_value"}
									_UpperCAmelCase									=  {"wandb": {"tags": ["my_tag"]}}
									accelerator.init_trackers(snake_case__					,					snake_case__					,					snake_case__   )
									_UpperCAmelCase									=  {"total_loss": 0.1, "iteration": 1, "my_text": "some_value"}
									accelerator.log(snake_case__					,					step=0   )
									accelerator.end_training()
									# The latest offline log is stored at wandb/latest-run/*.wandb
									for child in Path(F"""{self.tmpdir}/wandb/latest-run"""   ).glob("*"   ):
															if child.is_file() and child.suffix == ".wandb":
																					_UpperCAmelCase									=  subprocess.check_output(
																					    ["wandb", "sync", "--view", "--verbose", str(snake_case__   )]					,					env=os.environ.copy()   ).decode("utf8"					,					"ignore"   )
																					break

        # Check HPS through careful parsing and cleaning
									_UpperCAmelCase									=  self.parse_log(snake_case__					,					"config"   )
									self.assertEqual(logged_items["num_iterations"]					,					"12"   )
									self.assertEqual(logged_items["learning_rate"]					,					"0.01"   )
									self.assertEqual(logged_items["some_boolean"]					,					"false"   )
									self.assertEqual(logged_items["some_string"]					,					"some_value"   )
									self.assertEqual(logged_items["some_string"]					,					"some_value"   )

									# Run tags
									_UpperCAmelCase									=  self.parse_log(snake_case__					,					"run"					,					snake_case__   )
									self.assertEqual(logged_items["tags"]					,					"my_tag"   )

									# Actual logging
									_UpperCAmelCase									=  self.parse_log(snake_case__					,					"history"   )
									self.assertEqual(logged_items["total_loss"]					,					"0.1"   )
									self.assertEqual(logged_items["iteration"]					,					"1"   )
									self.assertEqual(logged_items["my_text"]					,					"some_value"   )
									self.assertEqual(logged_items["_step"]					,					"0"   )




def 							__SCREAMING_SNAKE_CASE						(    self			,					snake_case_			,					snake_case_		):

						'''simple docstring'''


						_UpperCAmelCase									=  run_name
						_UpperCAmelCase									=  OfflineExperiment(project_name=snake_case_			,					offline_directory=snake_case_		)
						logger.info(f"""Initialized offline CometML project {self.run_name}"""		)
						logger.info("Make sure to log any initial configurations with `self.store_init_configuration` before training!"		)






@require_comet_ml
@mock.patch.object(UpperCAmelCase__      ,	"__init__"      ,	UpperCAmelCase__				)
class 					__lowerCAmelCase				( unittest.TestCase				):







			@staticmethod
			def UpperCamelCase    ( snake_case__					:    Dict					,					snake_case__					:    str					,					snake_case__					:    bool = False   ):




									"""simple docstring"""


									for log in log_list:
															_UpperCAmelCase									=  json.loads(snake_case__   )["payload"]
															if is_param and "param" in j.keys():
																					if j["param"]["paramName"] == key:
																											return j["param"]["paramValue"]
															if "log_other" in j.keys():
																					if j["log_other"]["key"] == key:
																											return j["log_other"]["val"]
															if "metric" in j.keys():
																					if j["metric"]["metricName"] == key:
																											return j["metric"]["metricValue"]







			def UpperCamelCase    ( self					:    Any   ):




									"""simple docstring"""


									with tempfile.TemporaryDirectory() as d:
															_UpperCAmelCase									=  CometMLTracker("test_project_with_config"					,					snake_case__   )
															_UpperCAmelCase									=  Accelerator(log_with=snake_case__   )
															_UpperCAmelCase									=  {"num_iterations": 12, "learning_rate": 1e-2, "some_boolean": False, "some_string": "some_value"}
															accelerator.init_trackers(snake_case__					,					snake_case__   )
															accelerator.end_training()
															_UpperCAmelCase									=  os.listdir(snake_case__   )[0]  # Comet is nice, it's just a zip file here
															# We parse the raw logs
															_UpperCAmelCase									=  os.path.join(snake_case__					,					snake_case__   )
															_UpperCAmelCase									=  zipfile.ZipFile(snake_case__					,					"r"   )
															_UpperCAmelCase									=  archive.open("messages.json"   ).read().decode("utf-8"   )
									_UpperCAmelCase									=  log.split("\n"   )[:-1]
									self.assertEqual(self.get_value_from_key(snake_case__					,					"num_iterations"					,					snake_case__   )					,					12   )
									self.assertEqual(self.get_value_from_key(snake_case__					,					"learning_rate"					,					snake_case__   )					,					0.01   )
									self.assertEqual(self.get_value_from_key(snake_case__					,					"some_boolean"					,					snake_case__   )					,					snake_case__   )
									self.assertEqual(self.get_value_from_key(snake_case__					,					"some_string"					,					snake_case__   )					,					"some_value"   )







			def UpperCamelCase    ( self					:    List[str]   ):




									"""simple docstring"""


									with tempfile.TemporaryDirectory() as d:
															_UpperCAmelCase									=  CometMLTracker("test_project_with_config"					,					snake_case__   )
															_UpperCAmelCase									=  Accelerator(log_with=snake_case__   )
															accelerator.init_trackers(snake_case__   )
															_UpperCAmelCase									=  {"total_loss": 0.1, "iteration": 1, "my_text": "some_value"}
															accelerator.log(snake_case__					,					step=0   )
															accelerator.end_training()
															_UpperCAmelCase									=  os.listdir(snake_case__   )[0]  # Comet is nice, it's just a zip file here
															# We parse the raw logs
															_UpperCAmelCase									=  os.path.join(snake_case__					,					snake_case__   )
															_UpperCAmelCase									=  zipfile.ZipFile(snake_case__					,					"r"   )
															_UpperCAmelCase									=  archive.open("messages.json"   ).read().decode("utf-8"   )
									_UpperCAmelCase									=  log.split("\n"   )[:-1]
									self.assertEqual(self.get_value_from_key(snake_case__					,					"curr_step"					,					snake_case__   )					,					0   )
									self.assertEqual(self.get_value_from_key(snake_case__					,					"total_loss"   )					,					0.1   )
									self.assertEqual(self.get_value_from_key(snake_case__					,					"iteration"   )					,					1   )
									self.assertEqual(self.get_value_from_key(snake_case__					,					"my_text"   )					,					"some_value"   )






class 					__lowerCAmelCase				( UpperCAmelCase__				):
			snake_case_   :				Optional[Any]				  =  [
			    "total_loss",
			    "iteration",
			    "my_text",
			    "learning_rate",
			    "num_iterations",
			    "some_boolean",
			    "some_string",
			]

			snake_case_   :				Optional[Any]				  =  "my_custom_tracker"
			snake_case_   :				int				  =  False







			def __init__( self					:    Dict					,					snake_case__					:    str   ):




									"""simple docstring"""


									_UpperCAmelCase									=  open(F"""{dir}/log.csv"""					,					"w+"   )
									_UpperCAmelCase									=  csv.DictWriter(self.f					,					fieldnames=self._col_names   )
									self.writer.writeheader()







			@property
			def UpperCamelCase    ( self					:    Any   ):




									"""simple docstring"""


									return self.writer







			def UpperCamelCase    ( self					:    Any					,					snake_case__					:    dict   ):




									"""simple docstring"""


									logger.info("Call init"   )
									self.writer.writerow(snake_case__   )







			def UpperCamelCase    ( self					:    Optional[int]					,					snake_case__					:    dict					,					snake_case__					:    Optional[int]   ):




									"""simple docstring"""


									logger.info("Call log"   )
									self.writer.writerow(snake_case__   )







			def UpperCamelCase    ( self					:    List[Any]   ):




									"""simple docstring"""


									self.f.close()






class 					__lowerCAmelCase				( unittest.TestCase				):







			def UpperCamelCase    ( self					:    str   ):




									"""simple docstring"""


									with tempfile.TemporaryDirectory() as d:
															_UpperCAmelCase									=  MyCustomTracker(snake_case__   )
															_UpperCAmelCase									=  Accelerator(log_with=snake_case__   )
															_UpperCAmelCase									=  {"num_iterations": 12, "learning_rate": 1e-2, "some_boolean": False, "some_string": "some_value"}
															accelerator.init_trackers("Some name"					,					snake_case__   )
															accelerator.end_training()
															with open(F"""{d}/log.csv"""					,					"r"   ) as f:
																					_UpperCAmelCase									=  csv.DictReader(snake_case__   )
																					_UpperCAmelCase									=  next(snake_case__   )
																					_UpperCAmelCase									=  {
																					    "total_loss": "",
																					    "iteration": "",
																					    "my_text": "",
																					    "learning_rate": "0.01",
																					    "num_iterations": "12",
																					    "some_boolean": "False",
																					    "some_string": "some_value",
																					}
																					self.assertDictEqual(snake_case__					,					snake_case__   )







			def UpperCamelCase    ( self					:    Optional[int]   ):




									"""simple docstring"""


									with tempfile.TemporaryDirectory() as d:
															_UpperCAmelCase									=  MyCustomTracker(snake_case__   )
															_UpperCAmelCase									=  Accelerator(log_with=snake_case__   )
															accelerator.init_trackers("Some name"   )
															_UpperCAmelCase									=  {"total_loss": 0.1, "iteration": 1, "my_text": "some_value"}
															accelerator.log(snake_case__					,					step=0   )
															accelerator.end_training()
															with open(F"""{d}/log.csv"""					,					"r"   ) as f:
																					_UpperCAmelCase									=  csv.DictReader(snake_case__   )
																					_UpperCAmelCase									=  next(snake_case__   )
																					_UpperCAmelCase									=  {
																					    "total_loss": "0.1",
																					    "iteration": "1",
																					    "my_text": "some_value",
																					    "learning_rate": "",
																					    "num_iterations": "",
																					    "some_boolean": "",
																					    "some_string": "",
																					}
																					self.assertDictEqual(snake_case__					,					snake_case__   )










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:
				lowercase_			:       str						       =				False






class 					__lowerCAmelCase				:







			def __init__( self					:    Optional[Any]					,					snake_case__					:    Union[str, Any]					,					snake_case__					:    Union[str, Any]=14					,					snake_case__					:    Optional[Any]=7					,					snake_case__					:    int=True					,					snake_case__					:    Optional[int]=False					,					snake_case__					:    Any=True					,					snake_case__					:    Tuple=True					,					snake_case__					:    Union[str, Any]=True					,					snake_case__					:    Optional[int]=99					,					snake_case__					:    List[str]=32					,					snake_case__					:    Union[str, Any]=5					,					snake_case__					:    Dict=37					,					snake_case__					:    int="gelu"					,					snake_case__					:    List[str]=0.1					,					snake_case__					:    Dict=0.1					,					snake_case__					:    str=512					,					snake_case__					:    Any=16					,					snake_case__					:    Dict=2					,					snake_case__					:    int=3					,					snake_case__					:    str=4					,					snake_case__					:    int=None					,					):




									"""simple docstring"""


									_UpperCAmelCase									=  parent
									_UpperCAmelCase									=  batch_size
									_UpperCAmelCase									=  seq_length
									_UpperCAmelCase									=  is_training
									_UpperCAmelCase									=  use_token_type_ids
									_UpperCAmelCase									=  use_input_mask
									_UpperCAmelCase									=  use_labels
									_UpperCAmelCase									=  use_mc_token_ids
									_UpperCAmelCase									=  vocab_size
									_UpperCAmelCase									=  hidden_size
									_UpperCAmelCase									=  num_hidden_layers
									_UpperCAmelCase									=  intermediate_size
									_UpperCAmelCase									=  hidden_act
									_UpperCAmelCase									=  hidden_dropout_prob
									_UpperCAmelCase									=  attention_probs_dropout_prob
									_UpperCAmelCase									=  max_position_embeddings
									_UpperCAmelCase									=  type_vocab_size
									_UpperCAmelCase									=  type_sequence_label_size
									_UpperCAmelCase									=  num_labels
									_UpperCAmelCase									=  num_choices
									_UpperCAmelCase									=  scope
									_UpperCAmelCase									=  vocab_size - 1
									_UpperCAmelCase									=  vocab_size - 1
									_UpperCAmelCase									=  vocab_size - 1







			def UpperCamelCase    ( self					:    Any   ):




									"""simple docstring"""


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







			def UpperCamelCase    ( self					:    List[Any]					,					snake_case__					:    Dict=False					,					snake_case__					:    Union[str, Any]=False					,					snake_case__					:    Any=False   ):




									"""simple docstring"""


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

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

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

									_UpperCAmelCase									=  None
									if self.use_mc_token_ids:
															_UpperCAmelCase									=  ids_tensor([self.batch_size, self.num_choices]					,					self.seq_length   )

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

									_UpperCAmelCase									=  self.get_config(
									    gradient_checkpointing=snake_case__					,					scale_attn_by_inverse_layer_idx=snake_case__					,					reorder_and_upcast_attn=snake_case__					,					)

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







			def UpperCamelCase    ( self					:    Optional[Any]					,					snake_case__					:    Dict=False					,					snake_case__					:    Dict=False					,					snake_case__					:    Tuple=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=snake_case__					,					bos_token_id=self.bos_token_id					,					eos_token_id=self.eos_token_id					,					pad_token_id=self.pad_token_id					,					gradient_checkpointing=snake_case__					,					scale_attn_by_inverse_layer_idx=snake_case__					,					reorder_and_upcast_attn=snake_case__					,					)







			def UpperCamelCase    ( self					:    List[Any]   ):




									"""simple docstring"""


									_UpperCAmelCase									=  self.get_config()
									_UpperCAmelCase									=  300
									return config







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




									"""simple docstring"""


									(
									    (
									    _UpperCAmelCase
									)							,				(
									    _UpperCAmelCase
									)							,				(
									    _UpperCAmelCase
									)							,				(
									    _UpperCAmelCase
									)							,				(
									    _UpperCAmelCase
									)							,				(
									    _UpperCAmelCase
									)							,				(
									    _UpperCAmelCase
									)							,				(
									    _UpperCAmelCase
									)							,				(
									    _UpperCAmelCase
									)							,				
									)									=  self.prepare_config_and_inputs()

									_UpperCAmelCase									=  floats_tensor([self.batch_size, self.seq_length, self.hidden_size]   )
									_UpperCAmelCase									=  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 UpperCamelCase    ( self					:    Optional[Any]					,					snake_case__					:    List[str]					,					snake_case__					:    Optional[int]					,					snake_case__					:    int					,					snake_case__					:    Optional[Any]					,					snake_case__					:    str					,					*snake_case__					:    Optional[int]   ):




									"""simple docstring"""


									_UpperCAmelCase									=  True
									_UpperCAmelCase									=  RwkvModel(config=snake_case__   )
									model.to(snake_case__   )
									model.eval()

									_UpperCAmelCase									=  model(snake_case__   )

									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 UpperCamelCase    ( self					:    List[str]					,					snake_case__					:    str					,					snake_case__					:    Any					,					snake_case__					:    Any					,					snake_case__					:    Optional[Any]					,					snake_case__					:    List[Any]					,					*snake_case__					:    int   ):




									"""simple docstring"""


									_UpperCAmelCase									=  RwkvForCausalLM(snake_case__   )
									model.to(snake_case__   )
									model.eval()

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







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




									"""simple docstring"""


									_UpperCAmelCase									=  RwkvModel(config=snake_case__   )
									model.to(snake_case__   )
									model.eval()

									_UpperCAmelCase									=  model(snake_case__   )
									_UpperCAmelCase									=  outputs.last_hidden_state

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

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

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







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




									"""simple docstring"""


									_UpperCAmelCase									=  RwkvForCausalLM(snake_case__   )
									model.to(snake_case__   )
									if gradient_checkpointing:
															model.gradient_checkpointing_enable()

									_UpperCAmelCase									=  model(snake_case__					,					labels=snake_case__   )
									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 UpperCamelCase    ( self					:    Dict   ):




									"""simple docstring"""


									_UpperCAmelCase									=  self.prepare_config_and_inputs()

									(
									    (
									    _UpperCAmelCase
									)							,				(
									    _UpperCAmelCase
									)							,				(
									    _UpperCAmelCase
									)							,				(
									    _UpperCAmelCase
									)							,				(
									    _UpperCAmelCase
									)							,				(
									    _UpperCAmelCase
									)							,				(
									    _UpperCAmelCase
									)							,				(
									    _UpperCAmelCase
									)							,				(
									    _UpperCAmelCase
									)							,				
									)									=  config_and_inputs

									_UpperCAmelCase									=  {"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				( UpperCAmelCase__      ,	UpperCAmelCase__      ,	UpperCAmelCase__      ,	unittest.TestCase				):
			snake_case_   :				List[str]				  =  (RwkvModel, RwkvForCausalLM) if is_torch_available() else ()
			snake_case_   :				int				  =  (
			    {"feature-extraction": RwkvModel, "text-generation": RwkvForCausalLM} if is_torch_available() else {}
			)
			# all_generative_model_classes = (RwkvForCausalLM,) if is_torch_available() else ()
			snake_case_   :				Tuple				  =  False
			snake_case_   :				Dict				  =  False
			snake_case_   :				int				  =  False
			snake_case_   :				Dict				  =  False
			snake_case_   :				Any				  =  False  # Rwkv does not support head masking







			def UpperCamelCase    ( self					:    Optional[Any]   ):




									"""simple docstring"""


									_UpperCAmelCase									=  RwkvModelTester(self   )
									_UpperCAmelCase									=  ConfigTester(
									    self					,					config_class=snake_case__					,					n_embd=37					,					common_properties=["hidden_size", "num_hidden_layers"]   )







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




									"""simple docstring"""


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

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

									_UpperCAmelCase							,				_UpperCAmelCase									=  container

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

									if not is_inside_interval:
															_UpperCAmelCase									=  "%s not found in %s" % (safe_repr(snake_case__   ), safe_repr(snake_case__   ))
															self.fail(self._formatMessage(snake_case__					,					snake_case__   )   )







			def UpperCamelCase    ( self					:    List[Any]   ):




									"""simple docstring"""


									self.config_tester.run_common_tests()







			def UpperCamelCase    ( self					:    Tuple   ):




									"""simple docstring"""


									_UpperCAmelCase									=  self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_rwkv_model(*snake_case__   )







			def UpperCamelCase    ( self					:    Dict   ):




									"""simple docstring"""


									_UpperCAmelCase									=  self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_causl_lm(*snake_case__   )







			def UpperCamelCase    ( self					:    Any   ):




									"""simple docstring"""


									_UpperCAmelCase									=  self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_state_equivalency(*snake_case__   )







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




									"""simple docstring"""


									_UpperCAmelCase							,				_UpperCAmelCase									=  self.model_tester.prepare_config_and_inputs_for_common()

									for model_class in self.all_model_classes:
															_UpperCAmelCase									=  model_class(config=snake_case__   )
															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 UpperCamelCase    ( self					:    int   ):




									"""simple docstring"""


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

									_UpperCAmelCase									=  getattr(self.model_tester					,					"seq_length"					,					snake_case__   )

									for model_class in self.all_model_classes:
															_UpperCAmelCase									=  True
															_UpperCAmelCase									=  False
															_UpperCAmelCase									=  True
															_UpperCAmelCase									=  model_class(snake_case__   )
															model.to(snake_case__   )
															model.eval()

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

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

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

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

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

															_UpperCAmelCase									=  self._prepare_for_class(snake_case__					,					snake_case__   )
															_UpperCAmelCase									=  inputs["input_ids"].shape[0]
															with torch.no_grad():
																					_UpperCAmelCase									=  model(**snake_case__   )

															_UpperCAmelCase									=  1
															self.assertEqual(out_len + added_hidden_states					,					len(snake_case__   )   )

															_UpperCAmelCase									=  outputs.attentions

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







			@slow
			def UpperCamelCase    ( self					:    int   ):




									"""simple docstring"""


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






@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				):







			def UpperCamelCase    ( self					:    List[str]   ):




									"""simple docstring"""


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







			def UpperCamelCase    ( self					:    Dict   ):




									"""simple docstring"""


									_UpperCAmelCase									=  "Hello my name is Jasmine and I am a newbie to the"
									_UpperCAmelCase									=  RwkvForCausalLM.from_pretrained(self.model_id   ).to(snake_case__   )

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

									self.assertEqual(snake_case__					,					snake_case__   )







			def UpperCamelCase    ( self					:    str   ):




									"""simple docstring"""


									_UpperCAmelCase									=  "Hello my name is Jasmine and I am a newbie to the"

									_UpperCAmelCase									=  self.tokenizer("Hello my name is"					,					return_tensors="pt"   ).input_ids.to(snake_case__   )
									_UpperCAmelCase									=  RwkvForCausalLM.from_pretrained(self.model_id					,					torch_dtype=torch.bfloataa   ).to(snake_case__   )

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

									self.assertEqual(snake_case__					,					snake_case__   )




