





'''simple docstring'''



import time
from dataclasses import dataclass
from multiprocessing import Pool
from unittest import TestCase
from unittest.mock import patch

import multiprocess
import numpy as np
import pytest

from datasets.utils.py_utils import (
    NestedDataStructure,
    asdict,
    iflatmap_unordered,
    map_nested,
    temp_seed,
    temporary_assignment,
    zip_dict,
)

from .utils import require_tf, require_torch




def 				_lowerCAmelCase  (			lowerCamelCase_     :					List[str]			):  # picklable for multiprocessing
			return x.sum()




def 				_lowerCAmelCase  (			lowerCamelCase_     :					List[str]			):  # picklable for multiprocessing
			return i + 1






@dataclass
class 							__lowercase    :







	'''simple docstring'''

	a      :			int
	a      :			str






class 							__lowercase    ( lowerCAmelCase__  ):







	'''simple docstring'''

	def   _UpperCAmelCase       (self	)							->							Tuple:







				'''simple docstring'''





				__lowercase			  =							{}
				__lowercase			  =							[]
				__lowercase			  =							1
				__lowercase			  =							[1, 2]
				__lowercase			  =							{'a': 1, 'b': 2}
				__lowercase			  =							{'a': [1, 2], 'b': [3, 4]}
				__lowercase			  =							{'a': {'1': 1}, 'b': 2}
				__lowercase			  =							{'a': 1, 'b': 2, 'c': 3, 'd': 4}
				__lowercase			  =							{}
				__lowercase			  =							[]
				__lowercase			  =							2
				__lowercase			  =							[2, 3]
				__lowercase			  =							{'a': 2, 'b': 3}
				__lowercase			  =							{'a': [2, 3], 'b': [4, 5]}
				__lowercase			  =							{'a': {'1': 2}, 'b': 3}
				__lowercase			  =							{'a': 2, 'b': 3, 'c': 4, 'd': 5}
				self.assertEqual(map_nested(_A       ,_A	)       ,_A	)
				self.assertEqual(map_nested(_A       ,_A	)       ,_A	)
				self.assertEqual(map_nested(_A       ,_A	)       ,_A	)
				self.assertEqual(map_nested(_A       ,_A	)       ,_A	)
				self.assertEqual(map_nested(_A       ,_A	)       ,_A	)
				self.assertEqual(map_nested(_A       ,_A	)       ,_A	)
				self.assertEqual(map_nested(_A       ,_A	)       ,_A	)
				self.assertEqual(map_nested(_A       ,_A	)       ,_A	)

				__lowercase			  =							2
				self.assertEqual(map_nested(_A       ,_A       ,num_proc=_A	)       ,_A	)
				self.assertEqual(map_nested(_A       ,_A       ,num_proc=_A	)       ,_A	)
				self.assertEqual(map_nested(_A       ,_A       ,num_proc=_A	)       ,_A	)
				self.assertEqual(map_nested(_A       ,_A       ,num_proc=_A	)       ,_A	)
				self.assertEqual(map_nested(_A       ,_A       ,num_proc=_A	)       ,_A	)
				self.assertEqual(map_nested(_A       ,_A       ,num_proc=_A	)       ,_A	)
				self.assertEqual(map_nested(_A       ,_A       ,num_proc=_A	)       ,_A	)
				self.assertEqual(map_nested(_A       ,_A       ,num_proc=_A	)       ,_A	)

				__lowercase			  =							{'a': np.eye(2	), 'b': np.zeros(3	), 'c': np.ones(2	)}
				__lowercase			  =							{'a': 2, 'b': 0, 'c': 2}
				__lowercase			  =							{
				    'a': np.eye(2	).astype(_A	),
				    'b': np.zeros(3	).astype(_A	),
				    'c': np.ones(2	).astype(_A	),
				}
				self.assertEqual(map_nested(_A       ,_A       ,map_numpy=_A	)       ,_A	)
				self.assertEqual(
				    {k: v.tolist() for k, v in map_nested(_A       ,_A       ,map_numpy=_A	).items()}       ,{k: v.tolist() for k, v in expected_map_nested_sna_int.items()}       ,)
				self.assertEqual(map_nested(_A       ,_A       ,map_numpy=_A       ,num_proc=_A	)       ,_A	)
				self.assertEqual(
				    {k: v.tolist() for k, v in map_nested(_A       ,_A       ,map_numpy=_A       ,num_proc=_A	).items()}       ,{k: v.tolist() for k, v in expected_map_nested_sna_int.items()}       ,)
				with self.assertRaises(_A	):  # can't pickle a local lambda
							map_nested(lambda _lowerCamelCase	: x + 1       ,_A       ,num_proc=_A	)


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







				'''simple docstring'''





				__lowercase			  =							{'a': 1, 'b': 2}
				__lowercase			  =							{'a': 3, 'b': 4}
				__lowercase			  =							{'a': 5, 'b': 6}
				__lowercase			  =							sorted([('''a''', (1, 3, 5)), ('''b''', (2, 4, 6))]	)
				self.assertEqual(sorted(zip_dict(_A       ,_A       ,_A	)	)       ,_A	)


	def   _UpperCAmelCase       (self	)							->							str:







				'''simple docstring'''





				class 							__lowercase    :







					'''simple docstring'''

					a      :			Optional[int]     =    """bar"""





				__lowercase			  =							Foo()
				self.assertEqual(foo.my_attr       ,'''bar'''	)
				with temporary_assignment(_A       ,'''my_attr'''       ,'''BAR'''	):
							self.assertEqual(foo.my_attr       ,'''BAR'''	)
				self.assertEqual(foo.my_attr       ,'''bar'''	)





