


'''simple docstring'''


from __future__ import annotations

import unittest

from transformers import is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
						import numpy
						import tensorflow as tf

						from transformers import (
						    TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
						    TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
						    TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
						    BertConfig,
						    DPRConfig,
						    TFDPRContextEncoder,
						    TFDPRQuestionEncoder,
						    TFDPRReader,
						)






class 	lowerCAmelCase__			:




			def __init__(		self  :     int	,  lowerCamelCase__  :     Optional[int]	,  lowerCamelCase__  :     List[Any]=13	,  lowerCamelCase__  :     int=7	,  lowerCamelCase__  :     List[Any]=True	,  lowerCamelCase__  :     Dict=True	,  lowerCamelCase__  :     List[Any]=True	,  lowerCamelCase__  :     int=True	,  lowerCamelCase__  :     Optional[int]=99	,  lowerCamelCase__  :     Optional[Any]=32	,  lowerCamelCase__  :     Optional[int]=2	,  lowerCamelCase__  :     Any=4	,  lowerCamelCase__  :     int=37	,  lowerCamelCase__  :     Dict="gelu"	,  lowerCamelCase__  :     Optional[Any]=0.1	,  lowerCamelCase__  :     Optional[int]=0.1	,  lowerCamelCase__  :     int=5_12	,  lowerCamelCase__  :     List[Any]=16	,  lowerCamelCase__  :     Optional[int]=2	,  lowerCamelCase__  :     List[str]=0.0_2	,  lowerCamelCase__  :     Dict=3	,  lowerCamelCase__  :     str=4	,  lowerCamelCase__  :     str=None	,  lowerCamelCase__  :     List[Any]=0	,  )   ->Union[str, Any]:



					'''simple docstring'''







					_UpperCAmelCase		:       Dict								=      parent
					_UpperCAmelCase		:       Union[str, Any]								=      batch_size
					_UpperCAmelCase		:       Any								=      seq_length
					_UpperCAmelCase		:       int								=      is_training
					_UpperCAmelCase		:       Any								=      use_input_mask
					_UpperCAmelCase		:       Tuple								=      use_token_type_ids
					_UpperCAmelCase		:       int								=      use_labels
					_UpperCAmelCase		:       Tuple								=      vocab_size
					_UpperCAmelCase		:       Tuple								=      hidden_size
					_UpperCAmelCase		:       Any								=      num_hidden_layers
					_UpperCAmelCase		:       Tuple								=      num_attention_heads
					_UpperCAmelCase		:       int								=      intermediate_size
					_UpperCAmelCase		:       Tuple								=      hidden_act
					_UpperCAmelCase		:       int								=      hidden_dropout_prob
					_UpperCAmelCase		:       Any								=      attention_probs_dropout_prob
					_UpperCAmelCase		:       Optional[int]								=      max_position_embeddings
					_UpperCAmelCase		:       Optional[int]								=      type_vocab_size
					_UpperCAmelCase		:       List[Any]								=      type_sequence_label_size
					_UpperCAmelCase		:       Union[str, Any]								=      initializer_range
					_UpperCAmelCase		:       str								=      num_labels
					_UpperCAmelCase		:       Optional[int]								=      num_choices
					_UpperCAmelCase		:       List[str]								=      scope
					_UpperCAmelCase		:       int								=      projection_dim




			def 							lowerCAmelCase__    (		self  :     Union[str, Any] )   ->Tuple:



					'''simple docstring'''







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

					_UpperCAmelCase		:       Optional[Any]								=      None
					if self.use_input_mask:
							# follow test_modeling_tf_ctrl.py
							_UpperCAmelCase		:       Union[str, Any]								=      random_attention_mask([self.batch_size, self.seq_length] )

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

					_UpperCAmelCase		:       Optional[Any]								=      None
					_UpperCAmelCase		:       str								=      None
					_UpperCAmelCase		:       Tuple								=      None
					if self.use_labels:
							_UpperCAmelCase		:       Dict								=      ids_tensor([self.batch_size]	,  self.type_sequence_label_size )
							_UpperCAmelCase		:       int								=      ids_tensor([self.batch_size, self.seq_length]	,  self.num_labels )
							_UpperCAmelCase		:       Tuple								=      ids_tensor([self.batch_size]	,  self.num_choices )

					_UpperCAmelCase		:       Optional[int]								=      BertConfig(
					    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	,  )
					_UpperCAmelCase		:       List[Any]								=      DPRConfig(projection_dim=self.projection_dim	,  **config.to_dict() )

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




			def 							lowerCAmelCase__    (		self  :     Optional[Any]	,  lowerCamelCase__  :     Tuple	,  lowerCamelCase__  :     Dict	,  lowerCamelCase__  :     Optional[int]	,  lowerCamelCase__  :     Optional[int]	,  lowerCamelCase__  :     Tuple	,  lowerCamelCase__  :     Any	,  lowerCamelCase__  :     int )   ->int:



					'''simple docstring'''







					_UpperCAmelCase		:       Union[str, Any]								=      TFDPRContextEncoder(config=lowerCamelCase__ )
					_UpperCAmelCase		:       Dict								=      model(lowerCamelCase__	,  attention_mask=lowerCamelCase__	,  token_type_ids=lowerCamelCase__ )
					_UpperCAmelCase		:       Optional[Any]								=      model(lowerCamelCase__	,  token_type_ids=lowerCamelCase__ )
					_UpperCAmelCase		:       Any								=      model(lowerCamelCase__ )
					self.parent.assertEqual(result.pooler_output.shape	,  (self.batch_size, self.projection_dim or self.hidden_size) )




			def 							lowerCAmelCase__    (		self  :     List[str]	,  lowerCamelCase__  :     List[Any]	,  lowerCamelCase__  :     Union[str, Any]	,  lowerCamelCase__  :     str	,  lowerCamelCase__  :     Optional[Any]	,  lowerCamelCase__  :     Optional[Any]	,  lowerCamelCase__  :     Optional[int]	,  lowerCamelCase__  :     Any )   ->Any:



					'''simple docstring'''







					_UpperCAmelCase		:       Any								=      TFDPRQuestionEncoder(config=lowerCamelCase__ )
					_UpperCAmelCase		:       Optional[int]								=      model(lowerCamelCase__	,  attention_mask=lowerCamelCase__	,  token_type_ids=lowerCamelCase__ )
					_UpperCAmelCase		:       int								=      model(lowerCamelCase__	,  token_type_ids=lowerCamelCase__ )
					_UpperCAmelCase		:       Tuple								=      model(lowerCamelCase__ )
					self.parent.assertEqual(result.pooler_output.shape	,  (self.batch_size, self.projection_dim or self.hidden_size) )




			def 							lowerCAmelCase__    (		self  :     Optional[Any]	,  lowerCamelCase__  :     List[Any]	,  lowerCamelCase__  :     List[str]	,  lowerCamelCase__  :     Optional[Any]	,  lowerCamelCase__  :     Optional[Any]	,  lowerCamelCase__  :     Optional[int]	,  lowerCamelCase__  :     Optional[Any]	,  lowerCamelCase__  :     List[str] )   ->str:



					'''simple docstring'''







					_UpperCAmelCase		:       List[str]								=      TFDPRReader(config=lowerCamelCase__ )
					_UpperCAmelCase		:       Optional[int]								=      model(lowerCamelCase__	,  attention_mask=lowerCamelCase__ )

					self.parent.assertEqual(result.start_logits.shape	,  (self.batch_size, self.seq_length) )
					self.parent.assertEqual(result.end_logits.shape	,  (self.batch_size, self.seq_length) )
					self.parent.assertEqual(result.relevance_logits.shape	,  (self.batch_size,) )






			def 							lowerCAmelCase__    (		self  :     Optional[Any] )   ->str:



					'''simple docstring'''







					_UpperCAmelCase		:       Tuple								=      self.prepare_config_and_inputs()
					(
					    (
					    _UpperCAmelCase
					)					,    (
					    _UpperCAmelCase
					)					,    (
					    _UpperCAmelCase
					)					,    (
					    _UpperCAmelCase
					)					,    (
					    _UpperCAmelCase
					)					,    (
					    _UpperCAmelCase
					)					,    (
					    _UpperCAmelCase
					)					,    
					)		:       Optional[Any]								=      config_and_inputs
					_UpperCAmelCase		:       Tuple								=      {"input_ids": input_ids}
					return config, inputs_dict






