



"""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 	_SCREAMING_SNAKE_CASE  (__lowerCAmelCase							)       ->  int:  # picklable for multiprocessing







  '''simple docstring'''


  return x.sum()


def 	_SCREAMING_SNAKE_CASE  (__lowerCAmelCase							)       ->  Tuple:  # picklable for multiprocessing







  '''simple docstring'''


  return i + 1


@dataclass
class       SCREAMING_SNAKE_CASE__       :
    lowercase__          =      42
    lowercase__          =      42







class       SCREAMING_SNAKE_CASE__       (							snake_case_      ):

    def     _UpperCAmelCase		( self	: Any):







      """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(UpperCAmelCase__       ,     UpperCAmelCase__)       ,     UpperCAmelCase__)
      self.assertEqual(map_nested(UpperCAmelCase__       ,     UpperCAmelCase__)       ,     UpperCAmelCase__)
      self.assertEqual(map_nested(UpperCAmelCase__       ,     UpperCAmelCase__)       ,     UpperCAmelCase__)
      self.assertEqual(map_nested(UpperCAmelCase__       ,     UpperCAmelCase__)       ,     UpperCAmelCase__)
      self.assertEqual(map_nested(UpperCAmelCase__       ,     UpperCAmelCase__)       ,     UpperCAmelCase__)
      self.assertEqual(map_nested(UpperCAmelCase__       ,     UpperCAmelCase__)       ,     UpperCAmelCase__)
      self.assertEqual(map_nested(UpperCAmelCase__       ,     UpperCAmelCase__)       ,     UpperCAmelCase__)
      self.assertEqual(map_nested(UpperCAmelCase__       ,     UpperCAmelCase__)       ,     UpperCAmelCase__)

      lowercase_	    =       2
      self.assertEqual(map_nested(UpperCAmelCase__       ,     UpperCAmelCase__       ,     num_proc=UpperCAmelCase__)       ,     UpperCAmelCase__)
      self.assertEqual(map_nested(UpperCAmelCase__       ,     UpperCAmelCase__       ,     num_proc=UpperCAmelCase__)       ,     UpperCAmelCase__)
      self.assertEqual(map_nested(UpperCAmelCase__       ,     UpperCAmelCase__       ,     num_proc=UpperCAmelCase__)       ,     UpperCAmelCase__)
      self.assertEqual(map_nested(UpperCAmelCase__       ,     UpperCAmelCase__       ,     num_proc=UpperCAmelCase__)       ,     UpperCAmelCase__)
      self.assertEqual(map_nested(UpperCAmelCase__       ,     UpperCAmelCase__       ,     num_proc=UpperCAmelCase__)       ,     UpperCAmelCase__)
      self.assertEqual(map_nested(UpperCAmelCase__       ,     UpperCAmelCase__       ,     num_proc=UpperCAmelCase__)       ,     UpperCAmelCase__)
      self.assertEqual(map_nested(UpperCAmelCase__       ,     UpperCAmelCase__       ,     num_proc=UpperCAmelCase__)       ,     UpperCAmelCase__)
      self.assertEqual(map_nested(UpperCAmelCase__       ,     UpperCAmelCase__       ,     num_proc=UpperCAmelCase__)       ,     UpperCAmelCase__)

      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(UpperCAmelCase__),
          """b""": np.zeros(3).astype(UpperCAmelCase__),
          """c""": np.ones(2).astype(UpperCAmelCase__),
      }
      self.assertEqual(map_nested(UpperCAmelCase__       ,     UpperCAmelCase__       ,     map_numpy=UpperCAmelCase__)       ,     UpperCAmelCase__)
      self.assertEqual(
          {k: v.tolist() for k, v in map_nested(UpperCAmelCase__       ,     UpperCAmelCase__       ,     map_numpy=UpperCAmelCase__).items()}       ,     {k: v.tolist() for k, v in expected_map_nested_sna_int.items()}       ,     )
      self.assertEqual(map_nested(UpperCAmelCase__       ,     UpperCAmelCase__       ,     map_numpy=UpperCAmelCase__       ,     num_proc=UpperCAmelCase__)       ,     UpperCAmelCase__)
      self.assertEqual(
          {k: v.tolist() for k, v in map_nested(UpperCAmelCase__       ,     UpperCAmelCase__       ,     map_numpy=UpperCAmelCase__       ,     num_proc=UpperCAmelCase__).items()}       ,     {k: v.tolist() for k, v in expected_map_nested_sna_int.items()}       ,     )
      with self.assertRaises(UpperCAmelCase__):  # can't pickle a local lambda
        map_nested(lambda lowerCAmelCase_: x + 1       ,     UpperCAmelCase__       ,     num_proc=UpperCAmelCase__)

    def     _UpperCAmelCase		( self	: Optional[int]):







      """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(UpperCAmelCase__       ,     UpperCAmelCase__       ,     UpperCAmelCase__))       ,     UpperCAmelCase__)




    def     _UpperCAmelCase		( self	: Dict):







      """simple docstring"""







      class       SCREAMING_SNAKE_CASE__       :
          lowercase__          =      '''bar'''







      lowercase_	    =       Foo()
      self.assertEqual(foo.my_attr       ,     """bar""")
      with temporary_assignment(UpperCAmelCase__       ,     """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),
        (16, 16, 16),
        (16, 17, 16),
        (17, 16, 16),
    ]       ,     )
