"""Script to convert GPT2 models from llm.c binary format to Hugging Face It can optinally upload to your account on Hugging Face if you have the CLI:  pip install -U "huggingface_hub[cli]"  huggingface-cli login Export to a local HF model:  python export_hf.py --input input_file.bin --output output_dir Export to a local HF model and also push to your account on Hugging Face:  python export_hf.py --input input_file.bin --output output_dir --push true""" import numpy as npimport torchimport argparse, sysfrom transformers import GPT2Config, GPT2Tokenizer, GPT2LMHeadModel # -----------------------------------------------------------------------------# Tensor functions for both bfloat16 (from int16) and normal float32# Both return float32 tensors def tensor_bf16(data_int16, transpose=False):    if transpose:        data_int16 = data_int16.transpose(1,0)    return torch.tensor(data_int16).view(torch.bfloat16).to(torch.float32) def tensor_fp32(data_float32, transpose=False):    if transpose:        data_float32 = data_float32.transpose(1,0)    return torch.tensor(data_float32).view(torch.float32) # -----------------------------------------------------------------------------# Main conversion function def convert(filepath, output, push_to_hub=False, out_dtype="bfloat16"):    print(f"Converting model {filepath} to {output} in {out_dtype} format and pushing to Hugging Face: {push_to_hub}")     f = open(filepath, 'rb')    # Read in our header, checking the magic number and version    # version 3 = fp32, padded vocab    # version 5 = bf16, padded vocab    model_header = np.frombuffer(f.read(256*4), dtype=np.int32)    if model_header[0] != 20240326:        print("ERROR: magic number mismatch in the data .bin file!")        exit(1)    version = model_header[1]    if not version in [3, 5]:        print("Bad version in model file")        exit(1)     # Load in our model parameters    maxT = model_header[2].item() # max sequence length    V = model_header[3].item() # vocab size    L =  model_header[4].item() # num layers    H = model_header[5].item() # num heads    C = model_header[6].item() # channels    Vp = model_header[7].item() # padded vocab size     print(f"{version=}, {maxT=}, {V=}, {Vp=}, {L=}, {H=}, {C=}")     # Define the shapes of our parameters    shapes = {        'wte': (Vp, C),        'wpe': (maxT, C),        'ln1w': (L, C),        'ln1b': (L, C),        'qkvw': (L, 3 * C, C),        'qkvb': (L, 3 * C),        'attprojw': (L, C, C),        'attprojb': (L, C),        'ln2w': (L, C),        'ln2b': (L, C),        'fcw': (L, 4 * C, C),        'fcb': (L, 4 * C),        'fcprojw': (L, C, 4 * C),        'fcprojb': (L, C),        'lnfw': (C,),        'lnfb': (C,),    }     # Load in our weights given our parameter shapes    dtype = np.float32 if version == 3 else np.int16    w = {}    for key, shape in shapes.items():        num_elements = np.prod(shape)        data = np.frombuffer(f.read(num_elements * np.dtype(dtype).itemsize), dtype=dtype)        w[key] = data.reshape(shape)        # The binary file saves the padded vocab - drop the padding back to GPT2 size        if shape[0] == Vp:            w[key] = w[key].reshape(shape)[:(V-Vp), :]    # Ensure the file is fully read and then close    assert f.read() == b''    f.close()     # Map to our model dict, the tensors at this stage are always fp32    mk_tensor = {        3 : tensor_fp32,        5 : tensor_bf16,    }[version]    model_dict = {}    model_dict['transformer.wte.weight'] = mk_tensor(w['wte'])    model_dict['transformer.wpe.weight'] = mk_tensor(w['wpe'])    model_dict['lm_head.weight'] = model_dict['transformer.wte.weight'] # Tie weights    for i in range(L):        model_dict[f'transformer.h.{i}.ln_1.weight'] = mk_tensor(w['ln1w'][i])        model_dict[f'transformer.h.{i}.ln_1.bias'] = mk_tensor(w['ln1b'][i])        model_dict[f'transformer.h.{i}.attn.c_attn.weight'] = mk_tensor(w['qkvw'][i], True)        model_dict[f'transformer.h.{i}.attn.c_attn.bias'] = mk_tensor(w['qkvb'][i])        model_dict[f'transformer.h.{i}.attn.c_proj.weight'] = mk_tensor(w['attprojw'][i], True)        model_dict[f'transformer.h.{i}.attn.c_proj.bias'] = mk_tensor(w['attprojb'][i])        model_dict[f'transformer.h.{i}.ln_2.weight'] = mk_tensor(w['ln2w'][i])        model_dict[f'transformer.h.{i}.ln_2.bias'] = mk_tensor(w['ln2b'][i])        model_dict[f'transformer.h.{i}.mlp.c_fc.weight'] = mk_tensor(w['fcw'][i], True)        model_dict[f'transformer.h.{i}.mlp.c_fc.bias'] = mk_tensor(w['fcb'][i])        model_dict[f'transformer.h.{i}.mlp.c_proj.weight'] = mk_tensor(w['fcprojw'][i], True)        model_dict[f'transformer.h.{i}.mlp.c_proj.bias'] = mk_tensor(w['fcprojb'][i])    model_dict['transformer.ln_f.weight'] = mk_tensor(w['lnfw'])    model_dict['transformer.ln_f.bias'] = mk_tensor(w['lnfb'])     # Create a GPT-2 model instance, in the requested dtype    config = GPT2Config(vocab_size = V,                        n_positions = maxT,                        n_ctx = maxT,                        n_embd = C,                        n_layer = L,                        n_head = H)    model = GPT2LMHeadModel(config)    if out_dtype == "bfloat16":        model = model.to(torch.bfloat16)     # Set the model dict and save    model.load_state_dict(model_dict)    model.save_pretrained(output, max_shard_size="5GB", safe_serialization=True)     # Copy over a standard gpt2 tokenizer    tokenizer = GPT2Tokenizer.from_pretrained("gpt2")    tokenizer.save_pretrained(output)     if push_to_hub:        print(f"Uploading {output} to Hugging Face")        model.push_to_hub(output)        tokenizer.push_to_hub(output) def spin(output):    print("Taking the exported model for a spin...")    print('-'*80)    from transformers import AutoModelForCausalLM, AutoTokenizer    tokenizer = AutoTokenizer.from_pretrained(output)    model = AutoModelForCausalLM.from_pretrained(output, attn_implementation="flash_attention_2", torch_dtype=torch.bfloat16, device_map='cuda')    model.eval()    tokens = tokenizer.encode("During photosynthesis in green plants", return_tensors="pt")    tokens = tokens.to('cuda')    output = model.generate(tokens, max_new_tokens=64, repetition_penalty=1.3)    samples = tokenizer.batch_decode(output)    for sample in samples:        print('-'*30)        print(sample) # ----------------------------------------------------------------------------- if __name__== '__main__':    parser=argparse.ArgumentParser()    parser.add_argument("--input", "-i", help="The name of the llm.c model.bin file", type=str, required=True)    parser.add_argument("--output","-o",  help="The Hugging Face output model directory", type=str, required=True)    parser.add_argument("--dtype", "-d", help="Output as either float32 or bfloat16 (default)", type=str, default="bfloat16")    parser.add_argument("--push", "-p", help="Push the model to your Hugging Face account", type=bool, default=False)    parser.add_argument("--spin", "-s", help="Take the model for a spin at the end?", type=bool, default=True)    args = parser.parse_args()    convert(args.input, args.output, args.push, args.dtype)    if args.spin:        spin(args.output)