@require_tf
class 	lowerCAmelCase__			(			UpperCAmelCase__    ,       UpperCAmelCase__    ,       unittest.TestCase	):
			lowerCAmelCase :	Optional[int]          =     (
			    (
			        TFDPRContextEncoder,
			        TFDPRQuestionEncoder,
			        TFDPRReader,
			    )
			    if is_tf_available()
			    else ()
			)
			lowerCAmelCase :	List[str]          =     {"feature-extraction": TFDPRQuestionEncoder} if is_tf_available() else {}

			lowerCAmelCase :	int          =     False
			lowerCAmelCase :	Optional[Any]          =     False
			lowerCAmelCase :	Union[str, Any]          =     False
			lowerCAmelCase :	Tuple          =     False
			lowerCAmelCase :	int          =     False




			def 							lowerCAmelCase__    (		self  :     Dict )   ->str:



					'''simple docstring'''







					_UpperCAmelCase		:       Optional[int]								=      TFDPRModelTester(self )
					_UpperCAmelCase		:       Tuple								=      ConfigTester(self	,  config_class=lowerCamelCase__	,  hidden_size=37 )




			def 							lowerCAmelCase__    (		self  :     List[str] )   ->List[Any]:



					'''simple docstring'''







					self.config_tester.run_common_tests()




			def 							lowerCAmelCase__    (		self  :     int )   ->Dict:



					'''simple docstring'''







					_UpperCAmelCase		:       Union[str, Any]								=      self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_dpr_context_encoder(*lowerCamelCase__ )




			def 							lowerCAmelCase__    (		self  :     List[Any] )   ->Union[str, Any]:



					'''simple docstring'''







					_UpperCAmelCase		:       Dict								=      self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_dpr_question_encoder(*lowerCamelCase__ )




			def 							lowerCAmelCase__    (		self  :     Dict )   ->Dict:



					'''simple docstring'''







					_UpperCAmelCase		:       Dict								=      self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_dpr_reader(*lowerCamelCase__ )






			@slow
			def 							lowerCAmelCase__    (		self  :     List[str] )   ->Dict:



					'''simple docstring'''







					for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
							_UpperCAmelCase		:       Dict								=      TFDPRContextEncoder.from_pretrained(lowerCamelCase__ )
							self.assertIsNotNone(lowerCamelCase__ )

					for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
							_UpperCAmelCase		:       int								=      TFDPRContextEncoder.from_pretrained(lowerCamelCase__ )
							self.assertIsNotNone(lowerCamelCase__ )

					for model_name in TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
							_UpperCAmelCase		:       str								=      TFDPRQuestionEncoder.from_pretrained(lowerCamelCase__ )
							self.assertIsNotNone(lowerCamelCase__ )

					for model_name in TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
							_UpperCAmelCase		:       List[Any]								=      TFDPRReader.from_pretrained(lowerCamelCase__ )
							self.assertIsNotNone(lowerCamelCase__ )






