codingutf8 2023 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license convert a rwkv checkpoint from blinkdl to the hugging face format import argparse import gc import json import os import re import torch from huggingfacehub import hfhubdownload from transformers import automodelforcausallm autotokenizer pretrainedtokenizerfast rwkvconfig from transformers modelingutils import weightsindexname shardcheckpoint numhiddenlayersmapping 169m 12 430m 24 1b5 24 3b 32 7b 32 14b 40 hidensizemapping 169m 768 430m 1024 1b5 2048 3b 2560 7b 4096 14b 5120 def convertstatedictstatedict statedictkeys liststatedict keys for name in statedictkeys weight statedict popname emb embedding if name startswithemb name name replaceemb embeddings ln0 preln only present at block 0 if name startswithblocks 0 ln0 name name replaceblocks 0 ln0 blocks 0 preln att attention name re subrblocks d att rblocks 1 attention name ffn feedforward name re subrblocks d ffn rblocks 1 feedforward name timemixk timemixkey and reshape if name endswith timemixk name name replace timemixk timemixkey timemixv timemixvalue and reshape if name endswith timemixv name name replace timemixv timemixvalue timemixr timemixkey and reshape if name endswith timemixr name name replace timemixr timemixreceptance if name head weight name rwkv name statedictname weight return statedict def convertrmkvcheckpointtohfformat repoid checkpointfile outputdir sizenone tokenizerfilenone pushtohubfalse modelnamenone 1 if possible build the tokenizer if tokenizerfile is none printno tokenizerfile provided we will use the default tokenizer vocabsize 50277 tokenizer autotokenizer frompretrainedeleutheraigptneox20b else tokenizer pretrainedtokenizerfasttokenizerfiletokenizerfile vocabsize lentokenizer tokenizer savepretrainedoutputdir 2 build the config possiblesizes listnumhiddenlayersmapping keys if size is none try to infer size from the checkpoint name for candidate in possiblesizes if candidate in checkpointfile size candidate break if size is none raise valueerrorcould not infer the size please provide it with the size argument if size not in possiblesizes raise valueerrorfsize should be one of possiblesizes got size config rwkvconfig vocabsizevocabsize numhiddenlayersnumhiddenlayersmappingsize hiddensizehidensizemappingsize config savepretrainedoutputdir 3 download model file then convert statedict modelfile hfhubdownloadrepoid checkpointfile statedict torch loadmodelfile maplocationcpu statedict convertstatedictstatedict 4 split in shards and save shards index shardcheckpointstatedict for shardfile shard in shards items torch saveshard os path joinoutputdir shardfile if index is not none saveindexfile os path joinoutputdir weightsindexname save the index as well with opensaveindexfile w encodingutf8 as f content json dumpsindex indent2 sortkeystrue n f writecontent 5 clean up shards for some reason the file pytorch saves take the same space as the whole statedict 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 shardfiles listshards keys del statedict del shards gc collect for shardfile in shardfiles statedict torch loados path joinoutputdir shardfile torch savek v cpu clone for k v in statedict items os path joinoutputdir shardfile del statedict gc collect if pushtohub if modelname is none raise valueerrorplease provide a modelname to push the model to the hub model automodelforcausallm frompretrainedoutputdir model pushtohubmodelname maxshardsize2gb tokenizer pushtohubmodelname if name main parser argparse argumentparser required parameters parser addargument repoid defaultnone typestr requiredtrue helprepo id from which to pull the checkpoint parser addargument checkpointfile defaultnone typestr requiredtrue helpname of the checkpoint file in the repo parser addargument outputdir defaultnone typestr requiredtrue helpwhere to save the converted model parser addargument tokenizerfile defaultnone typestr helppath to the tokenizer file to use if not provided only the model is converted parser addargument size defaultnone typestr helpsize of the model will be inferred from the checkpointfile if not passed parser addargument pushtohub actionstoretrue helppush to the hub the converted model parser addargument modelname defaultnone typestr helpname of the pushed model on the hub including the username organization args parser parseargs convertrmkvcheckpointtohfformat args repoid args checkpointfile args outputdir sizeargs size tokenizerfileargs tokenizerfile pushtohubargs pushtohub modelnameargs modelname coding utf 8 2023 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license convert a rwkv checkpoint from blinkdl to the hugging face format emb embedding ln_0 pre_ln only present at block 0 att attention ffn feed_forward time_mix_k time_mix_key and reshape time_mix_v time_mix_value and reshape time_mix_r time_mix_key and reshape 1 if possible build the tokenizer 2 build the config try to infer size from the checkpoint name 3 download model file then convert state_dict 4 split in shards and save save the index as well 5 clean up shards for some reason the file pytorch saves take the same space as the whole state_dict required parameters
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


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

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


