
'''simple docstring'''






import gc
import unittest

from transformers import MODEL_FOR_MASKED_LM_MAPPING, TF_MODEL_FOR_MASKED_LM_MAPPING, FillMaskPipeline, pipeline
from transformers.pipelines import PipelineException
from transformers.testing_utils import (
    is_pipeline_test,
    is_torch_available,
    nested_simplify,
    require_tf,
    require_torch,
    require_torch_gpu,
    slow,
)

from .test_pipelines_common import ANY




@is_pipeline_test
class     __snake_case		(   unittest.TestCase):
					_lowerCAmelCase			     =	MODEL_FOR_MASKED_LM_MAPPING
					_lowerCAmelCase			     =	TF_MODEL_FOR_MASKED_LM_MAPPING
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									super().tearDown()
									# clean-up as much as possible GPU memory occupied by PyTorch
									gc.collect()
									if is_torch_available():
													import torch

													torch.cuda.empty_cache()
					@require_tf
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			str									= pipeline(task='fill-mask',    model='sshleifer/tiny-distilroberta-base',    top_k=2,    framework='tf'      )
									lowerCamelCase					:			Dict									= unmasker('My name is <mask>'      )
									self.assertEqual(
									    nested_simplify(A,    decimals=6      ),    [
									        {'sequence': 'My name is grouped', 'score': 2.1e-05, 'token': 3_8015, 'token_str': ' grouped'},
									        {'sequence': 'My name is accuser', 'score': 2.1e-05, 'token': 2_5506, 'token_str': ' accuser'},
									    ],    )

									lowerCamelCase					:			Tuple									= unmasker('The largest city in France is <mask>'      )
									self.assertEqual(
									    nested_simplify(A,    decimals=6      ),    [
									        {
									            'sequence': 'The largest city in France is grouped',
									            'score': 2.1e-05,
									            'token': 3_8015,
									            'token_str': ' grouped',
									        },
									        {
									            'sequence': 'The largest city in France is accuser',
									            'score': 2.1e-05,
									            'token': 2_5506,
									            'token_str': ' accuser',
									        },
									    ],    )

									lowerCamelCase					:			List[str]									= unmasker('My name is <mask>',    targets=[' Patrick', ' Clara', ' Teven'],    top_k=3      )
									self.assertEqual(
									    nested_simplify(A,    decimals=6      ),    [
									        {'sequence': 'My name is Clara', 'score': 2e-05, 'token': 1_3606, 'token_str': ' Clara'},
									        {'sequence': 'My name is Patrick', 'score': 2e-05, 'token': 3499, 'token_str': ' Patrick'},
									        {'sequence': 'My name is Te', 'score': 1.9e-05, 'token': 2941, 'token_str': ' Te'},
									    ],    )
					@require_torch
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			int									= pipeline(task='fill-mask',    model='sshleifer/tiny-distilroberta-base',    top_k=2,    framework='pt'      )

									lowerCamelCase					:			str									= unmasker('My name is <mask>'      )
									self.assertEqual(
									    nested_simplify(A,    decimals=6      ),    [
									        {'sequence': 'My name is Maul', 'score': 2.2e-05, 'token': 3_5676, 'token_str': ' Maul'},
									        {'sequence': 'My name isELS', 'score': 2.2e-05, 'token': 1_6416, 'token_str': 'ELS'},
									    ],    )

									lowerCamelCase					:			Optional[int]									= unmasker('The largest city in France is <mask>'      )
									self.assertEqual(
									    nested_simplify(A,    decimals=6      ),    [
									        {
									            'sequence': 'The largest city in France is Maul',
									            'score': 2.2e-05,
									            'token': 3_5676,
									            'token_str': ' Maul',
									        },
									        {'sequence': 'The largest city in France isELS', 'score': 2.2e-05, 'token': 1_6416, 'token_str': 'ELS'},
									    ],    )

									lowerCamelCase					:			Optional[int]									= unmasker('My name is <mask>',    targets=[' Patrick', ' Clara', ' Teven'],    top_k=3      )
									self.assertEqual(
									    nested_simplify(A,    decimals=6      ),    [
									        {'sequence': 'My name is Patrick', 'score': 2.1e-05, 'token': 3499, 'token_str': ' Patrick'},
									        {'sequence': 'My name is Te', 'score': 2e-05, 'token': 2941, 'token_str': ' Te'},
									        {'sequence': 'My name is Clara', 'score': 2e-05, 'token': 1_3606, 'token_str': ' Clara'},
									    ],    )

									lowerCamelCase					:			Optional[Any]									= unmasker('My name is <mask> <mask>',    top_k=2      )

									self.assertEqual(
									    nested_simplify(A,    decimals=6      ),    [
									        [
									            {
									                'score': 2.2e-05,
									                'token': 3_5676,
									                'token_str': ' Maul',
									                'sequence': '<s>My name is Maul<mask></s>',
									            },
									            {'score': 2.2e-05, 'token': 1_6416, 'token_str': 'ELS', 'sequence': '<s>My name isELS<mask></s>'},
									        ],
									        [
									            {
									                'score': 2.2e-05,
									                'token': 3_5676,
									                'token_str': ' Maul',
									                'sequence': '<s>My name is<mask> Maul</s>',
									            },
									            {'score': 2.2e-05, 'token': 1_6416, 'token_str': 'ELS', 'sequence': '<s>My name is<mask>ELS</s>'},
									        ],
									    ],    )
					@require_torch_gpu
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Dict									= pipeline('fill-mask',    model='hf-internal-testing/tiny-random-distilbert',    device=0,    framework='pt'      )

									# convert model to fp16
									pipe.model.half()

									lowerCamelCase					:			str									= pipe('Paris is the [MASK] of France.'      )
									# We actually don't care about the result, we just want to make sure
									# it works, meaning the float16 tensor got casted back to float32
									# for postprocessing.
									self.assertIsInstance(A,    A      )
					@slow
					@require_torch
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Any									= pipeline(task='fill-mask',    model='distilroberta-base',    top_k=2,    framework='pt'      )
									self.run_large_test(A      )
					@slow
					@require_tf
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Optional[int]									= pipeline(task='fill-mask',    model='distilroberta-base',    top_k=2,    framework='tf'      )
									self.run_large_test(A      )
					def     UpperCAmelCase_							(     self,    A      ):

									"""simple docstring"""


									lowerCamelCase					:			str									= unmasker('My name is <mask>'      )
									self.assertEqual(
									    nested_simplify(A      ),    [
									        {'sequence': 'My name is John', 'score': 0.008, 'token': 610, 'token_str': ' John'},
									        {'sequence': 'My name is Chris', 'score': 0.007, 'token': 1573, 'token_str': ' Chris'},
									    ],    )
									lowerCamelCase					:			Union[str, Any]									= unmasker('The largest city in France is <mask>'      )
									self.assertEqual(
									    nested_simplify(A      ),    [
									        {
									            'sequence': 'The largest city in France is Paris',
									            'score': 0.251,
									            'token': 2201,
									            'token_str': ' Paris',
									        },
									        {
									            'sequence': 'The largest city in France is Lyon',
									            'score': 0.214,
									            'token': 1_2790,
									            'token_str': ' Lyon',
									        },
									    ],    )

									lowerCamelCase					:			List[Any]									= unmasker('My name is <mask>',    targets=[' Patrick', ' Clara', ' Teven'],    top_k=3      )
									self.assertEqual(
									    nested_simplify(A      ),    [
									        {'sequence': 'My name is Patrick', 'score': 0.005, 'token': 3499, 'token_str': ' Patrick'},
									        {'sequence': 'My name is Clara', 'score': 0.000, 'token': 1_3606, 'token_str': ' Clara'},
									        {'sequence': 'My name is Te', 'score': 0.000, 'token': 2941, 'token_str': ' Te'},
									    ],    )
					@require_torch
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			int									= pipeline(task='fill-mask',    model='sshleifer/tiny-distilroberta-base',    framework='pt'      )
									lowerCamelCase					:			Tuple									= None
									lowerCamelCase					:			Union[str, Any]									= None
									self.run_pipeline_test(A,    []      )
					@require_tf
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Union[str, Any]									= pipeline(task='fill-mask',    model='sshleifer/tiny-distilroberta-base',    framework='tf'      )
									lowerCamelCase					:			Any									= None
									lowerCamelCase					:			Any									= None
									self.run_pipeline_test(A,    []      )
					def     UpperCAmelCase_							(     self,    A,    A,    A      ):

									"""simple docstring"""


									if tokenizer is None or tokenizer.mask_token_id is None:
													self.skipTest('The provided tokenizer has no mask token, (probably reformer or wav2vec2)'      )

									lowerCamelCase					:			Union[str, Any]									= FillMaskPipeline(model=A,    tokenizer=A      )
									lowerCamelCase					:			str									= [
									    F'''This is another {tokenizer.mask_token} test''',
									]
									return fill_masker, examples
					def     UpperCAmelCase_							(     self,    A,    A      ):

									"""simple docstring"""


									lowerCamelCase					:			List[str]									= fill_masker.tokenizer
									lowerCamelCase					:			str									= fill_masker.model

									lowerCamelCase					:			Optional[Any]									= fill_masker(
									    F'''This is a {tokenizer.mask_token}''',    )
									self.assertEqual(
									    A,    [
									        {'sequence': ANY(A      ), 'score': ANY(A      ), 'token': ANY(A      ), 'token_str': ANY(A      )},
									        {'sequence': ANY(A      ), 'score': ANY(A      ), 'token': ANY(A      ), 'token_str': ANY(A      )},
									        {'sequence': ANY(A      ), 'score': ANY(A      ), 'token': ANY(A      ), 'token_str': ANY(A      )},
									        {'sequence': ANY(A      ), 'score': ANY(A      ), 'token': ANY(A      ), 'token_str': ANY(A      )},
									        {'sequence': ANY(A      ), 'score': ANY(A      ), 'token': ANY(A      ), 'token_str': ANY(A      )},
									    ],    )

									lowerCamelCase					:			List[str]									= fill_masker([F'''This is a {tokenizer.mask_token}''']      )
									self.assertEqual(
									    A,    [
									        {'sequence': ANY(A      ), 'score': ANY(A      ), 'token': ANY(A      ), 'token_str': ANY(A      )},
									        {'sequence': ANY(A      ), 'score': ANY(A      ), 'token': ANY(A      ), 'token_str': ANY(A      )},
									        {'sequence': ANY(A      ), 'score': ANY(A      ), 'token': ANY(A      ), 'token_str': ANY(A      )},
									        {'sequence': ANY(A      ), 'score': ANY(A      ), 'token': ANY(A      ), 'token_str': ANY(A      )},
									        {'sequence': ANY(A      ), 'score': ANY(A      ), 'token': ANY(A      ), 'token_str': ANY(A      )},
									    ],    )

									lowerCamelCase					:			Dict									= fill_masker([F'''This is a {tokenizer.mask_token}''', F'''Another {tokenizer.mask_token} great test.''']      )
									self.assertEqual(
									    A,    [
									        [
									            {'sequence': ANY(A      ), 'score': ANY(A      ), 'token': ANY(A      ), 'token_str': ANY(A      )},
									            {'sequence': ANY(A      ), 'score': ANY(A      ), 'token': ANY(A      ), 'token_str': ANY(A      )},
									            {'sequence': ANY(A      ), 'score': ANY(A      ), 'token': ANY(A      ), 'token_str': ANY(A      )},
									            {'sequence': ANY(A      ), 'score': ANY(A      ), 'token': ANY(A      ), 'token_str': ANY(A      )},
									            {'sequence': ANY(A      ), 'score': ANY(A      ), 'token': ANY(A      ), 'token_str': ANY(A      )},
									        ],
									        [
									            {'sequence': ANY(A      ), 'score': ANY(A      ), 'token': ANY(A      ), 'token_str': ANY(A      )},
									            {'sequence': ANY(A      ), 'score': ANY(A      ), 'token': ANY(A      ), 'token_str': ANY(A      )},
									            {'sequence': ANY(A      ), 'score': ANY(A      ), 'token': ANY(A      ), 'token_str': ANY(A      )},
									            {'sequence': ANY(A      ), 'score': ANY(A      ), 'token': ANY(A      ), 'token_str': ANY(A      )},
									            {'sequence': ANY(A      ), 'score': ANY(A      ), 'token': ANY(A      ), 'token_str': ANY(A      )},
									        ],
									    ],    )

									with self.assertRaises(A      ):
													fill_masker([None]      )
									# No mask_token is not supported
									with self.assertRaises(A      ):
													fill_masker('This is'      )

									self.run_test_top_k(A,    A      )
									self.run_test_targets(A,    A      )
									self.run_test_top_k_targets(A,    A      )
									self.fill_mask_with_duplicate_targets_and_top_k(A,    A      )
									self.fill_mask_with_multiple_masks(A,    A      )
					def     UpperCAmelCase_							(     self,    A,    A      ):

									"""simple docstring"""


									lowerCamelCase					:			int									= tokenizer.get_vocab()
									lowerCamelCase					:			str									= sorted(vocab.keys()      )[:2]
									# Pipeline argument
									lowerCamelCase					:			List[str]									= FillMaskPipeline(model=A,    tokenizer=A,    targets=A      )
									lowerCamelCase					:			str									= fill_masker(F'''This is a {tokenizer.mask_token}'''      )
									self.assertEqual(
									    A,    [
									        {'sequence': ANY(A      ), 'score': ANY(A      ), 'token': ANY(A      ), 'token_str': ANY(A      )},
									        {'sequence': ANY(A      ), 'score': ANY(A      ), 'token': ANY(A      ), 'token_str': ANY(A      )},
									    ],    )
									lowerCamelCase					:			List[Any]									= {vocab[el] for el in targets}
									self.assertEqual({el['token'] for el in outputs},    A      )
									lowerCamelCase					:			Any									= [tokenizer.decode([x]      ) for x in target_ids]
									self.assertEqual({el['token_str'] for el in outputs},    set(A      )      )

									# Call argument
									lowerCamelCase					:			List[Any]									= FillMaskPipeline(model=A,    tokenizer=A      )
									lowerCamelCase					:			List[Any]									= fill_masker(F'''This is a {tokenizer.mask_token}''',    targets=A      )
									self.assertEqual(
									    A,    [
									        {'sequence': ANY(A      ), 'score': ANY(A      ), 'token': ANY(A      ), 'token_str': ANY(A      )},
									        {'sequence': ANY(A      ), 'score': ANY(A      ), 'token': ANY(A      ), 'token_str': ANY(A      )},
									    ],    )
									lowerCamelCase					:			str									= {vocab[el] for el in targets}
									self.assertEqual({el['token'] for el in outputs},    A      )
									lowerCamelCase					:			Optional[Any]									= [tokenizer.decode([x]      ) for x in target_ids]
									self.assertEqual({el['token_str'] for el in outputs},    set(A      )      )

									# Score equivalence
									lowerCamelCase					:			List[Any]									= fill_masker(F'''This is a {tokenizer.mask_token}''',    targets=A      )
									lowerCamelCase					:			List[str]									= [top_mask['token_str'] for top_mask in outputs]
									lowerCamelCase					:			List[str]									= [top_mask['score'] for top_mask in outputs]

									# For some BPE tokenizers, `</w>` is removed during decoding, so `token_str` won't be the same as in `targets`.
									if set(A      ) == set(A      ):
													lowerCamelCase					:			Union[str, Any]									= fill_masker(F'''This is a {tokenizer.mask_token}''',    targets=A      )
													lowerCamelCase					:			Dict									= [top_mask['score'] for top_mask in unmasked_targets]
													self.assertEqual(nested_simplify(A      ),    nested_simplify(A      )      )

									# Raises with invalid
									with self.assertRaises(A      ):
													lowerCamelCase					:			Union[str, Any]									= fill_masker(F'''This is a {tokenizer.mask_token}''',    targets=[]      )
									# For some tokenizers, `""` is actually in the vocabulary and the expected error won't raised
									if "" not in tokenizer.get_vocab():
													with self.assertRaises(A      ):
																	lowerCamelCase					:			Tuple									= fill_masker(F'''This is a {tokenizer.mask_token}''',    targets=['']      )
													with self.assertRaises(A      ):
																	lowerCamelCase					:			Dict									= fill_masker(F'''This is a {tokenizer.mask_token}''',    targets=''      )
					def     UpperCAmelCase_							(     self,    A,    A      ):

									"""simple docstring"""


									lowerCamelCase					:			int									= FillMaskPipeline(model=A,    tokenizer=A,    top_k=2      )
									lowerCamelCase					:			Dict									= fill_masker(F'''This is a {tokenizer.mask_token}'''      )
									self.assertEqual(
									    A,    [
									        {'sequence': ANY(A      ), 'score': ANY(A      ), 'token': ANY(A      ), 'token_str': ANY(A      )},
									        {'sequence': ANY(A      ), 'score': ANY(A      ), 'token': ANY(A      ), 'token_str': ANY(A      )},
									    ],    )

									lowerCamelCase					:			Optional[Any]									= FillMaskPipeline(model=A,    tokenizer=A      )
									lowerCamelCase					:			str									= fill_masker(F'''This is a {tokenizer.mask_token}''',    top_k=2      )
									self.assertEqual(
									    A,    [
									        {'sequence': ANY(A      ), 'score': ANY(A      ), 'token': ANY(A      ), 'token_str': ANY(A      )},
									        {'sequence': ANY(A      ), 'score': ANY(A      ), 'token': ANY(A      ), 'token_str': ANY(A      )},
									    ],    )
									self.assertEqual(nested_simplify(A      ),    nested_simplify(A      )      )
					def     UpperCAmelCase_							(     self,    A,    A      ):

									"""simple docstring"""


									lowerCamelCase					:			str									= tokenizer.get_vocab()
									lowerCamelCase					:			int									= FillMaskPipeline(model=A,    tokenizer=A      )

									# top_k=2, ntargets=3
									lowerCamelCase					:			str									= sorted(vocab.keys()      )[:3]
									lowerCamelCase					:			Dict									= fill_masker(F'''This is a {tokenizer.mask_token}''',    top_k=2,    targets=A      )

									# If we use the most probably targets, and filter differently, we should still
									# have the same results
									lowerCamelCase					:			Union[str, Any]									= [el['token_str'] for el in sorted(A,    key=lambda A      : x["score"],    reverse=A      )]
									# For some BPE tokenizers, `</w>` is removed during decoding, so `token_str` won't be the same as in `targets`.
									if set(A      ).issubset(A      ):
													lowerCamelCase					:			Tuple									= fill_masker(F'''This is a {tokenizer.mask_token}''',    top_k=3,    targets=A      )
													# They should yield exactly the same result
													self.assertEqual(nested_simplify(A      ),    nested_simplify(A      )      )
					def     UpperCAmelCase_							(     self,    A,    A      ):

									"""simple docstring"""


									lowerCamelCase					:			List[str]									= FillMaskPipeline(model=A,    tokenizer=A      )
									lowerCamelCase					:			Optional[int]									= tokenizer.get_vocab()
									# String duplicates + id duplicates
									lowerCamelCase					:			List[str]									= sorted(vocab.keys()      )[:3]
									lowerCamelCase					:			List[str]									= [targets[0], targets[1], targets[0], targets[2], targets[1]]
									lowerCamelCase					:			Dict									= fill_masker(F'''My name is {tokenizer.mask_token}''',    targets=A,    top_k=10      )

									# The target list contains duplicates, so we can't output more
									# than them
									self.assertEqual(len(A      ),    3      )




					def     UpperCAmelCase_							(     self,    A,    A      ):

									"""simple docstring"""


									lowerCamelCase					:			str									= FillMaskPipeline(model=A,    tokenizer=A      )

									lowerCamelCase					:			str									= fill_masker(
									    F'''This is a {tokenizer.mask_token} {tokenizer.mask_token} {tokenizer.mask_token}''',    top_k=2      )
									self.assertEqual(
									    A,    [
									        [
									            {'sequence': ANY(A      ), 'score': ANY(A      ), 'token': ANY(A      ), 'token_str': ANY(A      )},
									            {'sequence': ANY(A      ), 'score': ANY(A      ), 'token': ANY(A      ), 'token_str': ANY(A      )},
									        ],
									        [
									            {'sequence': ANY(A      ), 'score': ANY(A      ), 'token': ANY(A      ), 'token_str': ANY(A      )},
									            {'sequence': ANY(A      ), 'score': ANY(A      ), 'token': ANY(A      ), 'token_str': ANY(A      )},
									        ],
									        [
									            {'sequence': ANY(A      ), 'score': ANY(A      ), 'token': ANY(A      ), 'token_str': ANY(A      )},
									            {'sequence': ANY(A      ), 'score': ANY(A      ), 'token': ANY(A      ), 'token_str': ANY(A      )},
									        ],
									    ],    )


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