@require_tf
class 	lowerCAmelCase__			(			unittest.TestCase	):




			@slow
			def 							lowerCAmelCase__    (		self  :     int )   ->str:



					'''simple docstring'''







					_UpperCAmelCase		:       int								=      TFDPRQuestionEncoder.from_pretrained("facebook/dpr-question_encoder-single-nq-base" )

					_UpperCAmelCase		:       Tuple								=      tf.constant(
					    [[1_01, 75_92, 10_10, 20_03, 20_26, 38_99, 1_01_40, 10_29, 1_02]] )  # [CLS] hello, is my dog cute? [SEP]
					_UpperCAmelCase		:       int								=      model(lowerCamelCase__ )[0]  # embedding shape = (1, 768)
					# compare the actual values for a slice.
					_UpperCAmelCase		:       Dict								=      tf.constant(
					    [
					        [
					            0.0_3_2_3_6_2_5_3,
					            0.1_2_7_5_3_3_3_5,
					            0.1_6_8_1_8_5_0_9,
					            0.0_0_2_7_9_7_8_6,
					            0.3_8_9_6_9_3_3,
					            0.2_4_2_6_4_9_4_5,
					            0.2_1_7_8_9_7_1,
					            -0.0_2_3_3_5_2_2_7,
					            -0.0_8_4_8_1_9_5_9,
					            -0.1_4_3_2_4_1_1_7,
					        ]
					    ] )
					self.assertTrue(numpy.allclose(output[:, :10].numpy()	,  expected_slice.numpy()	,  atol=1E-4 ) )







