export/export_gguf.py

import os
import sys
import json
import torch
import struct
import numpy as np
from typing import Dict, List, Optional, Tuple

sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
from src.model import FSIEdgeModel, FSIEdgeConfig


# ============================================================================
# GGUF Format Specifications
# ============================================================================
# GGUF File Structure:
# - Magic: "GGUF" (4 bytes)
# - Version (int)
# - Tensor Count (int)
# - Metadata KV pairs
# - Alignment padding
# - Tensor data (quantized)

GGUF_MAGIC = 0x46554747  # "GGUF" as uint32
GGUF_VERSION = 3

# Tensor quantization types
GGML_TYPE_F32 = 0
GGML_TYPE_F16 = 1
GGML_TYPE_Q4_0 = 2
GGML_TYPE_Q4_1 = 3
GGML_TYPE_Q5_0 = 6
GGML_TYPE_Q5_1 = 7
GGML_TYPE_Q8_0 = 8
GGML_TYPE_Q6_K = 14
GGML_TYPE_Q4_K_M = 21
GGML_TYPE_Q5_K_M = 23  # Changed from 22 to avoid conflict
GGML_TYPE_Q3_K_M = 20
GGML_TYPE_Q2_K = 17
GGML_TYPE_Q8_K_M = 24  # Changed from 23

# Metadata keys
KEY_GENERAL_ARCHITECTURE = "general.architecture"
KEY_GENERAL_NAME = "general.name"
KEY_GENERAL_DESCRIPTION = "general.description"
KEY_GENERAL_FILE_TYPE = "general.file_type"
KEY_GENERAL_PARAMETER_COUNT = "general.parameter_count"
KEY_CONTEXT_LENGTH = "context_length"
KEY_EMBEDDING_LENGTH = "embedding_length"
KEY_BLOCK_COUNT = "block_count"
KEY_FEED_FORWARD_LENGTH = "feed_forward_length"
KEY_ATTENTION_HEAD_COUNT = "attention.head_count"
KEY_ATTENTION_HEAD_COUNT_KV = "attention.head_count_kv"
KEY_ATTENTION_LAYERNORM_EPS = "attention.layer_norm_epsilon"
KEY_ROPE_DIMENSION_COUNT = "rope.dimension_count"
KEY_ROPE_FREQ_BASE = "rope.freq_base"


