





from typing import List, Optional, Union

import numpy as np
import tensorflow as tf

from .utils import logging


UpperCAmelCase__						:			Optional[Any]									=logging.get_logger(__name__)





def      _lowercase					(						_UpperCAmelCase				)						->  List[int]:
				if isinstance(_UpperCAmelCase			,      np.ndarray				):
								return list(tensor.shape				)

				lowerCamelCase	 =tf.shape(_UpperCAmelCase				)

				if tensor.shape == tf.TensorShape(_UpperCAmelCase				):
								return dynamic

				lowerCamelCase	 =tensor.shape.as_list()

				return [dynamic[i] if s is None else s for i, s in enumerate(_UpperCAmelCase				)]





def      _lowercase					(						_UpperCAmelCase			,      _UpperCAmelCase = None			,      _UpperCAmelCase = None				)						->  tf.Tensor:
				return tf.nn.softmax(logits=logits + 1e-9			,      axis=_UpperCAmelCase			,      name=_UpperCAmelCase				)





def      _lowercase					(						_UpperCAmelCase			,      _UpperCAmelCase			,      _UpperCAmelCase			,      _UpperCAmelCase=1e-5			,      _UpperCAmelCase=-1				)						->  Optional[Any]:
				# This is a very simplified functional layernorm, designed to duplicate
				# the functionality of PyTorch nn.functional.layer_norm when this is needed to port
				# models in Transformers.

				if weight.shape.rank != 1 or bias.shape.rank != 1 or not isinstance(_UpperCAmelCase			,      _UpperCAmelCase				):
								raise NotImplementedError("""Only 1D weight and bias tensors are supported for now, with only a single axis."""				)

				# Get mean and variance on the axis to be normalized
				lowerCamelCase ,  lowerCamelCase	 =tf.nn.moments(_UpperCAmelCase			,      axes=[axis]			,      keepdims=_UpperCAmelCase				)

				if axis != -1:
								# Reshape scale and weight to have the same rank as inputs, but with 1 dimensions
								# on every dimension except axis
								lowerCamelCase	 =[1] * inputs.shape.rank
								lowerCamelCase	 =shape_list(_UpperCAmelCase				)[axis]
								lowerCamelCase	 =tf.reshape(_UpperCAmelCase			,      _UpperCAmelCase				)
								lowerCamelCase	 =tf.reshape(_UpperCAmelCase			,      _UpperCAmelCase				)

				# Compute layer normalization using the batch_normalization
				# function.
				lowerCamelCase	 =tf.nn.batch_normalization(
				    _UpperCAmelCase			,      _UpperCAmelCase			,      _UpperCAmelCase			,      offset=_UpperCAmelCase			,      scale=_UpperCAmelCase			,      variance_epsilon=_UpperCAmelCase			,      )
				return outputs





def      _lowercase					(						_UpperCAmelCase			,      _UpperCAmelCase=0			,      _UpperCAmelCase=-1				)						->  List[str]:
				# Replicates the behavior of torch.flatten in TF

				# If end_dim or start_dim is negative, count them from the end
				if end_dim < 0:
								end_dim += input.shape.rank
				if start_dim < 0:
								start_dim += input.shape.rank

				if start_dim == end_dim:
								return input

				lowerCamelCase	 =tf.shape(_UpperCAmelCase				)
				lowerCamelCase	 =tf.math.reduce_prod(in_shape[start_dim : end_dim + 1]				)
				lowerCamelCase	 =tf.concat([in_shape[:start_dim], [flattened_dim], in_shape[end_dim + 1 :]]			,      axis=0				)
				return tf.reshape(_UpperCAmelCase			,      _UpperCAmelCase				)