def convert_state_dict(state_dict):
    state_dict_keys = list(state_dict.keys())
    for name in state_dict_keys:
        weight = state_dict.pop(name)
        
        if name.startswith("emb."):
            name = name.replace("emb.", "embeddings.")
        
        if name.startswith("blocks.0.ln0"):
            name = name.replace("blocks.0.ln0", "blocks.0.pre_ln")
        
        name = re.sub(r"blocks\.(\d+)\.att", r"blocks.\1.attention", name)
        
        name = re.sub(r"blocks\.(\d+)\.ffn", r"blocks.\1.feed_forward", name)
        
        if name.endswith(".time_mix_k"):
            name = name.replace(".time_mix_k", ".time_mix_key")
        
        if name.endswith(".time_mix_v"):
            name = name.replace(".time_mix_v", ".time_mix_value")
        
        if name.endswith(".time_mix_r"):
            name = name.replace(".time_mix_r", ".time_mix_receptance")

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

        state_dict[name] = weight
    return state_dict


def convert_rmkv_checkpoint_to_hf_format(
    repo_id, checkpoint_file, output_dir, size=None, tokenizer_file=None, push_to_hub=False, model_name=None
):
    
    if tokenizer_file is None:
        print("No `--tokenizer_file` provided, we will use the default tokenizer.")
        vocab_size = 50277
        tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")
    else:
        tokenizer = PreTrainedTokenizerFast(tokenizer_file=tokenizer_file)
        vocab_size = len(tokenizer)
    tokenizer.save_pretrained(output_dir)

    
    possible_sizes = list(NUM_HIDDEN_LAYERS_MAPPING.keys())
    if size is None:
        
        for candidate in possible_sizes:
            if candidate in checkpoint_file:
                size = 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}.")

    config = RwkvConfig(
        vocab_size=vocab_size,
        num_hidden_layers=NUM_HIDDEN_LAYERS_MAPPING[size],
        hidden_size=HIDEN_SIZE_MAPPING[size],
    )
    config.save_pretrained(output_dir)

    
    model_file = hf_hub_download(repo_id, checkpoint_file)
    state_dict = torch.load(model_file, map_location="cpu")
    state_dict = convert_state_dict(state_dict)

    
    shards, index = shard_checkpoint(state_dict)
    for shard_file, shard in shards.items():
        torch.save(shard, os.path.join(output_dir, shard_file))

    if index is not None:
        save_index_file = os.path.join(output_dir, WEIGHTS_INDEX_NAME)
        
        with open(save_index_file, "w", encoding="utf-8") as f:
            content = json.dumps(index, indent=2, sort_keys=True) + "\n"
            f.write(content)

        
        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."
        )
        shard_files = list(shards.keys())

        del state_dict
        del shards
        gc.collect()

        for shard_file in shard_files:
            state_dict = torch.load(os.path.join(output_dir, shard_file))
            torch.save({k: v.cpu().clone() for k, v in state_dict.items()}, os.path.join(output_dir, shard_file))

    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.")
        model = AutoModelForCausalLM.from_pretrained(output_dir)
        model.push_to_hub(model_name, max_shard_size="2GB")
        tokenizer.push_to_hub(model_name)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    
    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.",
    )

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