from gguf import *

alignment = 32
embedding_length = 16
feed_forward_length = 64
patch_size = 14
num_channels = 3

kv = {
    'general.architecture': 'mllama',
    'general.file_type': 1,
    'general.name': 'SPRÅKMODELL',
    'general.description': "DIY LLM?!",
    'general.type': 'projector',
    'general.alignment': alignment,
    'mllama.vision.image_size': 560,
    'mllama.vision.patch_size': patch_size,
    'mllama.vision.embedding_length': embedding_length,
    'mllama.vision.feed_forward_length': feed_forward_length,
    'mllama.vision.projection_dim': 4096,
    'mllama.vision.attention.head_count': 16,
    'mllama.vision.attention.layer_norm_epsilon': GgufRaw(GgufFloat32.TYPE, b'\xAC\xC5\x27\x37'), # 1e-05 serialized
    'mllama.vision.block_count': 1,
    'mllama.vision.global.block_count': 0,
    'mllama.vision.max_num_tiles': 4,
    'mllama.vision.num_channels': num_channels,
    'mllama.vision.intermediate_layers_indices': [0],
}

tensors = [
    Tensor.f16('mm.0.weight', [2 * embedding_length, 1601], b'\0\0'),
    Tensor.f32('mm.0.bias', [1601], b'\0\0\0\0'),
    Tensor.f32('v.class_embd', [embedding_length], b'\0\0\0\0'),
    Tensor.f16('v.patch_embd.weight', [patch_size, patch_size, num_channels, embedding_length], b'\0\0'),
    Tensor.f16('v.position_embd.weight', [embedding_length, 1601], b'\0\0'),
    Tensor.f16(f'v.blk.0.attn_q.weight', [embedding_length, embedding_length], b'\0\0'),
    Tensor.f16(f'v.blk.0.attn_k.weight', [embedding_length, embedding_length], b'\0\0'),
    Tensor.f16(f'v.blk.0.attn_v.weight', [embedding_length, embedding_length], b'\0\0'),
    Tensor.f16(f'v.blk.0.attn_out.weight', [embedding_length, embedding_length], b'\0\0'),
    Tensor.f16(f'v.blk.0.ffn_down.weight', [embedding_length, feed_forward_length], b'\0\0'),
    Tensor.f32(f'v.blk.0.ffn_down.bias', [feed_forward_length], b'\0\0\0\0'),
    Tensor.f16(f'v.blk.0.ffn_up.weight', [feed_forward_length, embedding_length], b'\0\0'),
    Tensor.f32(f'v.blk.0.ffn_up.bias', [embedding_length], b'\0\0\0\0'),
    Tensor.f32(f'v.blk.0.ln1.weight', [embedding_length], b'\0\0\0\0'),
    Tensor.f32(f'v.blk.0.ln1.bias', [embedding_length], b'\0\0\0\0'),
    Tensor.f32(f'v.blk.0.ln2.weight', [embedding_length], b'\0\0\0\0'),
    Tensor.f32(f'v.blk.0.ln2.bias', [embedding_length], b'\0\0\0\0'),
]

serialize(alignment, kv, tensors)