def      _lowercase					(						_UpperCAmelCase				)						->  tf.Tensor:
				if not isinstance(_UpperCAmelCase			,      tf.Tensor				):
								lowerCamelCase	 =tf.convert_to_tensor(_UpperCAmelCase				)  # Catches stray NumPy inputs
				if encoder_attention_mask.shape.rank == 3:
								lowerCamelCase	 =encoder_attention_mask[:, None, :, :]
				if encoder_attention_mask.shape.rank == 2:
								lowerCamelCase	 =encoder_attention_mask[:, None, None, :]
				# T5 has a mask that can compare sequence ids, we can simulate this here with this transposition
				# Cf. https://github.com/tensorflow/mesh/blob/8d2465e9bc93129b913b5ccc6a59aa97abd96ec6/mesh_tensorflow
				# /transformer/transformer_layers.py#L270
				# encoder_extended_attention_mask = (encoder_extended_attention_mask ==
				# encoder_extended_attention_mask.transpose(-1, -2))
				lowerCamelCase	 =(
				    tf.cast(1			,      encoder_attention_mask.dtype				) - encoder_extended_attention_mask
				) * encoder_extended_attention_mask.dtype.min

				return encoder_extended_attention_mask





def      _lowercase					(						_UpperCAmelCase			,      _UpperCAmelCase			,      _UpperCAmelCase = "input_ids"				)						->  None:
				tf.debugging.assert_less(
				    _UpperCAmelCase			,      tf.cast(_UpperCAmelCase			,      dtype=tensor.dtype				)			,      message=(
				        F"""The maximum value of {tensor_name} ({tf.math.reduce_max(_UpperCAmelCase				)}) must be smaller than the embedding """
				        F"""layer's input dimension ({embed_dim}). The likely cause is some problem at tokenization time."""
				    )			,      )





def      _lowercase					(						_UpperCAmelCase			,      _UpperCAmelCase			,      _UpperCAmelCase				)						->  Any:
				lowerCamelCase	 =6_45_12
				# Check that no item in `data` is larger than `HDF5_OBJECT_HEADER_LIMIT`
				# because in that case even chunking the array would not make the saving
				# possible.
				lowerCamelCase	 =[x for x in data if len(_UpperCAmelCase				) > HDF5_OBJECT_HEADER_LIMIT]

				# Expecting this to never be true.
				if bad_attributes:
								raise RuntimeError(
								    """The following attributes cannot be saved to HDF5 file because """
								    F"""they are larger than {HDF5_OBJECT_HEADER_LIMIT} """
								    F"""bytes: {bad_attributes}"""				)

				lowerCamelCase	 =np.asarray(_UpperCAmelCase				)

				lowerCamelCase	 =1
				lowerCamelCase	 =np.array_split(_UpperCAmelCase			,      _UpperCAmelCase				)

				# This will never loop forever thanks to the test above.
				while any(x.nbytes > HDF5_OBJECT_HEADER_LIMIT for x in chunked_data				):
								num_chunks += 1
								lowerCamelCase	 =np.array_split(_UpperCAmelCase			,      _UpperCAmelCase				)

				if num_chunks > 1:
								for chunk_id, chunk_data in enumerate(_UpperCAmelCase				):
												lowerCamelCase	 =chunk_data
				else:
								lowerCamelCase	 =data





def      _lowercase					(						_UpperCAmelCase			,      _UpperCAmelCase				)						->  List[str]:
				if name in group.attrs:
								lowerCamelCase	 =[n.decode("""utf8"""				) if hasattr(_UpperCAmelCase			,      """decode"""				) else n for n in group.attrs[name]]
				else:
								lowerCamelCase	 =[]
								lowerCamelCase	 =0
								while "%s%d" % (name, chunk_id) in group.attrs:
												data.extend(
												    [n.decode("""utf8"""				) if hasattr(_UpperCAmelCase			,      """decode"""				) else n for n in group.attrs["""%s%d""" % (name, chunk_id)]]				)
												chunk_id += 1
				return data





def      _lowercase					(						_UpperCAmelCase				)						->  Dict:

				def _expand_single_ad_tensor(_UpperCAmelCase				):
								if isinstance(_UpperCAmelCase			,      tf.Tensor				) and t.shape.rank == 1:
												return tf.expand_dims(_UpperCAmelCase			,      axis=-1				)
								return t

				return tf.nest.map_structure(_expand_single_ad_tensor			,      _UpperCAmelCase				)












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


UpperCAmelCase__						:			Optional[int]									={
    '''169M''': 12,
    '''430M''': 24,
    '''1B5''': 24,
    '''3B''': 32,
    '''7B''': 32,
    '''14B''': 40,
}