@pytest.mark.parametrize(
    '''iterable_length, num_proc, expected_num_proc'''						,    [
        (1, None, 1),
        (1, 1, 1),
        (2, None, 1),
        (2, 1, 1),
        (2, 2, 1),
        (2, 3, 1),
        (3, 2, 1),
        (1_6, 1_6, 1_6),
        (1_6, 1_7, 1_6),
        (1_7, 1_6, 1_6),
    ]						,    )
def 				_lowerCAmelCase  (			lowerCamelCase_     :					Any						,    lowerCamelCase_     :					Tuple						,    lowerCamelCase_     :					List[str]			):





			with patch('''datasets.utils.py_utils._single_map_nested'''			) as mock_single_map_nested, patch(
			    '''datasets.parallel.parallel.Pool'''			) as mock_multiprocessing_pool:
						__lowercase			  =							{f"{i}": i for i in range(UpperCAmelCase_			)}
						__lowercase			  =							map_nested(lambda lowerCamelCase_			: x + 1_0						,    UpperCAmelCase_						,    num_proc=UpperCAmelCase_						,    parallel_min_length=1_6			)
						if expected_num_proc == 1:
									assert mock_single_map_nested.called
									assert not mock_multiprocessing_pool.called
						else:
									assert not mock_single_map_nested.called
									assert mock_multiprocessing_pool.called
									assert mock_multiprocessing_pool.call_args[0][0] == expected_num_proc






class 							__lowercase    ( lowerCAmelCase__  ):







	'''simple docstring'''

	@require_tf
	def   _UpperCAmelCase       (self	)							->							Optional[int]:







				'''simple docstring'''





				import tensorflow as tf
				from tensorflow.keras import layers

				__lowercase			  =							layers.Dense(2	)

				def gen_random_output():
							__lowercase			  =							tf.random.uniform((1, 3)	)
							return model(_A	).numpy()

				with temp_seed(42       ,set_tensorflow=_A	):
							__lowercase			  =							gen_random_output()
				with temp_seed(42       ,set_tensorflow=_A	):
							__lowercase			  =							gen_random_output()
				__lowercase			  =							gen_random_output()

				np.testing.assert_equal(_A       ,_A	)
				self.assertGreater(np.abs(outa - outa	).sum()       ,0	)


	@require_torch
	def   _UpperCAmelCase       (self	)							->							List[Any]:







				'''simple docstring'''





				import torch

				def gen_random_output():
							__lowercase			  =							torch.nn.Linear(3       ,2	)
							__lowercase			  =							torch.rand(1       ,3	)
							return model(_A	).detach().numpy()

				with temp_seed(42       ,set_pytorch=_A	):
							__lowercase			  =							gen_random_output()
				with temp_seed(42       ,set_pytorch=_A	):
							__lowercase			  =							gen_random_output()
				__lowercase			  =							gen_random_output()

				np.testing.assert_equal(_A       ,_A	)
				self.assertGreater(np.abs(outa - outa	).sum()       ,0	)


	def   _UpperCAmelCase       (self	)							->							Optional[Any]:







				'''simple docstring'''





				def gen_random_output():
							return np.random.rand(1       ,3	)

				with temp_seed(42	):
							__lowercase			  =							gen_random_output()
				with temp_seed(42	):
							__lowercase			  =							gen_random_output()
				__lowercase			  =							gen_random_output()

				np.testing.assert_equal(_A       ,_A	)
				self.assertGreater(np.abs(outa - outa	).sum()       ,0	)





@pytest.mark.parametrize('''input_data'''						,    [{}]			)
def 				_lowerCAmelCase  (			lowerCamelCase_     :					Optional[Any]			):
			__lowercase			  =							NestedDataStructure(UpperCAmelCase_			).data
			assert output_data == input_data