'''simple docstring'''


import argparse
import gc
import json
import os
import re

import torch
from huggingface_hub import hf_hub_download

from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedTokenizerFast, RwkvConfig
from transformers.modeling_utils import WEIGHTS_INDEX_NAME, shard_checkpoint


lowerCamelCase__    	=     {
    '169M': 12,
    '430M': 24,
    '1B5': 24,
    '3B': 32,
    '7B': 32,
    '14B': 40,
}

lowerCamelCase__    	=     {
    '169M': 768,
    '430M': 1_024,
    '1B5': 2_048,
    '3B': 2_560,
    '7B': 4_096,
    '14B': 5_120,
}






def  __lowerCAmelCase (__lowerCAmelCase   ):
		_UpperCAmelCase		:       List[str]								=      list(state_dict.keys()   )
		for name in state_dict_keys:
				_UpperCAmelCase		:       Optional[int]								=      state_dict.pop(__lowerCAmelCase   )
				# emb -> embedding
				if name.startswith("emb."   ):
						_UpperCAmelCase		:       Tuple								=      name.replace("emb."      , "embeddings."   )
				# ln_0 -> pre_ln (only present at block 0)
				if name.startswith("blocks.0.ln0"   ):
						_UpperCAmelCase		:       Optional[int]								=      name.replace("blocks.0.ln0"      , "blocks.0.pre_ln"   )
				# att -> attention
				_UpperCAmelCase		:       Union[str, Any]								=      re.sub(R"blocks\.(\d+)\.att"      , R"blocks.\1.attention"      , __lowerCAmelCase   )
				# ffn -> feed_forward
				_UpperCAmelCase		:       Dict								=      re.sub(R"blocks\.(\d+)\.ffn"      , R"blocks.\1.feed_forward"      , __lowerCAmelCase   )
				# time_mix_k -> time_mix_key and reshape
				if name.endswith(".time_mix_k"   ):
						_UpperCAmelCase		:       int								=      name.replace(".time_mix_k"      , ".time_mix_key"   )
				# time_mix_v -> time_mix_value and reshape
				if name.endswith(".time_mix_v"   ):
						_UpperCAmelCase		:       Union[str, Any]								=      name.replace(".time_mix_v"      , ".time_mix_value"   )
				# time_mix_r -> time_mix_key and reshape
				if name.endswith(".time_mix_r"   ):
						_UpperCAmelCase		:       int								=      name.replace(".time_mix_r"      , ".time_mix_receptance"   )

				if name != "head.weight":
						_UpperCAmelCase		:       List[str]								=      "rwkv." + name

				_UpperCAmelCase		:       Optional[Any]								=      weight
		return state_dict






