





'''simple docstring'''






import unittest

import numpy as np

from transformers import is_flax_available
from transformers.testing_utils import require_flax

from ..test_modeling_flax_common import ids_tensor


if is_flax_available():
					import jax
					import jax.numpy as jnp

					from transformers.generation import (
					    FlaxForcedBOSTokenLogitsProcessor,
					    FlaxForcedEOSTokenLogitsProcessor,
					    FlaxLogitsProcessorList,
					    FlaxMinLengthLogitsProcessor,
					    FlaxTemperatureLogitsWarper,
					    FlaxTopKLogitsWarper,
					    FlaxTopPLogitsWarper,
					)




@require_flax
class     __snake_case		(   unittest.TestCase):
					def     UpperCAmelCase_							(     self,    A,    A      ):

									"""simple docstring"""


									lowerCamelCase					:			str									= jnp.ones((batch_size, length)      ) / length
									return scores
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Union[str, Any]									= None
									lowerCamelCase					:			List[Any]									= 20

									lowerCamelCase					:			str									= self._get_uniform_logits(batch_size=2,    length=_A      )

									# tweak scores to not be uniform anymore
									lowerCamelCase					:			int									= scores.at[1, 5].set((1 / length) + 0.1      )  # peak, 1st batch
									lowerCamelCase					:			int									= scores.at[1, 10].set((1 / length) - 0.4      )  # valley, 1st batch

									# compute softmax
									lowerCamelCase					:			List[str]									= jax.nn.softmax(_A,    axis=-1      )

									lowerCamelCase					:			Union[str, Any]									= FlaxTemperatureLogitsWarper(temperature=0.5      )
									lowerCamelCase					:			Optional[Any]									= FlaxTemperatureLogitsWarper(temperature=1.3      )

									lowerCamelCase					:			Tuple									= jax.nn.softmax(temp_dist_warper_sharper(_A,    scores.copy(),    cur_len=_A      ),    axis=-1      )
									lowerCamelCase					:			str									= jax.nn.softmax(temp_dist_warper_smoother(_A,    scores.copy(),    cur_len=_A      ),    axis=-1      )

									# uniform distribution stays uniform
									self.assertTrue(jnp.allclose(probs[0, :],    warped_prob_sharp[0, :],    atol=1e-3      )      )
									self.assertTrue(jnp.allclose(probs[0, :],    warped_prob_smooth[0, :],    atol=1e-3      )      )

									# sharp peaks get higher, valleys get lower
									self.assertLess(probs[1, :].max(),    warped_prob_sharp[1, :].max()      )
									self.assertGreater(probs[1, :].min(),    warped_prob_sharp[1, :].min()      )

									# smooth peaks get lower, valleys get higher
									self.assertGreater(probs[1, :].max(),    warped_prob_smooth[1, :].max()      )
									self.assertLess(probs[1, :].min(),    warped_prob_smooth[1, :].min()      )
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			List[str]									= None
									lowerCamelCase					:			Dict									= 10
									lowerCamelCase					:			Optional[Any]									= 2

									# create ramp distribution
									lowerCamelCase					:			Union[str, Any]									= np.broadcast_to(np.arange(_A      )[None, :],    (batch_size, vocab_size)      ).copy()
									lowerCamelCase					:			Any									= ramp_logits[1:, : vocab_size // 2] + vocab_size

									lowerCamelCase					:			List[Any]									= FlaxTopKLogitsWarper(3      )

									lowerCamelCase					:			List[Any]									= top_k_warp(_A,    _A,    cur_len=_A      )

									# check that correct tokens are filtered
									self.assertListEqual(jnp.isinf(scores[0]      ).tolist(),    7 * [True] + 3 * [False]      )
									self.assertListEqual(jnp.isinf(scores[1]      ).tolist(),    2 * [True] + 3 * [False] + 5 * [True]      )

									# check special case
									lowerCamelCase					:			List[str]									= 5
									lowerCamelCase					:			int									= FlaxTopKLogitsWarper(top_k=1,    filter_value=0.0,    min_tokens_to_keep=3      )

									lowerCamelCase					:			List[str]									= np.broadcast_to(np.arange(_A      )[None, :],    (batch_size, length)      ).copy()
									lowerCamelCase					:			Optional[int]									= top_k_warp_safety_check(_A,    _A,    cur_len=_A      )

									# min_tokens overwrites k: 3 tokens are kept => 2 tokens are nullified
									self.assertListEqual((scores == 0.0).sum(axis=-1      ).tolist(),    [2, 2]      )
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Tuple									= None
									lowerCamelCase					:			Optional[Any]									= 10
									lowerCamelCase					:			Optional[int]									= 2

									# create distribution and take log (inverse to Softmax as taken in TopPLogitsWarper)
									lowerCamelCase					:			List[Any]									= np.log(np.array([[0.3, 0.1, 0.1, 0.5], [0.15, 0.3, 0.3, 0.25]]      )      )

									lowerCamelCase					:			Any									= FlaxTopPLogitsWarper(0.8      )
									lowerCamelCase					:			Any									= np.exp(top_p_warp(_A,    _A,    cur_len=_A      )      )

									# dist should be filtered to keep min num values so that sum is >= top_p
									# exp (-inf) => 0
									lowerCamelCase					:			Optional[int]									= np.array([[0.3, 0.0, 0.0, 0.5], [0.0, 0.3, 0.3, 0.25]]      )
									self.assertTrue(np.allclose(_A,    _A,    atol=1e-3      )      )

									# check edge cases with negative and extreme logits
									lowerCamelCase					:			str									= np.broadcast_to(np.arange(_A      )[None, :],    (batch_size, vocab_size)      ).copy() - (
									    vocab_size // 2
									)

									# make ramp_logits more extreme
									lowerCamelCase					:			Any									= ramp_logits[1] * 100.0

									# make sure at least 2 tokens are kept
									lowerCamelCase					:			Optional[Any]									= FlaxTopPLogitsWarper(0.9,    min_tokens_to_keep=2,    filter_value=0.0      )
									lowerCamelCase					:			int									= top_p_warp(_A,    _A,    cur_len=_A      )

									# first batch should keep three tokens, second batch would keep only 1, but due to `min_tokens_to_keep=2` keeps 2.
									self.assertListEqual((filtered_dist != 0.0).sum(axis=-1      ).tolist(),    [3, 2]      )
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			List[Any]									= 20
									lowerCamelCase					:			List[Any]									= 4
									lowerCamelCase					:			Dict									= 0

									lowerCamelCase					:			Union[str, Any]									= FlaxMinLengthLogitsProcessor(min_length=10,    eos_token_id=_A      )

									# check that min length is applied at length 5
									lowerCamelCase					:			Optional[int]									= ids_tensor((batch_size, 20),    vocab_size=20      )
									lowerCamelCase					:			str									= 5
									lowerCamelCase					:			Union[str, Any]									= self._get_uniform_logits(_A,    _A      )
									lowerCamelCase					:			str									= min_dist_processor(_A,    _A,    cur_len=_A      )
									self.assertListEqual(scores_before_min_length[:, eos_token_id].tolist(),    4 * [-float('inf'      )]      )

									# check that min length is not applied anymore at length 15
									lowerCamelCase					:			Tuple									= self._get_uniform_logits(_A,    _A      )
									lowerCamelCase					:			Any									= 15
									lowerCamelCase					:			Dict									= min_dist_processor(_A,    _A,    cur_len=_A      )
									self.assertFalse(jnp.isinf(_A      ).any()      )
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Optional[Any]									= 20
									lowerCamelCase					:			int									= 4
									lowerCamelCase					:			int									= 0

									lowerCamelCase					:			List[str]									= FlaxForcedBOSTokenLogitsProcessor(bos_token_id=_A      )

									# check that all scores are -inf except the bos_token_id score
									lowerCamelCase					:			Optional[int]									= ids_tensor((batch_size, 1),    vocab_size=20      )
									lowerCamelCase					:			Tuple									= 1
									lowerCamelCase					:			Tuple									= self._get_uniform_logits(_A,    _A      )
									lowerCamelCase					:			Optional[int]									= logits_processor(_A,    _A,    cur_len=_A      )
									self.assertTrue(jnp.isneginf(scores[:, bos_token_id + 1 :]      ).all()      )
									self.assertListEqual(scores[:, bos_token_id].tolist(),    4 * [0]      )  # score for bos_token_id shold be zero

									# check that bos_token_id is not forced if current length is greater than 1
									lowerCamelCase					:			str									= 3
									lowerCamelCase					:			Optional[int]									= self._get_uniform_logits(_A,    _A      )
									lowerCamelCase					:			Tuple									= logits_processor(_A,    _A,    cur_len=_A      )
									self.assertFalse(jnp.isinf(_A      ).any()      )
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			List[Any]									= 20
									lowerCamelCase					:			Optional[Any]									= 4
									lowerCamelCase					:			Dict									= 0
									lowerCamelCase					:			Optional[Any]									= 5

									lowerCamelCase					:			str									= FlaxForcedEOSTokenLogitsProcessor(max_length=_A,    eos_token_id=_A      )

									# check that all scores are -inf except the eos_token_id when max_length is reached
									lowerCamelCase					:			int									= ids_tensor((batch_size, 4),    vocab_size=20      )
									lowerCamelCase					:			int									= 4
									lowerCamelCase					:			Union[str, Any]									= self._get_uniform_logits(_A,    _A      )
									lowerCamelCase					:			List[str]									= logits_processor(_A,    _A,    cur_len=_A      )
									self.assertTrue(jnp.isneginf(scores[:, eos_token_id + 1 :]      ).all()      )
									self.assertListEqual(scores[:, eos_token_id].tolist(),    4 * [0]      )  # score for eos_token_id should be zero

									# check that eos_token_id is not forced if max_length is not reached
									lowerCamelCase					:			Any									= 3
									lowerCamelCase					:			Union[str, Any]									= self._get_uniform_logits(_A,    _A      )
									lowerCamelCase					:			Optional[Any]									= logits_processor(_A,    _A,    cur_len=_A      )
									self.assertFalse(jnp.isinf(_A      ).any()      )
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			int									= 4
									lowerCamelCase					:			List[str]									= 10
									lowerCamelCase					:			int									= 15
									lowerCamelCase					:			Any									= 2
									lowerCamelCase					:			Dict									= 1
									lowerCamelCase					:			Optional[int]									= 15

									# dummy input_ids and scores
									lowerCamelCase					:			Union[str, Any]									= ids_tensor((batch_size, sequence_length),    _A      )
									lowerCamelCase					:			str									= input_ids.copy()

									lowerCamelCase					:			Union[str, Any]									= self._get_uniform_logits(_A,    _A      )
									lowerCamelCase					:			int									= scores.copy()

									# instantiate all dist processors
									lowerCamelCase					:			List[str]									= FlaxTemperatureLogitsWarper(temperature=0.5      )
									lowerCamelCase					:			Optional[int]									= FlaxTopKLogitsWarper(3      )
									lowerCamelCase					:			Union[str, Any]									= FlaxTopPLogitsWarper(0.8      )

									# instantiate all logits processors
									lowerCamelCase					:			List[Any]									= FlaxMinLengthLogitsProcessor(min_length=10,    eos_token_id=_A      )
									lowerCamelCase					:			Tuple									= FlaxForcedBOSTokenLogitsProcessor(bos_token_id=_A      )
									lowerCamelCase					:			Any									= FlaxForcedEOSTokenLogitsProcessor(max_length=_A,    eos_token_id=_A      )

									lowerCamelCase					:			List[str]									= 10

									# no processor list
									lowerCamelCase					:			Dict									= temp_dist_warp(_A,    _A,    cur_len=_A      )
									lowerCamelCase					:			Optional[int]									= top_k_warp(_A,    _A,    cur_len=_A      )
									lowerCamelCase					:			Dict									= top_p_warp(_A,    _A,    cur_len=_A      )
									lowerCamelCase					:			int									= min_dist_proc(_A,    _A,    cur_len=_A      )
									lowerCamelCase					:			List[Any]									= bos_dist_proc(_A,    _A,    cur_len=_A      )
									lowerCamelCase					:			Optional[int]									= eos_dist_proc(_A,    _A,    cur_len=_A      )

									# with processor list
									lowerCamelCase					:			Dict									= FlaxLogitsProcessorList(
									    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]      )
									lowerCamelCase					:			List[str]									= processor(_A,    _A,    cur_len=_A      )

									# scores should be equal
									self.assertTrue(jnp.allclose(_A,    _A,    atol=1e-3      )      )

									# input_ids should never be changed
									self.assertListEqual(input_ids.tolist(),    input_ids_comp.tolist()      )




					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			int									= 4
									lowerCamelCase					:			List[Any]									= 10
									lowerCamelCase					:			Tuple									= 15
									lowerCamelCase					:			Union[str, Any]									= 2
									lowerCamelCase					:			Optional[int]									= 1
									lowerCamelCase					:			List[str]									= 15

									# dummy input_ids and scores
									lowerCamelCase					:			int									= ids_tensor((batch_size, sequence_length),    _A      )
									lowerCamelCase					:			Optional[int]									= input_ids.copy()

									lowerCamelCase					:			Optional[Any]									= self._get_uniform_logits(_A,    _A      )
									lowerCamelCase					:			List[Any]									= scores.copy()

									# instantiate all dist processors
									lowerCamelCase					:			Optional[int]									= FlaxTemperatureLogitsWarper(temperature=0.5      )
									lowerCamelCase					:			Optional[int]									= FlaxTopKLogitsWarper(3      )
									lowerCamelCase					:			Tuple									= FlaxTopPLogitsWarper(0.8      )

									# instantiate all logits processors
									lowerCamelCase					:			Optional[int]									= FlaxMinLengthLogitsProcessor(min_length=10,    eos_token_id=_A      )
									lowerCamelCase					:			Optional[int]									= FlaxForcedBOSTokenLogitsProcessor(bos_token_id=_A      )
									lowerCamelCase					:			Any									= FlaxForcedEOSTokenLogitsProcessor(max_length=_A,    eos_token_id=_A      )

									lowerCamelCase					:			List[str]									= 10

									# no processor list
									def run_no_processor_list(A,    A,    A      ):
													lowerCamelCase					:			Optional[Any]									= temp_dist_warp(_A,    _A,    cur_len=_A      )
													lowerCamelCase					:			Optional[Any]									= top_k_warp(_A,    _A,    cur_len=_A      )
													lowerCamelCase					:			List[Any]									= top_p_warp(_A,    _A,    cur_len=_A      )
													lowerCamelCase					:			Optional[Any]									= min_dist_proc(_A,    _A,    cur_len=_A      )
													lowerCamelCase					:			Tuple									= bos_dist_proc(_A,    _A,    cur_len=_A      )
													lowerCamelCase					:			Dict									= eos_dist_proc(_A,    _A,    cur_len=_A      )
													return scores

									# with processor list
									def run_processor_list(A,    A,    A      ):
													lowerCamelCase					:			Optional[Any]									= FlaxLogitsProcessorList(
													    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]      )
													lowerCamelCase					:			List[Any]									= processor(_A,    _A,    cur_len=_A      )
													return scores

									lowerCamelCase					:			Dict									= jax.jit(_A      )
									lowerCamelCase					:			Optional[Any]									= jax.jit(_A      )

									lowerCamelCase					:			List[str]									= jitted_run_no_processor_list(_A,    _A,    _A      )
									lowerCamelCase					:			List[Any]									= jitted_run_processor_list(_A,    _A,    _A      )

									# scores should be equal
									self.assertTrue(jnp.allclose(_A,    _A,    atol=1e-3      )      )

									# input_ids should never be changed
									self.assertListEqual(input_ids.tolist(),    input_ids_comp.tolist()      )


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


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

