





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


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

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


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

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

    _snake_case   :				Optional[int]														=	weight
  return state_dict


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

  # 2. Build the config
  _snake_case   :				Union[str, 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:
        _snake_case   :				Optional[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}."""				)

  _snake_case   :				Optional[Any]														=	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
  _snake_case   :				int														=	hf_hub_download(lowerCAmelCase    ,							lowerCAmelCase				)
  _snake_case   :				Dict														=	torch.load(lowerCAmelCase    ,							map_location='cpu'				)
  _snake_case   :				str														=	convert_state_dict(lowerCAmelCase				)

  # 4. Split in shards and save
  _snake_case   :				Optional[Any]														=	shard_checkpoint(lowerCAmelCase				)
  for shard_file, shard in shards.items():
    torch.save(lowerCAmelCase    ,							os.path.join(lowerCAmelCase    ,							lowerCAmelCase				)				)

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

    del state_dict
    del shards
    gc.collect()

    for shard_file in shard_files:
      _snake_case   :				Optional[Any]														=	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.'				)
    _snake_case   :				Optional[Any]														=	AutoModelForCausalLM.from_pretrained(lowerCAmelCase				)
    model.push_to_hub(lowerCAmelCase    ,							max_shard_size='2GB'				)
    tokenizer.push_to_hub(lowerCAmelCase				)


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

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











import json
from typing import List, Optional, Tuple

from tokenizers import normalizers

from ...tokenization_utils_base import BatchEncoding
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import PaddingStrategy, logging
from .tokenization_realm import RealmTokenizer


lowerCAmelCase_						   =    logging.get_logger(__name__)

lowerCAmelCase_						   =    {"""vocab_file""": """vocab.txt""", """tokenizer_file""": """tokenizer.json"""}

lowerCAmelCase_						   =    {
    """vocab_file""": {
        """google/realm-cc-news-pretrained-embedder""": (
            """https://huggingface.co/google/realm-cc-news-pretrained-embedder/resolve/main/vocab.txt"""
        ),
        """google/realm-cc-news-pretrained-encoder""": (
            """https://huggingface.co/google/realm-cc-news-pretrained-encoder/resolve/main/vocab.txt"""
        ),
        """google/realm-cc-news-pretrained-scorer""": (
            """https://huggingface.co/google/realm-cc-news-pretrained-scorer/resolve/main/vocab.txt"""
        ),
        """google/realm-cc-news-pretrained-openqa""": (
            """https://huggingface.co/google/realm-cc-news-pretrained-openqa/aresolve/main/vocab.txt"""
        ),
        """google/realm-orqa-nq-openqa""": """https://huggingface.co/google/realm-orqa-nq-openqa/resolve/main/vocab.txt""",
        """google/realm-orqa-nq-reader""": """https://huggingface.co/google/realm-orqa-nq-reader/resolve/main/vocab.txt""",
        """google/realm-orqa-wq-openqa""": """https://huggingface.co/google/realm-orqa-wq-openqa/resolve/main/vocab.txt""",
        """google/realm-orqa-wq-reader""": """https://huggingface.co/google/realm-orqa-wq-reader/resolve/main/vocab.txt""",
    },
    """tokenizer_file""": {
        """google/realm-cc-news-pretrained-embedder""": (
            """https://huggingface.co/google/realm-cc-news-pretrained-embedder/resolve/main/tokenizer.jsont"""
        ),
        """google/realm-cc-news-pretrained-encoder""": (
            """https://huggingface.co/google/realm-cc-news-pretrained-encoder/resolve/main/tokenizer.json"""
        ),
        """google/realm-cc-news-pretrained-scorer""": (
            """https://huggingface.co/google/realm-cc-news-pretrained-scorer/resolve/main/tokenizer.json"""
        ),
        """google/realm-cc-news-pretrained-openqa""": (
            """https://huggingface.co/google/realm-cc-news-pretrained-openqa/aresolve/main/tokenizer.json"""
        ),
        """google/realm-orqa-nq-openqa""": (
            """https://huggingface.co/google/realm-orqa-nq-openqa/resolve/main/tokenizer.json"""
        ),
        """google/realm-orqa-nq-reader""": (
            """https://huggingface.co/google/realm-orqa-nq-reader/resolve/main/tokenizer.json"""
        ),
        """google/realm-orqa-wq-openqa""": (
            """https://huggingface.co/google/realm-orqa-wq-openqa/resolve/main/tokenizer.json"""
        ),
        """google/realm-orqa-wq-reader""": (
            """https://huggingface.co/google/realm-orqa-wq-reader/resolve/main/tokenizer.json"""
        ),
    },
}

