#!/usr/bin/env python3
import sys
from pathlib import Path
import json

import numpy as np

# Necessary to load the local gguf package
sys.path.insert(0, str(Path(__file__).parent.parent))

from gguf import GGUFWriter  # noqa: E402

# Example usage:
def writer_example() -> None:
    # Example usage with a file
    gguf_writer = GGUFWriter("cal_test.gguf", "cal_test")

    tokens = [
        "a",
        "b",
        "c",
        "d",
        "e",
        "f",
        "g",
        "h",
        "i",
        "j",
        "k",
        "l",
        "m",
        "n",
        "o",
        "p",
        "q",
        "r",
        "s",
        "t",
        "u",
        "v",
        "w",
        "x",
        "y",
        "z",
        "\u2581", #空格
        "\n"
    ]

    vocab_size = len(tokens)
    embedding_len = 256
    ctx_len = 2048

    gguf_writer.add_block_count(1)
    gguf_writer.add_context_length(ctx_len)
    gguf_writer.add_embedding_length(embedding_len)
    gguf_writer.add_vocab_size(vocab_size)
    
    gguf_writer.add_head_count_kv(8)
    gguf_writer.add_head_count(16)
    gguf_writer.add_key_length(16)
    gguf_writer.add_value_length(16)
    gguf_writer.add_tokenizer_model("llama")
    # 添加 tokenizer 列表到元数据
    gguf_writer.add_token_list(tokens)
    output_weight = np.ones((vocab_size, embedding_len), dtype=np.float32)

    print(f"token 数量：{vocab_size}")

    # token_embd_weight = np.ones((128000,embedding_len), dtype=np.float32)
    # print(f"创建的张量形状: {output_weight.shape}")
    # gguf_writer.add_tensor("output.weight", output_weight)

    # attn_norm = np.ones((1))
    # gguf_writer.add_tensor("blk.0.attn_norm.weight", attn_norm)

    gguf_writer.write_header_to_file()
    gguf_writer.write_kv_data_to_file()
    gguf_writer.write_tensors_to_file()

    gguf_writer.close()


if __name__ == '__main__':
    writer_example()