def quantize_q4_0(tensor: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
    """Quantize to Q4_0: 4-bit block quantization."""
    assert tensor.ndim >= 1
    shape = tensor.shape
    
    # Flatten
    flat = tensor.astype(np.float32).ravel()
    n = flat.shape[0]
    
    # Pad to block size
    block_size = 32
    if n % block_size != 0:
        pad_len = block_size - (n % block_size)
        flat = np.pad(flat, (0, pad_len))
    
    n_blocks = flat.shape[0] // block_size
    blocks = flat.reshape(n_blocks, block_size)
    
    # Per-block quantization
    scales = np.zeros(n_blocks, dtype=np.float16)
    quants = np.zeros(n_blocks * block_size // 2, dtype=np.uint8)
    
    for i in range(n_blocks):
        block = blocks[i]
        scale = np.max(np.abs(block)) / 7.0
        if scale == 0:
            scale = 1.0
        scales[i] = np.float16(scale)
        
        quant_block = np.clip(np.round(block / scale), -8, 7).astype(np.int8)
        
        # Pack into uint8 pairs
        for j in range(block_size // 2):
            lo = quant_block[2*j] & 0x0F
            hi = (quant_block[2*j + 1] & 0x0F) << 4
            quants[i * (block_size // 2) + j] = lo | hi
    
    return scales, quants


def quantize_q8_0(tensor: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
    """Quantize to Q8_0: 8-bit block quantization."""
    flat = tensor.astype(np.float32).ravel()
    block_size = 32
    
    if flat.shape[0] % block_size != 0:
        pad_len = block_size - (flat.shape[0] % block_size)
        flat = np.pad(flat, (0, pad_len))
    
    n_blocks = flat.shape[0] // block_size
    blocks = flat.reshape(n_blocks, block_size)
    
    scales = np.zeros(n_blocks, dtype=np.float16)
    quants = np.zeros(n_blocks * block_size, dtype=np.uint8)
    
    for i in range(n_blocks):
        block = blocks[i]
        scale = np.max(np.abs(block)) / 127.0
        if scale == 0:
            scale = 1.0
        scales[i] = np.float16(scale)
        
        quant_block = np.clip(np.round(block / scale), -128, 127).astype(np.int8)
        quants[i * block_size:(i + 1) * block_size] = (quant_block & 0xFF).astype(np.uint8)
    
    return scales, quants


def write_gguf_tensor(f, name: str, tensor: np.ndarray, quant_type: int):
    """Write tensor to GGUF file."""
    name_bytes = name.encode('utf-8') + b'\x00'
    
    if quant_type == GGML_TYPE_F32:
        data = tensor.astype(np.float32).tobytes()
    elif quant_type == GGML_TYPE_F16:
        data = tensor.astype(np.float16).tobytes()
    elif quant_type == GGML_TYPE_Q4_0:
        scales, quants = quantize_q4_0(tensor)
        data = scales.tobytes() + quants.tobytes()
    elif quant_type == GGML_TYPE_Q8_0:
        scales, quants = quantize_q8_0(tensor)
        data = scales.tobytes() + quants.tobytes()
    else:
        raise ValueError(f"Unsupported quant type: {quant_type}")
    
    # Tensor info
    n_dims = len(tensor.shape)
    f.write(struct.pack('I', len(name_bytes)))
    f.write(name_bytes)
    f.write(struct.pack('I', n_dims))
    f.write(struct.pack('I', quant_type))
    
    for d in tensor.shape:
        f.write(struct.pack('Q', d))
    
    # Pad to alignment
    alignment = 32
    offset = f.tell()
    padded_offset = (offset + alignment - 1) // alignment * alignment
    f.write(b'\x00' * (padded_offset - offset))
    
    f.write(data)


def export_to_gguf(model, output_path, quant_type=GGML_TYPE_Q4_0, config=None):
    """Export FSI_Edge model to GGUF format."""
    
    architecture = "fsi_edge"
    
    # Build module name mapping
    param_map = {}
    state_dict = model.state_dict()
    
    n_params = sum(p.numel() for p in model.parameters())
    
    with open(output_path, 'wb') as f:
        # Header
        f.write(struct.pack('I', GGUF_MAGIC))
        f.write(struct.pack('I', GGUF_VERSION))
        
        # Count tensors
        tensor_count = len(state_dict)
        f.write(struct.pack('Q', tensor_count))
        
        # Metadata count (key-value pairs)
        metadata = {
            KEY_GENERAL_ARCHITECTURE: architecture,
            KEY_GENERAL_NAME: f"fsi_edge-{config.d_model if config else '800m'}",
            KEY_GENERAL_DESCRIPTION: "FSI_Edge: Novel DNA Helix Memory coding model",
            KEY_GENERAL_FILE_TYPE: quant_type,
            KEY_GENERAL_PARAMETER_COUNT: n_params,
            KEY_CONTEXT_LENGTH: config.max_seq_len if config else 16384,
            KEY_EMBEDDING_LENGTH: config.d_model if config else 1536,
            KEY_BLOCK_COUNT: config.n_layers if config else 28,
            KEY_FEED_FORWARD_LENGTH: config.d_ff if config else 6144,
            KEY_ATTENTION_HEAD_COUNT: config.n_heads if config else 24,
            KEY_ATTENTION_HEAD_COUNT_KV: config.kv_heads if config else 6,
            KEY_ATTENTION_LAYERNORM_EPS: 1e-5,
            KEY_ROPE_DIMENSION_COUNT: (config.d_model // config.n_heads) if config else 64,
            KEY_ROPE_FREQ_BASE: 10000.0,
        }
        
        f.write(struct.pack('Q', len(metadata)))
        
        for key, value in metadata.items():
            key_bytes = key.encode('utf-8') + b'\x00'
            f.write(struct.pack('I', len(key_bytes)))
            f.write(key_bytes)
            
            if isinstance(value, str):
                val_bytes = value.encode('utf-8') + b'\x00'
                f.write(struct.pack('I', 8))  # string type
                f.write(struct.pack('Q', len(val_bytes)))
                f.write(val_bytes)
            elif isinstance(value, (int, np.integer)):
                f.write(struct.pack('I', 4))  # int32 type
                f.write(struct.pack('i', int(value)))
            elif isinstance(value, float):
                f.write(struct.pack('I', 10))  # float32 type
                f.write(struct.pack('f', float(value)))
            else:
                f.write(struct.pack('I', 4))
                f.write(struct.pack('i', int(value)))
        
        # Write tensors
        for name, tensor in state_dict.items():
            np_tensor = tensor.cpu().numpy()
            gguf_name = name.replace('.', '_')
            
            # Skip non-essential layers for smaller file
            if any(skip in name for skip in ['helix_step', 'buffer']):
                continue
            
            write_gguf_tensor(f, gguf_name, np_tensor, quant_type)
    
    return output_path


def convert_pytorch_to_gguf(model_path, model_size='800M', output_path=None, quant='q4_0'):
    """Convert saved PyTorch checkpoint to GGUF."""
    
    quant_map = {
        'f32': GGML_TYPE_F32,
        'f16': GGML_TYPE_F16,
        'q4_0': GGML_TYPE_Q4_0,
        'q8_0': GGML_TYPE_Q8_0,
    }
    quant_type = quant_map.get(quant, GGML_TYPE_Q4_0)
    
    size_config = {
        '360M': {'d_model': 1024, 'n_layers': 24, 'n_heads': 16, 'kv_heads': 4, 'd_ff': 4096, 'max_seq_len': 8192},
        '800M': {'d_model': 1536, 'n_layers': 28, 'n_heads': 24, 'kv_heads': 6, 'd_ff': 6144, 'max_seq_len': 16384},
        '1.5B': {'d_model': 2048, 'n_layers': 32, 'n_heads': 32, 'kv_heads': 8, 'd_ff': 8192, 'max_seq_len': 32768},
    }
    
    sc = size_config[model_size]
    config = FSIEdgeConfig(**sc)
    
    model = FSIEdgeModel(config)
    state = torch.load(model_path, map_location='cpu')
    if 'model_state_dict' in state:
        model.load_state_dict(state['model_state_dict'])
    else:
        model.load_state_dict(state)
    
    if output_path is None:
        output_path = f'fsi_edge-{model_size}-{quant}.gguf'
    
    export_to_gguf(model, output_path, quant_type, config)
    print(f"Exported GGUF: {output_path}")
    return output_path


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--model-path', type=str, required=True)
    parser.add_argument('--model-size', type=str, default='800M')
    parser.add_argument('--output', type=str, default=None)
    parser.add_argument('--quant', type=str, default='q4_0')
    args = parser.parse_args()
    
    convert_pytorch_to_gguf(args.model_path, args.model_size, args.output, args.quant)