添加chat template到gguf模型文件

添加chat_template到gguf模型文件

目前使用的很多gguf模型文件缺少chat_template, 为了更好的使用，需要添加chat_template到gguf模型文件中

查询了gguf-py代码，hack如下

• 需要知道该model的chat_template, 以openhermes-2.5-mistral-7b.Q5_K_M.gguf为例, chat_template如下

“{% for message in messages %}{{‘<|im_start|>’ + message[‘role’] + ‘\n’ + message[‘content’] + ‘<|im_end|>’ + ‘\n’}}{% endfor %}{% if add_generation_prompt %}{{ ‘<|im_start|>assistant\n’ }}{% endif %}”

• 查看gguf模型中kv的结构，找到需要插入的位置，确认chat_template的bytes的长度，根据model对齐长度，添加空格到chat_template的bytes的长度

• 构件chat_template的bytes，添加到模型文件中

• 重新写入新的gguf文件

代码如下

import struct
import numpy as np
from gguf import GGUFReader, GGUFValueType, GGUF_DEFAULT_ALIGNMENT

# 原始gguf文件路径和新的gguf文件路径
file_path = "../openhermes.gguf"
new_file_path = "../add_chat_model.gguf"
#
reader = GGUFReader(file_path, "r+")
#

# 定义chat_template, 查询该模型的tokenizer_config.json文件
CHAT_TEMPLATE = "tokenizer.chat_template"
chat_template = "{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}"

# 生成对齐的chat_template的bytes
alignment = GGUF_DEFAULT_ALIGNMENT
new_align = reader.fields.get('general.alignment')
if new_align is not None:
    alignment = new_align.parts[-1][0]

add_data = bytearray()

name_data = CHAT_TEMPLATE.encode("utf-8")
add_data += struct.pack("Q", len(name_data))
add_data += name_data
add_data += struct.pack("I", GGUFValueType.STRING.value)

raw_len = len(add_data) + 8 + len(chat_template)
add_len = alignment - (raw_len % alignment)
if add_len != 0:
    chat_template += " " * add_len

raw_data = chat_template.encode("utf-8")
add_data += struct.pack("Q", len(raw_data))
add_data += raw_data

# 找到插入的位置
kv = reader.fields
last_field = list(kv.values())[-1]
insert_offset = last_field.offset

# 复制新模型文件，并插入chat_template
new_data = reader.data.copy()
new_data = np.concatenate(
    (new_data[:insert_offset], add_data, new_data[insert_offset:]))

# kv_count 加1
kv_count_idx = reader.fields["GGUF.kv_count"].parts[0][0]
new_data[kv_count_idx] += 1

# 保存文件
with open(new_file_path, "wb") as file:
    file.write(new_data.tobytes())