lowerCAmelCase_						   =    {
    """google/realm-cc-news-pretrained-embedder""": 512,
    """google/realm-cc-news-pretrained-encoder""": 512,
    """google/realm-cc-news-pretrained-scorer""": 512,
    """google/realm-cc-news-pretrained-openqa""": 512,
    """google/realm-orqa-nq-openqa""": 512,
    """google/realm-orqa-nq-reader""": 512,
    """google/realm-orqa-wq-openqa""": 512,
    """google/realm-orqa-wq-reader""": 512,
}

lowerCAmelCase_						   =    {
    """google/realm-cc-news-pretrained-embedder""": {"""do_lower_case""": True},
    """google/realm-cc-news-pretrained-encoder""": {"""do_lower_case""": True},
    """google/realm-cc-news-pretrained-scorer""": {"""do_lower_case""": True},
    """google/realm-cc-news-pretrained-openqa""": {"""do_lower_case""": True},
    """google/realm-orqa-nq-openqa""": {"""do_lower_case""": True},
    """google/realm-orqa-nq-reader""": {"""do_lower_case""": True},
    """google/realm-orqa-wq-openqa""": {"""do_lower_case""": True},
    """google/realm-orqa-wq-reader""": {"""do_lower_case""": True},
}



class        _lowerCAmelCase    (			UpperCAmelCase_							):


  '''simple docstring'''
  a_							:       Union[str, Any]      					=VOCAB_FILES_NAMES
  a_							:       List[str]      					=PRETRAINED_VOCAB_FILES_MAP
  a_							:       str      					=PRETRAINED_INIT_CONFIGURATION
  a_							:       Optional[Any]      					=PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
  a_							:       List[Any]      					=RealmTokenizer







  def __init__(      self      :					List[str]      ,				UpperCamelCase      :					Optional[int]=None      ,				UpperCamelCase      :					List[Any]=None      ,				UpperCamelCase      :					Optional[Any]=True      ,				UpperCamelCase      :					Optional[Any]="[UNK]"      ,				UpperCamelCase      :					Any="[SEP]"      ,				UpperCamelCase      :					Optional[Any]="[PAD]"      ,				UpperCamelCase      :					Optional[int]="[CLS]"      ,				UpperCamelCase      :					Optional[Any]="[MASK]"      ,				UpperCamelCase      :					Dict=True      ,				UpperCamelCase      :					Optional[int]=None      ,				**UpperCamelCase      :					Union[str, Any]      ,				):


    '''simple docstring'''





    super().__init__(
        UpperCamelCase      ,				tokenizer_file=UpperCamelCase      ,				do_lower_case=UpperCamelCase      ,				unk_token=UpperCamelCase      ,				sep_token=UpperCamelCase      ,				pad_token=UpperCamelCase      ,				cls_token=UpperCamelCase      ,				mask_token=UpperCamelCase      ,				tokenize_chinese_chars=UpperCamelCase      ,				strip_accents=UpperCamelCase      ,				**UpperCamelCase      ,				)

    _snake_case   :				int														=	json.loads(self.backend_tokenizer.normalizer.__getstate__()							)
    if (
        normalizer_state.get('lowercase'      ,				UpperCamelCase							) != do_lower_case
        or normalizer_state.get('strip_accents'      ,				UpperCamelCase							) != strip_accents
        or normalizer_state.get('handle_chinese_chars'      ,				UpperCamelCase							) != tokenize_chinese_chars
    ):
      _snake_case   :				int														=	getattr(UpperCamelCase      ,				normalizer_state.pop('type'							)							)
      _snake_case   :				List[str]														=	do_lower_case
      _snake_case   :				List[Any]														=	strip_accents
      _snake_case   :				Dict														=	tokenize_chinese_chars
      _snake_case   :				Any														=	normalizer_class(**UpperCamelCase							)

    _snake_case   :				Optional[int]														=	do_lower_case







  def UpperCamelCase_ (      self      :					Dict      ,				UpperCamelCase      :					Any      ,				**UpperCamelCase      :					Union[str, Any]							):


    '''simple docstring'''





    _snake_case   :				Tuple														=	PaddingStrategy.MAX_LENGTH

    _snake_case   :				Any														=	text
    _snake_case   :				List[str]														=	kwargs.pop('text_pair'      ,				UpperCamelCase							)
    _snake_case   :				int														=	kwargs.pop('return_tensors'      ,				UpperCamelCase							)

    _snake_case   :				Optional[int]														=	{
        'input_ids': [],
        'attention_mask': [],
        'token_type_ids': [],
    }

    for idx, candidate_text in enumerate(UpperCamelCase							):
      if batch_text_pair is not None:
        _snake_case   :				List[Any]														=	batch_text_pair[idx]
      else:
        _snake_case   :				Optional[Any]														=	None

      _snake_case   :				Optional[int]														=	super().__call__(UpperCamelCase      ,				UpperCamelCase      ,				return_tensors=UpperCamelCase      ,				**UpperCamelCase							)

      _snake_case   :				str														=	encoded_candidates.get('input_ids'							)
      _snake_case   :				Tuple														=	encoded_candidates.get('attention_mask'							)
      _snake_case   :				List[str]														=	encoded_candidates.get('token_type_ids'							)

      if encoded_input_ids is not None:
        output_data["input_ids"].append(UpperCamelCase							)
      if encoded_attention_mask is not None:
        output_data["attention_mask"].append(UpperCamelCase							)
      if encoded_token_type_ids is not None:
        output_data["token_type_ids"].append(UpperCamelCase							)

    _snake_case   :				str														=	{key: item for key, item in output_data.items() if len(UpperCamelCase							) != 0}

    return BatchEncoding(UpperCamelCase      ,				tensor_type=UpperCamelCase							)







  def UpperCamelCase_ (      self      :					str      ,				UpperCamelCase      :					Optional[Any]      ,				UpperCamelCase      :					Union[str, Any]=None							):


    '''simple docstring'''





    _snake_case   :				Dict														=	[self.cls_token_id] + token_ids_a + [self.sep_token_id]

    if token_ids_a:
      output += token_ids_a + [self.sep_token_id]

    return output







  def UpperCamelCase_ (      self      :					Union[str, Any]      ,				UpperCamelCase      :					List[int]      ,				UpperCamelCase      :					Optional[List[int]] = None							):


    '''simple docstring'''





    _snake_case   :				int														=	[self.sep_token_id]
    _snake_case   :				Dict														=	[self.cls_token_id]
    if token_ids_a is None:
      return len(cls + token_ids_a + sep							) * [0]
    return len(cls + token_ids_a + sep							) * [0] + len(token_ids_a + sep							) * [1]







  def UpperCamelCase_ (      self      :					Union[str, Any]      ,				UpperCamelCase      :					str      ,				UpperCamelCase      :					Optional[str] = None							):


    '''simple docstring'''





    _snake_case   :				Optional[Any]														=	self._tokenizer.model.save(UpperCamelCase      ,				name=UpperCamelCase							)
    return tuple(UpperCamelCase							)




