├── README.md ├── convert.py └── test.py /README.md: -------------------------------------------------------------------------------- 1 | **WARNING:** Project is abandonned and in a very early/unsuable state 2 | 3 | # :fire: Torch to ggml 4 | 5 | a python tool to convert any (hopefully) pytorch model file to a gguf file 6 | and generate as much of the c code to use it as possible. 7 | 8 | ## Usage 9 | 10 | ```sh 11 | ./convert.py [model_name] 12 | ``` 13 | 14 | This will generate a model_name.gguf model file and a model_name.c file. 15 | -------------------------------------------------------------------------------- /convert.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | import sys 4 | import gguf 5 | import torch 6 | import numpy as np 7 | from typing import Dict 8 | 9 | STRUCT_TEMPLATE = """struct {struct_name} {{ 10 | {struct_fields} 11 | }}""" 12 | 13 | 14 | # https://stackoverflow.com/questions/651794/whats-the-best-way-to-initialize-a-dict-of-dicts-in-python 15 | class AutoVivification(dict): 16 | """Implementation of perl's autovivification feature.""" 17 | 18 | def __getitem__(self, item): 19 | try: 20 | return dict.__getitem__(self, item) 21 | except KeyError: 22 | value = self[item] = type(self)() 23 | return value 24 | 25 | 26 | DEFAULT_MODEL_NAME = "untitled" 27 | GGML_TENSOR_PTR = "struct ggml_tensor*" 28 | MAIN_STRUCT: str = "module" 29 | TAB_WIDTH = 4 30 | TAB = " " * TAB_WIDTH 31 | 32 | 33 | def make_structs(o: Dict) -> Dict[str, str]: 34 | """ 35 | :param o: the model 36 | :return: the structs used to represent the model in c 37 | """ 38 | 39 | structs = AutoVivification() 40 | 41 | def get_module_add_structs(struct_name_to_struct_fields): 42 | _main_struct = None 43 | for struct_name, struct_fields in struct_name_to_struct_fields.items(): 44 | if struct_name == MAIN_STRUCT: 45 | _main_struct = struct_fields 46 | continue 47 | structs[struct_name] = struct_fields # FIXME: override is bad 48 | return _main_struct 49 | 50 | for k in o.keys(): 51 | if type(o[k]) is dict: 52 | if all(map(lambda x: x.isnumeric(), o[k].keys())): 53 | arr_length = int(max(o[k].keys())) + 1 54 | structs[MAIN_STRUCT][k] = f"struct {k.capitalize()}[{arr_length}]" 55 | structs[k.capitalize()] = get_module_add_structs(make_structs(o[k]["0"])) 56 | else: 57 | structs[MAIN_STRUCT][k] = f"struct {k.capitalize()}" 58 | structs[k.capitalize()] = get_module_add_structs(make_structs(o[k])) 59 | else: 60 | structs[MAIN_STRUCT][k] = type(o[k]) 61 | 62 | return structs 63 | 64 | 65 | def c_gen(state_dict, model_name): 66 | code = '#include \n#include \n#include "ggml/ggml.h"\n\n\n' 67 | 68 | def k(line, tabs=0): 69 | nonlocal code 70 | code += f"{TAB * tabs}{line}\n" 71 | 72 | # gen model AST 73 | model = {} 74 | for param_name, v in state_dict.items(): 75 | parts = param_name.split(".") 76 | 77 | last = model 78 | for i in range(len(parts)): 79 | p = parts[i] 80 | if p not in last: 81 | last[p] = {} 82 | # last p 83 | if i == len(parts) - 1: 84 | last[p] = v.numpy() 85 | else: 86 | last = last[p] 87 | 88 | # gen model structs 89 | model_structs = make_structs(model) 90 | model_structs[MAIN_STRUCT]["ctx"] = "struct ggml_context*" 91 | for struct_name, struct_fields in model_structs.items(): 92 | if struct_name == MAIN_STRUCT: 93 | struct_name = model_name 94 | 95 | struct_fields = {k: (GGML_TENSOR_PTR if type is np.ndarray else type) for k, type in struct_fields.items()} 96 | 97 | struct_fields_str = "\n ".join(map(lambda x: f"{x[1]} {x[0]};", struct_fields.items())) 98 | code += STRUCT_TEMPLATE.format(struct_name=struct_name, struct_fields=struct_fields_str) 99 | code += "\n\n" 100 | 101 | # gen model load function 102 | code += f""" 103 | // returns a pointer to a loaded model struct. user is responsible of freeing it later 104 | struct {model_name}* {model_name}_model_load(const char *model_file, mnist_model & model) {{ 105 | struct {model_name} *model = malloc(sizeof(*model)); 106 | if (!model) {{ 107 | fprintf(stderr, "%s: malloc(sizeof(*model)) failed, Out of memory\\n", __func__); 108 | return NULL; 109 | }} 110 | 111 | struct gguf_init_params params = {{ 112 | /*.no_alloc =*/ false, 113 | /*.ctx =*/ &model.ctx, 114 | }}; 115 | gguf_context *ctx = gguf_init_from_file(model_file, params); 116 | if (!ctx) {{ 117 | fprintf(stderr, "%s: gguf_init_from_file() failed\\n", __func__); 118 | return NULL; 119 | }} 120 | """ 121 | 122 | for param_name in state_dict.keys(): 123 | param_name_to_accessor = "" 124 | for p in param_name.split("."): 125 | if p.isnumeric(): 126 | param_name_to_accessor += f"[{p}]" 127 | else: 128 | param_name_to_accessor += f".{p}" 129 | 130 | k(f'model{param_name_to_accessor} = ggml_get_tensor(model.ctx, "{param_name}");', tabs=1) 131 | 132 | k("") 133 | k("return model;\n}", tabs=1) 134 | 135 | return code 136 | 137 | 138 | def convert(model_path, model_name=DEFAULT_MODEL_NAME): 139 | if model_name == DEFAULT_MODEL_NAME: 140 | print(f"Warning: no provided model_name, default={DEFAULT_MODEL_NAME}") 141 | 142 | state_dict = torch.load(model_path, map_location="cpu") 143 | state_dict = {k.replace("module.", ""): v for k, v in state_dict.items() if "module." in k} 144 | 145 | code = c_gen(state_dict, model_name) 146 | c_output_path = f"./{model_name}.c" 147 | with open(c_output_path, "w") as f: 148 | f.write(code) 149 | print(f"Model ggml code generated and saved to '{c_output_path}'") 150 | 151 | gguf_output_path = f"./{model_name}.gguf" 152 | gguf_writer = gguf.GGUFWriter(gguf_output_path, model_name) 153 | for param_name, param_value in state_dict.items(): 154 | gguf_writer.add_tensor(param_name, param_value.numpy()) 155 | 156 | # kernel1 = model.layers[0].weights[0].numpy() 157 | # kernel1 = np.moveaxis(kernel1, [2, 3], [0, 1]) 158 | # kernel1 = kernel1.astype(np.float16) 159 | # gguf_writer.add_tensor("kernel1", kernel1, raw_shape=(32, 1, 3, 3)) 160 | # 161 | # bias1 = model.layers[0].weights[1].numpy() 162 | # bias1 = np.repeat(bias1, 26 * 26) 163 | # gguf_writer.add_tensor("bias1", bias1, raw_shape=(1, 32, 26, 26)) 164 | # 165 | # kernel2 = model.layers[2].weights[0].numpy() 166 | # kernel2 = np.moveaxis(kernel2, [0, 1, 2, 3], [2, 3, 1, 0]) 167 | # kernel2 = kernel2.astype(np.float16) 168 | # gguf_writer.add_tensor("kernel2", kernel2, raw_shape=(64, 32, 3, 3)) 169 | # 170 | # bias2 = model.layers[2].weights[1].numpy() 171 | # bias2 = np.repeat(bias2, 11 * 11) 172 | # gguf_writer.add_tensor("bias2", bias2, raw_shape=(1, 64, 11, 11)) 173 | # 174 | # dense_w = model.layers[-1].weights[0].numpy() 175 | # dense_w = dense_w.transpose() 176 | # gguf_writer.add_tensor("dense_w", dense_w, raw_shape=(10, 1600)) 177 | # 178 | # dense_b = model.layers[-1].weights[1].numpy() 179 | # gguf_writer.add_tensor("dense_b", dense_b) 180 | 181 | gguf_writer.write_header_to_file() 182 | gguf_writer.write_kv_data_to_file() 183 | gguf_writer.write_tensors_to_file() 184 | gguf_writer.close() 185 | print(f"Model converted and saved to '{gguf_output_path}'") 186 | 187 | 188 | if __name__ == '__main__': 189 | if len(sys.argv) < 2: 190 | print(f"Usage: ./{sys.argv[0]} [model_name]") 191 | sys.exit(1) 192 | 193 | if len(sys.argv) < 3: 194 | convert(sys.argv[1]) 195 | else: 196 | convert(sys.argv[1], sys.argv[2]) 197 | -------------------------------------------------------------------------------- /test.py: -------------------------------------------------------------------------------- 1 | import unittest 2 | from convert import make_structs 3 | 4 | 5 | class TestParse(unittest.TestCase): 6 | def test_parse_simple(self): 7 | parser_input = {"weight": None, "bias": None} 8 | expected = {"module": {"weight": type(None), "bias": type(None)}} 9 | 10 | self.assertEqual(make_structs(parser_input), expected) 11 | 12 | def test_parse_struct(self): 13 | parser_input = {"linear": {"weight": None, "bias": None}} 14 | expected = {"module": {"linear": "struct Linear"}, "Linear": {"weight": type(None), "bias": type(None)}} 15 | 16 | self.assertEqual(make_structs(parser_input), expected) 17 | 18 | def test_parse_nested_struct(self): 19 | parser_input = {"block": {"linear": {"weight": None, "bias": None}}} 20 | expected = {"module": {"block": "struct Block"}, "Block": {"linear": "struct Linear"}, "Linear": {"weight": type(None), "bias": type(None)}} 21 | 22 | self.assertEqual(make_structs(parser_input), expected) 23 | 24 | def test_parse_array(self): 25 | parser_input = {"conv0": {'0': {"weight": None, "bias": None}}} 26 | expected = {"module": {"conv0": "struct Conv0[1]"}, "Conv0": {"weight": type(None), "bias": type(None)}} 27 | 28 | self.assertEqual(make_structs(parser_input), expected) 29 | 30 | 31 | if __name__ == '__main__': 32 | unittest.main() 33 | --------------------------------------------------------------------------------