




'''simple docstring'''





from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
    is_vision_available,
)


UpperCamelCase_						:      Dict       =						{
    '''configuration_clip''': [
        '''CLIP_PRETRAINED_CONFIG_ARCHIVE_MAP''',
        '''CLIPConfig''',
        '''CLIPOnnxConfig''',
        '''CLIPTextConfig''',
        '''CLIPVisionConfig''',
    ],
    '''processing_clip''': ['''CLIPProcessor'''],
    '''tokenization_clip''': ['''CLIPTokenizer'''],
}

try:
   if not is_tokenizers_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   UpperCamelCase_						:      Optional[int]       =						['''CLIPTokenizerFast''']

try:
   if not is_vision_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   UpperCamelCase_						:      List[str]       =						['''CLIPFeatureExtractor''']
   UpperCamelCase_						:      str       =						['''CLIPImageProcessor''']

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   UpperCamelCase_						:      Union[str, Any]       =						[
       '''CLIP_PRETRAINED_MODEL_ARCHIVE_LIST''',
       '''CLIPModel''',
       '''CLIPPreTrainedModel''',
       '''CLIPTextModel''',
       '''CLIPTextModelWithProjection''',
       '''CLIPVisionModel''',
       '''CLIPVisionModelWithProjection''',
   ]

try:
   if not is_tf_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   UpperCamelCase_						:      List[str]       =						[
       '''TF_CLIP_PRETRAINED_MODEL_ARCHIVE_LIST''',
       '''TFCLIPModel''',
       '''TFCLIPPreTrainedModel''',
       '''TFCLIPTextModel''',
       '''TFCLIPVisionModel''',
   ]

try:
   if not is_flax_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   UpperCamelCase_						:      Dict       =						[
       '''FlaxCLIPModel''',
       '''FlaxCLIPPreTrainedModel''',
       '''FlaxCLIPTextModel''',
       '''FlaxCLIPTextPreTrainedModel''',
       '''FlaxCLIPVisionModel''',
       '''FlaxCLIPVisionPreTrainedModel''',
   ]


if TYPE_CHECKING:
   from .configuration_clip import (
       CLIP_PRETRAINED_CONFIG_ARCHIVE_MAP,
       CLIPConfig,
       CLIPOnnxConfig,
       CLIPTextConfig,
       CLIPVisionConfig,
   )
   from .processing_clip import CLIPProcessor
   from .tokenization_clip import CLIPTokenizer

   try:
      if not is_tokenizers_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .tokenization_clip_fast import CLIPTokenizerFast

   try:
      if not is_vision_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .feature_extraction_clip import CLIPFeatureExtractor
      from .image_processing_clip import CLIPImageProcessor

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_clip import (
          CLIP_PRETRAINED_MODEL_ARCHIVE_LIST,
          CLIPModel,
          CLIPPreTrainedModel,
          CLIPTextModel,
          CLIPTextModelWithProjection,
          CLIPVisionModel,
          CLIPVisionModelWithProjection,
      )

   try:
      if not is_tf_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_tf_clip import (
          TF_CLIP_PRETRAINED_MODEL_ARCHIVE_LIST,
          TFCLIPModel,
          TFCLIPPreTrainedModel,
          TFCLIPTextModel,
          TFCLIPVisionModel,
      )

   try:
      if not is_flax_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_flax_clip import (
          FlaxCLIPModel,
          FlaxCLIPPreTrainedModel,
          FlaxCLIPTextModel,
          FlaxCLIPTextPreTrainedModel,
          FlaxCLIPVisionModel,
          FlaxCLIPVisionPreTrainedModel,
      )


else:
   import sys

   UpperCamelCase_						:      Tuple       =						_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)








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


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

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





def      __a     (       _UpperCamelCase:   str   )						->			Any:

     """simple docstring"""




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

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

          _snake_case							=				weight
     return state_dict





def      __a     (       _UpperCamelCase:   Any ,						_UpperCamelCase:   List[Any] ,						_UpperCamelCase:   List[Any] ,						_UpperCamelCase:   str=None ,						_UpperCamelCase:   Optional[Any]=None ,						_UpperCamelCase:   List[str]=False ,						_UpperCamelCase:   Dict=None   )						->			Dict:

     """simple docstring"""




     if tokenizer_file is None:
          print("No `--tokenizer_file` provided, we will use the default tokenizer."   )
          _snake_case							=				50_277
          _snake_case							=				AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b"   )
     else:
          _snake_case							=				PreTrainedTokenizerFast(tokenizer_file=_UpperCamelCase   )
          _snake_case							=				len(_UpperCamelCase   )
     tokenizer.save_pretrained(_UpperCamelCase   )

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

     _snake_case							=				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
     _snake_case							=				hf_hub_download(_UpperCamelCase ,						_UpperCamelCase   )
     _snake_case							=				torch.load(_UpperCamelCase ,						map_location="cpu"   )
     _snake_case							=				convert_state_dict(_UpperCamelCase   )

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

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

          del state_dict
          del shards
          gc.collect()

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


if __name__ == "__main__":
   UpperCamelCase_						:      Any       =						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_						:      int       =						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,
   )