def  __lowerCAmelCase (__lowerCAmelCase      , __lowerCAmelCase      , __lowerCAmelCase      , __lowerCAmelCase=None      , __lowerCAmelCase=None      , __lowerCAmelCase=False      , __lowerCAmelCase=None   ):
		# 1. If possible, build the tokenizer.
		if tokenizer_file is None:
				print("No `--tokenizer_file` provided, we will use the default tokenizer."   )
				_UpperCAmelCase		:       str								=      50_277
				_UpperCAmelCase		:       Optional[Any]								=      AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b"   )
		else:
				_UpperCAmelCase		:       Tuple								=      PreTrainedTokenizerFast(tokenizer_file=__lowerCAmelCase   )
				_UpperCAmelCase		:       List[Any]								=      len(__lowerCAmelCase   )
		tokenizer.save_pretrained(__lowerCAmelCase   )

		# 2. Build the config
		_UpperCAmelCase		:       Optional[int]								=      list(NUM_HIDDEN_LAYERS_MAPPING.keys()   )
		if size is None:
				# Try to infer size from the checkpoint name
				for candidate in possible_sizes:
						if candidate in checkpoint_file:
								_UpperCAmelCase		:       Optional[Any]								=      candidate
								break
				if size is None:
						raise ValueError("Could not infer the size, please provide it with the `--size` argument."   )
		if size not in possible_sizes:
				raise ValueError(F"""`size` should be one of {possible_sizes}, got {size}."""   )

		_UpperCAmelCase		:       Any								=      RwkvConfig(
		    vocab_size=__lowerCAmelCase      , num_hidden_layers=NUM_HIDDEN_LAYERS_MAPPING[size]      , hidden_size=HIDEN_SIZE_MAPPING[size]      , )
		config.save_pretrained(__lowerCAmelCase   )

		# 3. Download model file then convert state_dict
		_UpperCAmelCase		:       str								=      hf_hub_download(__lowerCAmelCase      , __lowerCAmelCase   )
		_UpperCAmelCase		:       Optional[int]								=      torch.load(__lowerCAmelCase      , map_location="cpu"   )
		_UpperCAmelCase		:       Any								=      convert_state_dict(__lowerCAmelCase   )

		# 4. Split in shards and save
		_UpperCAmelCase					,    _UpperCAmelCase		:       List[str]								=      shard_checkpoint(__lowerCAmelCase   )
		for shard_file, shard in shards.items():
				torch.save(__lowerCAmelCase      , os.path.join(__lowerCAmelCase      , __lowerCAmelCase   )   )

		if index is not None:
				_UpperCAmelCase		:       int								=      os.path.join(__lowerCAmelCase      , __lowerCAmelCase   )
				# Save the index as well
				with open(__lowerCAmelCase      , "w"      , encoding="utf-8"   ) as f:
						_UpperCAmelCase		:       int								=      json.dumps(__lowerCAmelCase      , indent=2      , sort_keys=__lowerCAmelCase   ) + "\n"
						f.write(__lowerCAmelCase   )

				# 5. Clean up shards (for some reason the file PyTorch saves take the same space as the whole state_dict
				print(
				    "Cleaning up shards. This may error with an OOM error, it this is the case don't worry you still have converted the model."   )
				_UpperCAmelCase		:       Union[str, Any]								=      list(shards.keys()   )

				del state_dict
				del shards
				gc.collect()

				for shard_file in shard_files:
						_UpperCAmelCase		:       Union[str, Any]								=      torch.load(os.path.join(__lowerCAmelCase      , __lowerCAmelCase   )   )
						torch.save({k: v.cpu().clone() for k, v in state_dict.items()}      , os.path.join(__lowerCAmelCase      , __lowerCAmelCase   )   )

		del state_dict
		gc.collect()

		if push_to_hub:
				if model_name is None:
						raise ValueError("Please provide a `model_name` to push the model to the Hub."   )
				_UpperCAmelCase		:       Optional[Any]								=      AutoModelForCausalLM.from_pretrained(__lowerCAmelCase   )
				model.push_to_hub(__lowerCAmelCase      , max_shard_size="2GB"   )
				tokenizer.push_to_hub(__lowerCAmelCase   )


if __name__ == "__main__":
						lowerCamelCase__    	=     argparse.ArgumentParser()
						# Required parameters
						parser.add_argument(
						    '--repo_id', default=None, type=str, required=True, help='Repo ID from which to pull the checkpoint.'
						)
						parser.add_argument(
						    '--checkpoint_file', default=None, type=str, required=True, help='Name of the checkpoint file in the repo.'
						)
						parser.add_argument(
						    '--output_dir', default=None, type=str, required=True, help='Where to save the converted model.'
						)
						parser.add_argument(
						    '--tokenizer_file',
						    default=None,
						    type=str,
						    help='Path to the tokenizer file to use (if not provided, only the model is converted).',
						)
						parser.add_argument(
						    '--size',
						    default=None,
						    type=str,
						    help='Size of the model. Will be inferred from the `checkpoint_file` if not passed.',
						)
						parser.add_argument(
						    '--push_to_hub',
						    action='store_true',
						    help='Push to the Hub the converted model.',
						)
						parser.add_argument(
						    '--model_name',
						    default=None,
						    type=str,
						    help='Name of the pushed model on the Hub, including the username / organization.',
						)

						lowerCamelCase__    	=     parser.parse_args()
						convert_rmkv_checkpoint_to_hf_format(
						    args.repo_id,
						    args.checkpoint_file,
						    args.output_dir,
						    size=args.size,
						    tokenizer_file=args.tokenizer_file,
						    push_to_hub=args.push_to_hub,
						    model_name=args.model_name,
						)