A			      =      {
    '169M': 768,
    '430M': 1024,
    '1B5': 2048,
    '3B': 2560,
    '7B': 4096,
    '14B': 5120,
}





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

								if name != "head.weight":
												lowerCamelCase					:			Any									= 'rwkv.' + name

								lowerCamelCase					:			Any									= weight
				return state_dict





def    UpperCAmelCase							(  UpperCAmelCase__							:							str       ,							UpperCAmelCase__							:							List[str]       ,							UpperCAmelCase__							:							Optional[int]       ,							UpperCAmelCase__							:							Optional[Any]=None       ,							UpperCAmelCase__							:							Tuple=None       ,							UpperCAmelCase__							:							List[Any]=False       ,							UpperCAmelCase__							:							Dict=None):
				# 1. If possible, build the tokenizer.
				if tokenizer_file is None:
								print('No `--tokenizer_file` provided, we will use the default tokenizer.')
								lowerCamelCase					:			Dict									= 5_02_77
								lowerCamelCase					:			Union[str, Any]									= AutoTokenizer.from_pretrained('EleutherAI/gpt-neox-20b')
				else:
								lowerCamelCase					:			int									= PreTrainedTokenizerFast(tokenizer_file=UpperCAmelCase__)
								lowerCamelCase					:			List[Any]									= len(UpperCAmelCase__)
				tokenizer.save_pretrained(UpperCAmelCase__)

				# 2. Build the config
				lowerCamelCase					:			Tuple									= 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:
																lowerCamelCase					:			Union[str, 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}.''')

				lowerCamelCase					:			List[Any]									= RwkvConfig(
				    vocab_size=UpperCAmelCase__       ,							num_hidden_layers=NUM_HIDDEN_LAYERS_MAPPING[size]       ,							hidden_size=HIDEN_SIZE_MAPPING[size]       ,							)
				config.save_pretrained(UpperCAmelCase__)

				# 3. Download model file then convert state_dict
				lowerCamelCase					:			Tuple									= hf_hub_download(UpperCAmelCase__       ,							UpperCAmelCase__)
				lowerCamelCase					:			Any									= torch.load(UpperCAmelCase__       ,							map_location='cpu')
				lowerCamelCase					:			str									= convert_state_dict(UpperCAmelCase__)

				# 4. Split in shards and save
				lowerCamelCase			, lowerCamelCase					:			Optional[int]									= shard_checkpoint(UpperCAmelCase__)
				for shard_file, shard in shards.items():
								torch.save(UpperCAmelCase__       ,							os.path.join(UpperCAmelCase__       ,							UpperCAmelCase__))

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

								# 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.')
								lowerCamelCase					:			List[str]									= list(shards.keys())

								del state_dict
								del shards
								gc.collect()

								for shard_file in shard_files:
												lowerCamelCase					:			Dict									= torch.load(os.path.join(UpperCAmelCase__       ,							UpperCAmelCase__))
												torch.save({k: v.cpu().clone() for k, v in state_dict.items()}       ,							os.path.join(UpperCAmelCase__       ,							UpperCAmelCase__))

				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.')
								lowerCamelCase					:			List[Any]									= AutoModelForCausalLM.from_pretrained(UpperCAmelCase__)
								model.push_to_hub(UpperCAmelCase__       ,							max_shard_size='2GB')
								tokenizer.push_to_hub(UpperCAmelCase__)


if __name__ == "__main__":
					A			      =      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.',
					)

					A			      =      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,
					)