def 	_SCREAMING_SNAKE_CASE  (__lowerCAmelCase       ,     __lowerCAmelCase       ,     __lowerCAmelCase							)       ->  Dict:







  '''simple docstring'''


  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 + 10       ,     UpperCamelCase__       ,     num_proc=UpperCamelCase__       ,     parallel_min_length=16							)
    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       SCREAMING_SNAKE_CASE__       (							snake_case_      ):

    @require_tf
    def     _UpperCAmelCase		( self	: List[Any]):







      """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(UpperCAmelCase__).numpy()

      with temp_seed(4_2       ,     set_tensorflow=UpperCAmelCase__):
        lowercase_	    =       gen_random_output()
      with temp_seed(4_2       ,     set_tensorflow=UpperCAmelCase__):
        lowercase_	    =       gen_random_output()
      lowercase_	    =       gen_random_output()

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

    @require_torch
    def     _UpperCAmelCase		( self	: Optional[int]):







      """simple docstring"""







      import torch

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

      with temp_seed(4_2       ,     set_pytorch=UpperCAmelCase__):
        lowercase_	    =       gen_random_output()
      with temp_seed(4_2       ,     set_pytorch=UpperCAmelCase__):
        lowercase_	    =       gen_random_output()
      lowercase_	    =       gen_random_output()

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




    def     _UpperCAmelCase		( self	: str):







      """simple docstring"""







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

      with temp_seed(4_2):
        lowercase_	    =       gen_random_output()
      with temp_seed(4_2):
        lowercase_	    =       gen_random_output()
      lowercase_	    =       gen_random_output()

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







@pytest.mark.parametrize("""input_data"""       ,     [{}]							)
def 	_SCREAMING_SNAKE_CASE  (__lowerCAmelCase							)       ->  Optional[int]:







  '''simple docstring'''


  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 	_SCREAMING_SNAKE_CASE  (__lowerCAmelCase       ,     __lowerCAmelCase							)       ->  Dict:







  '''simple docstring'''


  lowercase_	    =       NestedDataStructure(UpperCamelCase__							).flatten()
  assert output == expected_output


def 	_SCREAMING_SNAKE_CASE  ()       ->  Optional[Any]:







  '''simple docstring'''


  lowercase_	    =       A(x=1       ,     y="""foobar"""							)
  lowercase_	    =       {"""x""": 1, """y""": """foobar"""}
  assert asdict(UpperCamelCase__							) == expected_output

  lowercase_	    =       {"""a""": {"""b""": A(x=10       ,     y="""foo"""							)}, """c""": [A(x=20       ,     y="""bar"""							)]}
  lowercase_	    =       {"""a""": {"""b""": {"""x""": 10, """y""": """foo"""}}, """c""": [{"""x""": 20, """y""": """bar"""}]}
  assert asdict(UpperCamelCase__							) == expected_output

  with pytest.raises(UpperCamelCase__							):
    asdict([1, A(x=10       ,     y="""foo"""							)]							)


def 	_SCREAMING_SNAKE_CASE  (__lowerCAmelCase							)       ->  Union[str, Any]:







  '''simple docstring'''


  return text.split()


def 	_SCREAMING_SNAKE_CASE  (__lowerCAmelCase							)       ->  Tuple:







  '''simple docstring'''


  yield (time.time(), content)
  time.sleep(2							)
  yield (time.time(), content)


def 	_SCREAMING_SNAKE_CASE  ()       ->  Tuple:







  '''simple docstring'''


  with Pool(2							) as pool:
    lowercase_	    =       list(iflatmap_unordered(UpperCamelCase__       ,     _split_text       ,     kwargs_iterable=[{"""text""": """hello there"""}] * 10							)							)
    assert out.count("""hello"""							) == 10
    assert out.count("""there"""							) == 10
    assert len(UpperCamelCase__							) == 20

  # 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"""}] * 10							)							)
    assert out.count("""hello"""							) == 10
    assert out.count("""there"""							) == 10
    assert len(UpperCamelCase__							) == 20

  # 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


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

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

