File size: 2,923 Bytes
0cb29e4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 |
import argparse
from gguf.gguf_reader import GGUFReader
from gguf import GGUFWriter,GGUFValueType,ReaderField,OrderedDict # noqa: E402
def add_keys(writer: GGUFWriter, fields: "OrderedDict[str, ReaderField]"):
for key, field in fields.items():
if(key not in ["general.architecture","GGUF.version","GGUF.tensor_count","GGUF.kv_count"]):
if field.types[0] == GGUFValueType.STRING:
writer.add_string(key=key, val=''.join(chr(i) for i in field.parts[field.data[0]]))
elif field.types[0] == GGUFValueType.ARRAY:
writer.add_array(key=key, val=field.contents())
else:
writer.add_key_value(key=key, val=field.parts[field.data[0]][0],vtype=field.types[0])
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Merge GGUF models, especially embedding tables.")
parser.add_argument("model_src_path", help="Path to the main model GGUF file.")
parser.add_argument("embed_src_path", help="Path to the model GGUF file to take the embeddings table from (or other tensors).")
parser.add_argument("dst_path", help="Path to the output GGUF file.")
parser.add_argument("--target_blocks", nargs="+", default=["token_embd.weight"],
help="List of tensor names to merge from the embedding file. Default: token_embd.weight")
args = parser.parse_args()
reader_model = GGUFReader(args.model_src_path)
reader_embed = GGUFReader(args.embed_src_path)
archField = reader_model.get_field("general.architecture")
if archField is None:
print("Couldn't get arch from src0 file")
exit(-1)
arch = str(''.join(chr(i) for i in archField.parts[archField.data[0]]))
archField = reader_model.get_field("general.architecture")
if archField is None:
print("Couldn't get arch from src1 file")
exit(-1)
if str(''.join(chr(i) for i in archField.parts[archField.data[0]])) != arch:
print("src0 and sdc1 have different architectures")
exit(-1)
writer = GGUFWriter(path=args.dst_path, arch=arch)
add_keys(writer,reader_model.fields)
for tensor in reader_model.tensors:
# print(tensor.name)
if tensor.name in args.target_blocks:
name = tensor.name
for tensorSrc in reader_embed.tensors:
if tensorSrc.name == name:
writer.add_tensor(name = tensorSrc.name, tensor=tensorSrc.data, raw_shape=tensorSrc.shape.tolist().reverse(),raw_dtype= tensorSrc.tensor_type)
break
else:
writer.add_tensor(name = tensor.name, tensor=tensor.data, raw_shape=tensor.shape.tolist().reverse(),raw_dtype= tensor.tensor_type)
writer.write_header_to_file()
writer.write_kv_data_to_file()
writer.write_tensors_to_file()
# exit(0)
writer.close()
|