haoline 0.3.0__py3-none-any.whl

haoline/formats/gguf.py

@@ -0,0 +1,483 @@
+# Copyright (c) 2025 HaoLine Contributors
+# SPDX-License-Identifier: MIT
+
+"""
+GGUF format reader for llama.cpp models.
+
+GGUF (GGML Universal Format) is the standard format for llama.cpp
+and other GGML-based inference engines. This reader extracts:
+- Model metadata (architecture, context length, etc.)
+- Tensor information (names, shapes, quantization types)
+- Memory footprint estimates
+
+Reference: https://github.com/ggerganov/ggml/blob/master/docs/gguf.md
+"""
+
+from __future__ import annotations
+
+import struct
+from dataclasses import dataclass, field
+from enum import IntEnum
+from pathlib import Path
+from typing import Any, BinaryIO
+
+# GGUF magic number
+GGUF_MAGIC = b"GGUF"
+
+# Supported GGUF versions
+GGUF_VERSION_MIN = 2
+GGUF_VERSION_MAX = 3
+
+
+class GGMLType(IntEnum):
+    """GGML tensor data types with their properties."""
+
+    F32 = 0
+    F16 = 1
+    Q4_0 = 2
+    Q4_1 = 3
+    Q5_0 = 6
+    Q5_1 = 7
+    Q8_0 = 8
+    Q8_1 = 9
+    Q2_K = 10
+    Q3_K = 11
+    Q4_K = 12
+    Q5_K = 13
+    Q6_K = 14
+    Q8_K = 15
+    IQ2_XXS = 16
+    IQ2_XS = 17
+    IQ3_XXS = 18
+    IQ1_S = 19
+    IQ4_NL = 20
+    IQ3_S = 21
+    IQ2_S = 22
+    IQ4_XS = 23
+    I8 = 24
+    I16 = 25
+    I32 = 26
+    I64 = 27
+    F64 = 28
+    BF16 = 29
+
+
+# Bits per weight for each quantization type
+GGML_TYPE_BITS: dict[int, float] = {
+    GGMLType.F32: 32.0,
+    GGMLType.F16: 16.0,
+    GGMLType.BF16: 16.0,
+    GGMLType.Q4_0: 4.5,  # 4 bits + 0.5 for scales
+    GGMLType.Q4_1: 5.0,
+    GGMLType.Q5_0: 5.5,
+    GGMLType.Q5_1: 6.0,
+    GGMLType.Q8_0: 8.5,
+    GGMLType.Q8_1: 9.0,
+    GGMLType.Q2_K: 2.5625,
+    GGMLType.Q3_K: 3.4375,
+    GGMLType.Q4_K: 4.5,
+    GGMLType.Q5_K: 5.5,
+    GGMLType.Q6_K: 6.5625,
+    GGMLType.Q8_K: 8.5,
+    GGMLType.IQ2_XXS: 2.0625,
+    GGMLType.IQ2_XS: 2.3125,
+    GGMLType.IQ3_XXS: 3.0625,
+    GGMLType.IQ1_S: 1.5,
+    GGMLType.IQ4_NL: 4.5,
+    GGMLType.IQ3_S: 3.4375,
+    GGMLType.IQ2_S: 2.5,
+    GGMLType.IQ4_XS: 4.25,
+    GGMLType.I8: 8.0,
+    GGMLType.I16: 16.0,
+    GGMLType.I32: 32.0,
+    GGMLType.I64: 64.0,
+    GGMLType.F64: 64.0,
+}
+
+
+def ggml_type_name(type_id: int) -> str:
+    """Get human-readable name for a GGML type."""
+    try:
+        return GGMLType(type_id).name
+    except ValueError:
+        return f"UNKNOWN_{type_id}"
+
+
+class GGUFValueType(IntEnum):
+    """GGUF metadata value types."""
+
+    UINT8 = 0
+    INT8 = 1
+    UINT16 = 2
+    INT16 = 3
+    UINT32 = 4
+    INT32 = 5
+    FLOAT32 = 6
+    BOOL = 7
+    STRING = 8
+    ARRAY = 9
+    UINT64 = 10
+    INT64 = 11
+    FLOAT64 = 12
+
+
+@dataclass
+class TensorInfo:
+    """Information about a single tensor in the GGUF file."""
+
+    name: str
+    n_dims: int
+    dims: tuple[int, ...]
+    type_id: int
+    offset: int
+
+    @property
+    def type_name(self) -> str:
+        """Human-readable type name."""
+        return ggml_type_name(self.type_id)
+
+    @property
+    def n_elements(self) -> int:
+        """Total number of elements."""
+        result = 1
+        for d in self.dims:
+            result *= d
+        return result
+
+    @property
+    def bits_per_element(self) -> float:
+        """Bits per element for this tensor's type."""
+        return GGML_TYPE_BITS.get(self.type_id, 32.0)
+
+    @property
+    def size_bytes(self) -> int:
+        """Estimated size in bytes."""
+        return int(self.n_elements * self.bits_per_element / 8)
+
+
+@dataclass
+class GGUFInfo:
+    """Parsed GGUF file information."""
+
+    path: Path
+    version: int
+    tensor_count: int
+    metadata: dict[str, Any] = field(default_factory=dict)
+    tensors: list[TensorInfo] = field(default_factory=list)
+
+    @property
+    def architecture(self) -> str:
+        """Model architecture (e.g., 'llama', 'mistral')."""
+        return str(self.metadata.get("general.architecture", "unknown"))
+
+    @property
+    def model_name(self) -> str:
+        """Model name from metadata."""
+        return str(self.metadata.get("general.name", self.path.stem))
+
+    @property
+    def context_length(self) -> int | None:
+        """Maximum context length."""
+        arch = self.architecture
+        return self.metadata.get(f"{arch}.context_length")
+
+    @property
+    def embedding_length(self) -> int | None:
+        """Hidden size / embedding dimension."""
+        arch = self.architecture
+        return self.metadata.get(f"{arch}.embedding_length")
+
+    @property
+    def block_count(self) -> int | None:
+        """Number of transformer blocks/layers."""
+        arch = self.architecture
+        return self.metadata.get(f"{arch}.block_count")
+
+    @property
+    def head_count(self) -> int | None:
+        """Number of attention heads."""
+        arch = self.architecture
+        return self.metadata.get(f"{arch}.attention.head_count")
+
+    @property
+    def head_count_kv(self) -> int | None:
+        """Number of KV heads (for GQA/MQA)."""
+        arch = self.architecture
+        return self.metadata.get(f"{arch}.attention.head_count_kv")
+
+    @property
+    def vocab_size(self) -> int | None:
+        """Vocabulary size."""
+        arch = self.architecture
+        return self.metadata.get(f"{arch}.vocab_size")
+
+    @property
+    def total_params(self) -> int:
+        """Total parameter count."""
+        return sum(t.n_elements for t in self.tensors)
+
+    @property
+    def total_size_bytes(self) -> int:
+        """Total model size in bytes."""
+        return sum(t.size_bytes for t in self.tensors)
+
+    @property
+    def quantization_breakdown(self) -> dict[str, int]:
+        """Count of tensors by quantization type."""
+        breakdown: dict[str, int] = {}
+        for t in self.tensors:
+            type_name = t.type_name
+            breakdown[type_name] = breakdown.get(type_name, 0) + 1
+        return breakdown
+
+    @property
+    def size_breakdown(self) -> dict[str, int]:
+        """Size in bytes by quantization type."""
+        breakdown: dict[str, int] = {}
+        for t in self.tensors:
+            type_name = t.type_name
+            breakdown[type_name] = breakdown.get(type_name, 0) + t.size_bytes
+        return breakdown
+
+    def estimate_vram(self, context_length: int | None = None) -> dict[str, int]:
+        """
+        Estimate VRAM requirements.
+
+        Args:
+            context_length: Context length to use for KV cache estimation.
+                           Defaults to model's context_length metadata.
+
+        Returns:
+            Dict with 'weights', 'kv_cache', and 'total' in bytes.
+        """
+        ctx = context_length or self.context_length or 2048
+        weights = self.total_size_bytes
+
+        # Estimate KV cache
+        kv_cache = 0
+        n_layers = self.block_count or 32
+        hidden = self.embedding_length or 4096
+        n_kv_heads = self.head_count_kv or self.head_count or 32
+        head_dim = hidden // (self.head_count or 32)
+
+        # KV cache: 2 (K+V) * layers * ctx * kv_heads * head_dim * 2 bytes (fp16)
+        kv_cache = 2 * n_layers * ctx * n_kv_heads * head_dim * 2
+
+        return {
+            "weights": weights,
+            "kv_cache": kv_cache,
+            "total": weights + kv_cache,
+        }
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary for JSON serialization."""
+        return {
+            "path": str(self.path),
+            "version": self.version,
+            "architecture": self.architecture,
+            "model_name": self.model_name,
+            "tensor_count": self.tensor_count,
+            "total_params": self.total_params,
+            "total_size_bytes": self.total_size_bytes,
+            "context_length": self.context_length,
+            "embedding_length": self.embedding_length,
+            "block_count": self.block_count,
+            "head_count": self.head_count,
+            "head_count_kv": self.head_count_kv,
+            "vocab_size": self.vocab_size,
+            "quantization_breakdown": self.quantization_breakdown,
+            "size_breakdown": self.size_breakdown,
+            "metadata": {
+                k: v
+                for k, v in self.metadata.items()
+                if not k.startswith("tokenizer.")  # Skip large tokenizer data
+            },
+        }
+
+
+class GGUFReader:
+    """Reader for GGUF format files."""
+
+    def __init__(self, path: str | Path):
+        """
+        Initialize reader with file path.
+
+        Args:
+            path: Path to the GGUF file.
+        """
+        self.path = Path(path)
+        if not self.path.exists():
+            raise FileNotFoundError(f"GGUF file not found: {self.path}")
+
+    def read(self) -> GGUFInfo:
+        """
+        Read and parse the GGUF file.
+
+        Returns:
+            GGUFInfo with parsed metadata and tensor information.
+
+        Raises:
+            ValueError: If the file is not a valid GGUF file.
+        """
+        with open(self.path, "rb") as f:
+            return self._parse(f)
+
+    def _parse(self, f: BinaryIO) -> GGUFInfo:
+        """Parse the GGUF file structure."""
+        # Read magic
+        magic = f.read(4)
+        if magic != GGUF_MAGIC:
+            raise ValueError(f"Not a GGUF file: expected {GGUF_MAGIC!r}, got {magic!r}")
+
+        # Read version
+        version = struct.unpack("<I", f.read(4))[0]
+        if version < GGUF_VERSION_MIN or version > GGUF_VERSION_MAX:
+            raise ValueError(
+                f"Unsupported GGUF version {version} "
+                f"(supported: {GGUF_VERSION_MIN}-{GGUF_VERSION_MAX})"
+            )
+
+        # Read counts
+        tensor_count = struct.unpack("<Q", f.read(8))[0]
+        metadata_kv_count = struct.unpack("<Q", f.read(8))[0]
+
+        # Read metadata
+        metadata = {}
+        for _ in range(metadata_kv_count):
+            key = self._read_string(f)
+            value = self._read_value(f)
+            metadata[key] = value
+
+        # Read tensor info
+        tensors = []
+        for _ in range(tensor_count):
+            tensor = self._read_tensor_info(f)
+            tensors.append(tensor)
+
+        return GGUFInfo(
+            path=self.path,
+            version=version,
+            tensor_count=tensor_count,
+            metadata=metadata,
+            tensors=tensors,
+        )
+
+    def _read_string(self, f: BinaryIO) -> str:
+        """Read a length-prefixed string."""
+        length = struct.unpack("<Q", f.read(8))[0]
+        return f.read(length).decode("utf-8")
+
+    def _read_value(self, f: BinaryIO) -> Any:
+        """Read a typed metadata value."""
+        value_type = struct.unpack("<I", f.read(4))[0]
+
+        if value_type == GGUFValueType.UINT8:
+            return struct.unpack("<B", f.read(1))[0]
+        elif value_type == GGUFValueType.INT8:
+            return struct.unpack("<b", f.read(1))[0]
+        elif value_type == GGUFValueType.UINT16:
+            return struct.unpack("<H", f.read(2))[0]
+        elif value_type == GGUFValueType.INT16:
+            return struct.unpack("<h", f.read(2))[0]
+        elif value_type == GGUFValueType.UINT32:
+            return struct.unpack("<I", f.read(4))[0]
+        elif value_type == GGUFValueType.INT32:
+            return struct.unpack("<i", f.read(4))[0]
+        elif value_type == GGUFValueType.FLOAT32:
+            return struct.unpack("<f", f.read(4))[0]
+        elif value_type == GGUFValueType.BOOL:
+            return struct.unpack("<B", f.read(1))[0] != 0
+        elif value_type == GGUFValueType.STRING:
+            return self._read_string(f)
+        elif value_type == GGUFValueType.ARRAY:
+            return self._read_array(f)
+        elif value_type == GGUFValueType.UINT64:
+            return struct.unpack("<Q", f.read(8))[0]
+        elif value_type == GGUFValueType.INT64:
+            return struct.unpack("<q", f.read(8))[0]
+        elif value_type == GGUFValueType.FLOAT64:
+            return struct.unpack("<d", f.read(8))[0]
+        else:
+            raise ValueError(f"Unknown value type: {value_type}")
+
+    def _read_array(self, f: BinaryIO) -> list[Any]:
+        """Read an array value."""
+        element_type = struct.unpack("<I", f.read(4))[0]
+        length = struct.unpack("<Q", f.read(8))[0]
+
+        # Read elements based on type
+        if element_type == GGUFValueType.UINT8:
+            return list(struct.unpack(f"<{length}B", f.read(length)))
+        elif element_type == GGUFValueType.INT8:
+            return list(struct.unpack(f"<{length}b", f.read(length)))
+        elif element_type == GGUFValueType.UINT16:
+            return list(struct.unpack(f"<{length}H", f.read(length * 2)))
+        elif element_type == GGUFValueType.INT16:
+            return list(struct.unpack(f"<{length}h", f.read(length * 2)))
+        elif element_type == GGUFValueType.UINT32:
+            return list(struct.unpack(f"<{length}I", f.read(length * 4)))
+        elif element_type == GGUFValueType.INT32:
+            return list(struct.unpack(f"<{length}i", f.read(length * 4)))
+        elif element_type == GGUFValueType.FLOAT32:
+            return list(struct.unpack(f"<{length}f", f.read(length * 4)))
+        elif element_type == GGUFValueType.BOOL:
+            return [b != 0 for b in struct.unpack(f"<{length}B", f.read(length))]
+        elif element_type == GGUFValueType.STRING:
+            return [self._read_string(f) for _ in range(length)]
+        elif element_type == GGUFValueType.UINT64:
+            return list(struct.unpack(f"<{length}Q", f.read(length * 8)))
+        elif element_type == GGUFValueType.INT64:
+            return list(struct.unpack(f"<{length}q", f.read(length * 8)))
+        elif element_type == GGUFValueType.FLOAT64:
+            return list(struct.unpack(f"<{length}d", f.read(length * 8)))
+        else:
+            # Skip unknown array types
+            return []
+
+    def _read_tensor_info(self, f: BinaryIO) -> TensorInfo:
+        """Read tensor metadata."""
+        name = self._read_string(f)
+        n_dims = struct.unpack("<I", f.read(4))[0]
+        dims = tuple(struct.unpack("<Q", f.read(8))[0] for _ in range(n_dims))
+        type_id = struct.unpack("<I", f.read(4))[0]
+        offset = struct.unpack("<Q", f.read(8))[0]
+
+        return TensorInfo(
+            name=name,
+            n_dims=n_dims,
+            dims=dims,
+            type_id=type_id,
+            offset=offset,
+        )
+
+
+def is_gguf_file(path: str | Path) -> bool:
+    """
+    Check if a file is a valid GGUF file.
+
+    Args:
+        path: Path to check.
+
+    Returns:
+        True if the file starts with GGUF magic bytes.
+    """
+    path = Path(path)
+    if not path.exists() or not path.is_file():
+        return False
+
+    try:
+        with open(path, "rb") as f:
+            magic = f.read(4)
+            return magic == GGUF_MAGIC
+    except Exception:
+        return False
+
+
+def format_size(size_bytes: int) -> str:
+    """Format bytes as human-readable string."""
+    size_float = float(size_bytes)
+    for unit in ["B", "KB", "MB", "GB", "TB"]:
+        if size_float < 1024:
+            return f"{size_float:.2f} {unit}"
+        size_float /= 1024
+    return f"{size_float:.2f} PB"