@pytest.mark.parametrize(
    '''data, expected_output'''						,    [
        ({}, []),
        ([], []),
        ('''foo''', ['''foo''']),
        (['''foo''', '''bar'''], ['''foo''', '''bar''']),
        ([['''foo''', '''bar''']], ['''foo''', '''bar''']),
        ([[['''foo'''], ['''bar''']]], ['''foo''', '''bar''']),
        ([[['''foo'''], '''bar''']], ['''foo''', '''bar''']),
        ({'''a''': 1, '''b''': 2}, [1, 2]),
        ({'''a''': [1, 2], '''b''': [3, 4]}, [1, 2, 3, 4]),
        ({'''a''': [[1, 2]], '''b''': [[3, 4]]}, [1, 2, 3, 4]),
        ({'''a''': [[1, 2]], '''b''': [3, 4]}, [1, 2, 3, 4]),
        ({'''a''': [[[1], [2]]], '''b''': [[[3], [4]]]}, [1, 2, 3, 4]),
        ({'''a''': [[[1], [2]]], '''b''': [[3, 4]]}, [1, 2, 3, 4]),
        ({'''a''': [[[1], [2]]], '''b''': [3, 4]}, [1, 2, 3, 4]),
        ({'''a''': [[[1], [2]]], '''b''': [3, [4]]}, [1, 2, 3, 4]),
        ({'''a''': {'''1''': 1}, '''b''': 2}, [1, 2]),
        ({'''a''': {'''1''': [1]}, '''b''': 2}, [1, 2]),
        ({'''a''': {'''1''': [1]}, '''b''': [2]}, [1, 2]),
    ]						,    )
def 				_lowerCAmelCase  (			lowerCamelCase_     :					str						,    lowerCamelCase_     :					str			):
			__lowercase			  =							NestedDataStructure(UpperCAmelCase_			).flatten()
			assert output == expected_output




def 				_lowerCAmelCase  (			):
			__lowercase			  =							A(x=1						,    y='''foobar'''			)
			__lowercase			  =							{'x': 1, 'y': 'foobar'}
			assert asdict(UpperCAmelCase_			) == expected_output

			__lowercase			  =							{'a': {'b': A(x=1_0						,    y='''foo'''			)}, 'c': [A(x=2_0						,    y='''bar'''			)]}
			__lowercase			  =							{'a': {'b': {'x': 1_0, 'y': 'foo'}}, 'c': [{'x': 2_0, 'y': 'bar'}]}
			assert asdict(UpperCAmelCase_			) == expected_output

			with pytest.raises(UpperCAmelCase_			):
						asdict([1, A(x=1_0						,    y='''foo'''			)]			)




def 				_lowerCAmelCase  (			lowerCamelCase_     :					str			):
			return text.split()




def 				_lowerCAmelCase  (			lowerCamelCase_     :					Dict			):
			yield (time.time(), content)
			time.sleep(2			)
			yield (time.time(), content)




def 				_lowerCAmelCase  (			):





			with Pool(2			) as pool:
						__lowercase			  =							list(iflatmap_unordered(UpperCAmelCase_						,    _split_text						,    kwargs_iterable=[{'''text''': '''hello there'''}] * 1_0			)			)
						assert out.count('''hello'''			) == 1_0
						assert out.count('''there'''			) == 1_0
						assert len(UpperCAmelCase_			) == 2_0

			# check multiprocess from pathos (uses dill for pickling)
			with multiprocess.Pool(2			) as pool:
						__lowercase			  =							list(iflatmap_unordered(UpperCAmelCase_						,    _split_text						,    kwargs_iterable=[{'''text''': '''hello there'''}] * 1_0			)			)
						assert out.count('''hello'''			) == 1_0
						assert out.count('''there'''			) == 1_0
						assert len(UpperCAmelCase_			) == 2_0

			# check that we get items as fast as possible
			with Pool(2			) as pool:
						__lowercase			  =							[]
						for yield_time, content in iflatmap_unordered(
						    UpperCAmelCase_						,    _aseconds_generator_of_aitems_with_timing						,    kwargs_iterable=[{'''content''': '''a'''}, {'''content''': '''b'''}]			):
									assert yield_time < time.time() + 0.1, "we should each item directly after it was yielded"
									out.append(UpperCAmelCase_			)
						assert out.count('''a'''			) == 2
						assert out.count('''b'''			) == 2
						assert len(UpperCAmelCase_			) == 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


_SCREAMING_SNAKE_CASE   =      {
    '''169M''': 1_2,
    '''430M''': 2_4,
    '''1B5''': 2_4,
    '''3B''': 3_2,
    '''7B''': 3_2,
    '''14B''': 4_0,
}

_SCREAMING_SNAKE_CASE   =      {
    '''169M''': 7_6_8,
    '''430M''': 1_0_2_4,
    '''1B5''': 2_0_4_8,
    '''3B''': 2_5_6_0,
    '''7B''': 4_0_9_6,
    '''14B''': 5_1_2_0,
}




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

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

						__lowercase			  =							weight
			return state_dict




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

			# 2. Build the config
			__lowercase			  =							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:
												__lowercase			  =							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}."			)

			__lowercase			  =							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
			__lowercase			  =							hf_hub_download(lowerCamelCase_						,    lowerCamelCase_			)
			__lowercase			  =							torch.load(lowerCamelCase_						,    map_location='''cpu'''			)
			__lowercase			  =							convert_state_dict(lowerCamelCase_			)

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

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

						del state_dict
						del shards
						gc.collect()

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


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

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