UpperCAmelCase__						:			str									={
    '''169M''': 7_68,
    '''430M''': 10_24,
    '''1B5''': 20_48,
    '''3B''': 25_60,
    '''7B''': 40_96,
    '''14B''': 51_20,
}





def      _lowercase					(						_UpperCAmelCase				)						->  Tuple:
				lowerCamelCase	 =list(state_dict.keys()				)
				for name in state_dict_keys:
								lowerCamelCase	 =state_dict.pop(_UpperCAmelCase				)
								# emb -> embedding
								if name.startswith("""emb."""				):
												lowerCamelCase	 =name.replace("""emb."""			,      """embeddings."""				)
								# ln_0 -> pre_ln (only present at block 0)
								if name.startswith("""blocks.0.ln0"""				):
												lowerCamelCase	 =name.replace("""blocks.0.ln0"""			,      """blocks.0.pre_ln"""				)
								# att -> attention
								lowerCamelCase	 =re.sub(r"""blocks\.(\d+)\.att"""			,      r"""blocks.\1.attention"""			,      _UpperCAmelCase				)
								# ffn -> feed_forward
								lowerCamelCase	 =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	 =name.replace(""".time_mix_k"""			,      """.time_mix_key"""				)
								# time_mix_v -> time_mix_value and reshape
								if name.endswith(""".time_mix_v"""				):
												lowerCamelCase	 =name.replace(""".time_mix_v"""			,      """.time_mix_value"""				)
								# time_mix_r -> time_mix_key and reshape
								if name.endswith(""".time_mix_r"""				):
												lowerCamelCase	 =name.replace(""".time_mix_r"""			,      """.time_mix_receptance"""				)

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

								lowerCamelCase	 =weight
				return state_dict





def      _lowercase					(						_UpperCAmelCase			,      _UpperCAmelCase			,      _UpperCAmelCase			,      _UpperCAmelCase=None			,      _UpperCAmelCase=None			,      _UpperCAmelCase=False			,      _UpperCAmelCase=None				)						->  Tuple:
				# 1. If possible, build the tokenizer.
				if tokenizer_file is None:
								print("""No `--tokenizer_file` provided, we will use the default tokenizer."""				)
								lowerCamelCase	 =5_02_77
								lowerCamelCase	 =AutoTokenizer.from_pretrained("""EleutherAI/gpt-neox-20b"""				)
				else:
								lowerCamelCase	 =PreTrainedTokenizerFast(tokenizer_file=_UpperCAmelCase				)
								lowerCamelCase	 =len(_UpperCAmelCase				)
				tokenizer.save_pretrained(_UpperCAmelCase				)

				# 2. Build the config
				lowerCamelCase	 =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	 =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	 =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	 =hf_hub_download(_UpperCAmelCase			,      _UpperCAmelCase				)
				lowerCamelCase	 =torch.load(_UpperCAmelCase			,      map_location="""cpu"""				)
				lowerCamelCase	 =convert_state_dict(_UpperCAmelCase				)

				# 4. Split in shards and save
				lowerCamelCase ,  lowerCamelCase	 =shard_checkpoint(_UpperCAmelCase				)
				for shard_file, shard in shards.items():
								torch.save(_UpperCAmelCase			,      os.path.join(_UpperCAmelCase			,      _UpperCAmelCase				)				)

				if index is not None:
								lowerCamelCase	 =os.path.join(_UpperCAmelCase			,      _UpperCAmelCase				)
								# Save the index as well
								with open(_UpperCAmelCase			,      """w"""			,      encoding="""utf-8"""				) as f:
												lowerCamelCase	 =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(shards.keys()				)

								del state_dict
								del shards
								gc.collect()

								for shard_file in shard_files:
												lowerCamelCase	 =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	 =AutoModelForCausalLM.from_pretrained(_UpperCAmelCase				)
								model.push_to_hub(_UpperCAmelCase			,      max_shard_size="""2GB"""				)
								tokenizer.push_to_hub(_UpperCAmelCase				)


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

							UpperCAmelCase__						:			List[Any]									=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,
							)





