






"""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					:				Optional[Any]					    =						{
    '''169M''': 12,
    '''430M''': 24,
    '''1B5''': 24,
    '''3B''': 32,
    '''7B''': 32,
    '''14B''': 40,
}

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







def      lowerCamelCase__       (							_lowerCamelCase       :		Dict )      ->				List[str]:
  lowerCamelCase_         =						list(state_dict.keys() )
  for name in state_dict_keys:
    lowerCamelCase_         =						state_dict.pop(_lowerCamelCase )
    # 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'		,				_lowerCamelCase )
    # ffn -> feed_forward
    lowerCamelCase_         =						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' ):
      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      lowerCamelCase__       (							_lowerCamelCase       :		Dict		,				_lowerCamelCase       :		Optional[Any]		,				_lowerCamelCase       :		int		,				_lowerCamelCase       :		Dict=None		,				_lowerCamelCase       :		Any=None		,				_lowerCamelCase       :		Optional[Any]=False		,				_lowerCamelCase       :		Any=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_         =						50277
    lowerCamelCase_         =						AutoTokenizer.from_pretrained('EleutherAI/gpt-neox-20b' )
  else:
    lowerCamelCase_         =						PreTrainedTokenizerFast(tokenizer_file=_lowerCamelCase )
    lowerCamelCase_         =						len(_lowerCamelCase )
  tokenizer.save_pretrained(_lowerCamelCase )

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

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

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

    del state_dict
    del shards
    gc.collect()

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


if __name__ == "__main__":
  _SCREAMING_SNAKE_CASE					:				List[str]					    =						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					:				Dict					    =						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,
  )












"""simple docstring"""





import unittest

from transformers import load_tool

from .test_tools_common import ToolTesterMixin


_SCREAMING_SNAKE_CASE					:				Optional[int]					    =						'''
Hugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf originally as a company that developed a chatbot app targeted at teenagers.[2] After open-sourcing the model behind the chatbot, the company pivoted to focus on being a platform for machine learning.

In March 2021, Hugging Face raised $40 million in a Series B funding round.[3]

On April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model.[4] In 2022, the workshop concluded with the announcement of BLOOM, a multilingual large language model with 176 billion parameters.[5]
'''






class      a						(					unittest.TestCase  ,      __snake_case		):






    def    UpperCamelCase       (       self     :							int					)     ->		Optional[int]:
      lowerCamelCase_         =						load_tool('text-question-answering'					)
      self.tool.setup()
      lowerCamelCase_         =						load_tool('text-question-answering'       ,						remote=__SCREAMING_SNAKE_CASE					)






    def    UpperCamelCase       (       self     :							Dict					)     ->		Tuple:
      lowerCamelCase_         =						self.tool(__SCREAMING_SNAKE_CASE       ,						'What did Hugging Face do in April 2021?'					)
      self.assertEqual(__SCREAMING_SNAKE_CASE       ,						'launched the BigScience Research Workshop'					)






    def    UpperCamelCase       (       self     :							List[Any]					)     ->		Optional[int]:
      lowerCamelCase_         =						self.remote_tool(__SCREAMING_SNAKE_CASE       ,						'What did Hugging Face do in April 2021?'					)
      self.assertEqual(__SCREAMING_SNAKE_CASE       ,						'launched the BigScience Research Workshop'					)






    def    UpperCamelCase       (       self     :							List[str]					)     ->		Optional[int]:
      lowerCamelCase_         =						self.tool(text=__SCREAMING_SNAKE_CASE       ,						question='What did Hugging Face do in April 2021?'					)
      self.assertEqual(__SCREAMING_SNAKE_CASE       ,						'launched the BigScience Research Workshop'					)






    def    UpperCamelCase       (       self     :							Union[str, Any]					)     ->		int:
      lowerCamelCase_         =						self.remote_tool(text=__SCREAMING_SNAKE_CASE       ,						question='What did Hugging Face do in April 2021?'					)
      self.assertEqual(__SCREAMING_SNAKE_CASE       ,						'launched the BigScience Research Workshop'					)