def 			lowerCAmelCase__     (     UpperCamelCase__				):



  '''simple docstring'''
  _a  :       int             =	list(state_dict.keys()				)
  for name in state_dict_keys:
    _a  :       str             =	state_dict.pop(UpperCamelCase__				)
    # emb -> embedding
    if name.startswith("""emb."""				):
      _a  :       Dict             =	name.replace("""emb."""     ,   """embeddings."""				)
    # ln_0 -> pre_ln (only present at block 0)
    if name.startswith("""blocks.0.ln0"""				):
      _a  :       Dict             =	name.replace("""blocks.0.ln0"""     ,   """blocks.0.pre_ln"""				)
    # att -> attention
    _a  :       Any             =	re.sub(R"""blocks\.(\d+)\.att"""     ,   R"""blocks.\1.attention"""     ,   UpperCamelCase__				)
    # ffn -> feed_forward
    _a  :       int             =	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"""				):
      _a  :       List[str]             =	name.replace(""".time_mix_k"""     ,   """.time_mix_key"""				)
    # time_mix_v -> time_mix_value and reshape
    if name.endswith(""".time_mix_v"""				):
      _a  :       Tuple             =	name.replace(""".time_mix_v"""     ,   """.time_mix_value"""				)
    # time_mix_r -> time_mix_key and reshape
    if name.endswith(""".time_mix_r"""				):
      _a  :       Dict             =	name.replace(""".time_mix_r"""     ,   """.time_mix_receptance"""				)

    if name != "head.weight":
      _a  :       Optional[int]             =	"""rwkv.""" + name

    _a  :       Any             =	weight
  return state_dict

def 			lowerCAmelCase__     (     UpperCamelCase__     ,   UpperCamelCase__     ,   UpperCamelCase__     ,   UpperCamelCase__=None     ,   UpperCamelCase__=None     ,   UpperCamelCase__=False     ,   UpperCamelCase__=None				):



  '''simple docstring'''
  # 1. If possible, build the tokenizer.
  if tokenizer_file is None:
    print("""No `--tokenizer_file` provided, we will use the default tokenizer."""				)
    _a  :       Tuple             =	5_0_2_7_7
    _a  :       str             =	AutoTokenizer.from_pretrained("""EleutherAI/gpt-neox-20b"""				)
  else:
    _a  :       int             =	PreTrainedTokenizerFast(tokenizer_file=UpperCamelCase__				)
    _a  :       int             =	len(UpperCamelCase__				)
  tokenizer.save_pretrained(UpperCamelCase__				)

  # 2. Build the config
  _a  :       Optional[Any]             =	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:
        _a  :       Tuple             =	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}."""				)

  _a  :       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
  _a  :       str             =	hf_hub_download(UpperCamelCase__     ,   UpperCamelCase__				)
  _a  :       int             =	torch.load(UpperCamelCase__     ,   map_location="""cpu"""				)
  _a  :       List[str]             =	convert_state_dict(UpperCamelCase__				)

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

  if index is not None:
    _a  :       Optional[Any]             =	os.path.join(UpperCamelCase__     ,   UpperCamelCase__				)
    # Save the index as well
    with open(UpperCamelCase__     ,   """w"""     ,   encoding="""utf-8"""				) as f:
      _a  :       Dict             =	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."""				)
    _a  :       List[Any]             =	list(shards.keys()				)

    del state_dict
    del shards
    gc.collect()

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


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

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

