model-explorer-gguf 1.0.0.post1.dev0__py3-none-any.whl

model_explorer_gguf/__init__.py

@@ -0,0 +1,10 @@
+"""GGUF adapter for Model Explorer."""
+
+from .main import GGUFAdapter
+
+try:
+    from ._version import __version__  # type: ignore
+except ImportError:
+    __version__ = "0.0.0"
+
+__all__ = ["GGUFAdapter", "__version__"]

model_explorer_gguf/_version.py

@@ -0,0 +1 @@
+__version__ = "1.0.0.post1.dev0"

model_explorer_gguf/bin/ggufvis.py

@@ -0,0 +1,115 @@
+"""A shortcut to run model explorer with GGUF extension."""
+
+import argparse
+import json
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+
+def load_config(config_path, model_path):
+    """Load config from file, checking standard locations."""
+    candidates = []
+    if config_path:
+        candidates.append(Path(config_path))
+    candidates.append(Path.cwd() / ".ggufvis.json")
+    candidates.append(Path(model_path).parent / ".ggufvis.json")
+    for path in candidates:
+        if path.is_file():
+            try:
+                return json.loads(path.read_text())
+            except (json.JSONDecodeError, OSError):
+                continue
+    return {}
+
+
+def parse_settings_args(settings_args):
+    """Parse KEY=VALUE settings from CLI."""
+    if not settings_args:
+        return {}
+    result = {}
+    for arg in settings_args:
+        if "=" not in arg:
+            continue
+        key, value = arg.split("=", 1)
+        value = value.strip()
+        if value.lower() in ("true", "false"):
+            result[key.strip()] = value.lower() == "true"
+        else:
+            try:
+                result[key.strip()] = int(value)
+            except ValueError:
+                try:
+                    result[key.strip()] = float(value)
+                except ValueError:
+                    result[key.strip()] = value
+    return result
+
+
+def format_stats_text(stats):
+    """Format stats dict as readable text."""
+    return "\n".join(f"  {k}: {v}" for k, v in stats.items())
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Visualize GGUF models with Model Explorer")
+    parser.add_argument("model", help="Path to .gguf file")
+    parser.add_argument("--compare", metavar="MODEL", help="Path to second .gguf for comparison")
+    parser.add_argument("--stats", action="store_true", help="Print model statistics to stdout")
+    parser.add_argument(
+        "--format",
+        choices=["json", "text"],
+        default="text",
+        help="Output format for --stats (default: text)",
+    )
+    parser.add_argument("--config", metavar="PATH", help="Path to config file")
+    parser.add_argument(
+        "--settings",
+        nargs="*",
+        metavar="KEY=VALUE",
+        help="Override settings (e.g., show_tensor_values=true)",
+    )
+    args = parser.parse_args()
+
+    file_config = load_config(args.config, args.model)
+    cli_settings = parse_settings_args(args.settings)
+    merged = {**file_config, **cli_settings}
+
+    if args.stats:
+        from model_explorer_gguf.gguf_parser import GGUFParser
+        from model_explorer_gguf.settings import Settings
+
+        gguf_parser = GGUFParser(args.model)
+        gguf_parser.parse()
+        stats = {
+            "model_name": gguf_parser.get_model_name(),
+            "architecture": gguf_parser.get_architecture(),
+            "context_length": gguf_parser.get_context_length(),
+            "embedding_length": gguf_parser.get_embedding_length(),
+            "block_count": gguf_parser.get_block_count(),
+            "tensor_count": gguf_parser.header.tensor_count,
+            "gguf_version": gguf_parser.header.version,
+        }
+        if args.format == "json":
+            print(json.dumps(stats, indent=2))
+        else:
+            print(format_stats_text(stats))
+        return
+
+    cmd = ["model-explorer", "--extensions", "model_explorer_gguf", args.model]
+    if args.compare:
+        merged["comparison_model_path"] = args.compare
+
+    env = os.environ.copy()
+    if merged:
+        env["GGUFVIS_SETTINGS"] = json.dumps(merged)
+
+    try:
+        subprocess.run(cmd, env=env)
+    except KeyboardInterrupt:
+        pass
+
+
+if __name__ == "__main__":
+    main()

model_explorer_gguf/compare.py

@@ -0,0 +1,102 @@
+"""Model comparison and diff view for GGUF models."""
+
+from typing import Any, Dict, Optional
+
+from model_explorer.graph_builder import Graph
+
+from . import graph_builder as _gb
+from .gguf_parser import GGUFParser
+from .settings import Settings
+
+
+def compare_models(
+    model_a_path: str, model_b_path: str, settings: Optional[Settings] = None
+) -> Dict[str, Any]:
+    """Compare two GGUF models and return diff graphs for Model Explorer.
+
+    Returns a ModelExplorerGraphs dict with a graphCollections containing
+    both models side-by-side.
+    """
+    if settings is None:
+        settings = Settings()
+
+    parser_a = GGUFParser(model_a_path)
+    parser_a.parse()
+    parser_b = GGUFParser(model_b_path)
+    parser_b.parse()
+
+    if settings.show_tensor_values:
+        parser_a.load_all_tensor_data()
+        parser_b.load_all_tensor_data()
+
+    graph_a = _gb.build_graph(parser_a, settings=settings)
+    graph_b = _gb.build_graph(parser_b, settings=settings)
+
+    diff_attrs = _diff_metadata(parser_a, parser_b)
+    _apply_diff_metadata(graph_a, diff_attrs)
+    _apply_diff_metadata(graph_b, diff_attrs)
+
+    label_a = parser_a.get_model_name()
+    label_b = parser_b.get_model_name()
+
+    return {
+        "graphCollections": [
+            {
+                "label": f"{label_a} vs {label_b}",
+                "graphs": [graph_a, graph_b],
+            }
+        ]
+    }
+
+
+def _apply_diff_metadata(graph: Graph, diff_attrs: Dict[str, str]) -> None:
+    if graph.groupNodeAttributes:
+        graph.groupNodeAttributes[""].update(diff_attrs)
+    else:
+        graph.groupNodeAttributes = {"": diff_attrs}
+
+
+def _diff_metadata(parser_a: GGUFParser, parser_b: GGUFParser) -> Dict[str, str]:
+    """Compute diff summary between two parsers for GroupNodeAttributes."""
+    names_a = {t.name for t in parser_a.tensors}
+    names_b = {t.name for t in parser_b.tensors}
+
+    matching = names_a & names_b
+    only_a = names_a - names_b
+    only_b = names_b - names_a
+
+    attrs: Dict[str, str] = {
+        "diff.matching_count": str(len(matching)),
+        "diff.only_in_a_count": str(len(only_a)),
+        "diff.only_in_b_count": str(len(only_b)),
+    }
+
+    if only_a:
+        attrs["diff.only_in_a"] = ", ".join(sorted(only_a)[:20])
+    if only_b:
+        attrs["diff.only_in_b"] = ", ".join(sorted(only_b)[:20])
+
+    shape_diffs = []
+    type_diffs = []
+    tensor_a_map = {t.name: t for t in parser_a.tensors}
+    tensor_b_map = {t.name: t for t in parser_b.tensors}
+
+    for name in sorted(matching):
+        ta = tensor_a_map[name]
+        tb = tensor_b_map[name]
+        if ta.dimensions != tb.dimensions:
+            shape_diffs.append(name)
+        if ta.ggml_type != tb.ggml_type:
+            type_diffs.append(name)
+
+    if shape_diffs:
+        attrs["diff.shape_diffs"] = ", ".join(shape_diffs[:20])
+        attrs["diff.shape_diffs_count"] = str(len(shape_diffs))
+    if type_diffs:
+        attrs["diff.type_diffs"] = ", ".join(type_diffs[:20])
+        attrs["diff.type_diffs_count"] = str(len(type_diffs))
+
+    attrs["diff.tensor_count_a"] = str(len(parser_a.tensors))
+    attrs["diff.tensor_count_b"] = str(len(parser_b.tensors))
+
+    return attrs

model_explorer_gguf/gguf_parser.py

@@ -0,0 +1,376 @@
+"""
+GGUF file format parser.
+
+This module implements a complete parser for the GGUF (GGML Universal File) format
+as specified in the llama.cpp project.
+"""
+
+import glob
+import re
+import struct
+from dataclasses import dataclass, field
+from enum import IntEnum
+from pathlib import Path
+from typing import Any, BinaryIO, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+
+
+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
+
+
+class GGMLQuantizationType(IntEnum):
+    """GGML quantization types supported in GGUF."""
+
+    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
+    IQ1_M = 29
+    BF16 = 30
+    Q4_0_4_4 = 31
+    Q4_0_4_8 = 32
+    Q4_0_8_8 = 33
+    TQ1_0 = 34
+    TQ2_0 = 35
+
+
+GGML_TYPE_NAMES = {
+    GGMLQuantizationType.F32: "F32",
+    GGMLQuantizationType.F16: "F16",
+    GGMLQuantizationType.Q4_0: "Q4_0",
+    GGMLQuantizationType.Q4_1: "Q4_1",
+    GGMLQuantizationType.Q5_0: "Q5_0",
+    GGMLQuantizationType.Q5_1: "Q5_1",
+    GGMLQuantizationType.Q8_0: "Q8_0",
+    GGMLQuantizationType.Q8_1: "Q8_1",
+    GGMLQuantizationType.Q2_K: "Q2_K",
+    GGMLQuantizationType.Q3_K: "Q3_K",
+    GGMLQuantizationType.Q4_K: "Q4_K",
+    GGMLQuantizationType.Q5_K: "Q5_K",
+    GGMLQuantizationType.Q6_K: "Q6_K",
+    GGMLQuantizationType.Q8_K: "Q8_K",
+    GGMLQuantizationType.IQ2_XXS: "IQ2_XXS",
+    GGMLQuantizationType.IQ2_XS: "IQ2_XS",
+    GGMLQuantizationType.IQ3_XXS: "IQ3_XXS",
+    GGMLQuantizationType.IQ1_S: "IQ1_S",
+    GGMLQuantizationType.IQ4_NL: "IQ4_NL",
+    GGMLQuantizationType.IQ3_S: "IQ3_S",
+    GGMLQuantizationType.IQ2_S: "IQ2_S",
+    GGMLQuantizationType.IQ4_XS: "IQ4_XS",
+    GGMLQuantizationType.I8: "I8",
+    GGMLQuantizationType.I16: "I16",
+    GGMLQuantizationType.I32: "I32",
+    GGMLQuantizationType.I64: "I64",
+    GGMLQuantizationType.F64: "F64",
+    GGMLQuantizationType.IQ1_M: "IQ1_M",
+    GGMLQuantizationType.BF16: "BF16",
+    GGMLQuantizationType.Q4_0_4_4: "Q4_0_4_4",
+    GGMLQuantizationType.Q4_0_4_8: "Q4_0_4_8",
+    GGMLQuantizationType.Q4_0_8_8: "Q4_0_8_8",
+    GGMLQuantizationType.TQ1_0: "TQ1_0",
+    GGMLQuantizationType.TQ2_0: "TQ2_0",
+}
+
+
+@dataclass
+class GGUFHeader:
+    """GGUF file header."""
+
+    magic: bytes
+    version: int
+    tensor_count: int
+    metadata_kv_count: int
+
+
+@dataclass
+class GGUFTensorInfo:
+    """GGUF tensor information."""
+
+    name: str
+    n_dimensions: int
+    dimensions: List[int]
+    ggml_type: int
+    offset: int
+    shard_idx: int = 0
+    shard_file: str = ""
+
+
+class GGUFParser:
+    """Parser for GGUF files."""
+
+    GGUF_MAGIC = b"GGUF"
+    GGUF_VERSION = 3
+    DEFAULT_ALIGNMENT = 32
+
+    def __init__(self, file_path: str):
+        """Initialize the GGUF parser.
+
+        Args:
+            file_path: Path to the GGUF file
+        """
+        self.file_path = Path(file_path)
+        self.header: Optional[GGUFHeader] = None
+        self.metadata: Dict[str, Any] = {}
+        self.tensors: List[GGUFTensorInfo] = []
+        self.alignment = self.DEFAULT_ALIGNMENT
+        self.tensor_data_offset = 0
+        self.tensor_data: Optional[dict] = None
+
+    def parse(self) -> None:
+        """Parse the GGUF file."""
+        with open(self.file_path, "rb") as f:
+            self._parse_header(f)
+            self._parse_metadata(f)
+            self._parse_tensor_info(f)
+
+    def _parse_header(self, f: BinaryIO) -> None:
+        """Parse the GGUF header."""
+        magic = f.read(4)
+        if magic != self.GGUF_MAGIC:
+            raise ValueError(f"Invalid GGUF magic number: {magic}")
+
+        version = struct.unpack("<I", f.read(4))[0]
+        if version > self.GGUF_VERSION:
+            raise ValueError(f"Unsupported GGUF version: {version}")
+
+        tensor_count = struct.unpack("<Q", f.read(8))[0]
+        metadata_kv_count = struct.unpack("<Q", f.read(8))[0]
+
+        self.header = GGUFHeader(
+            magic=magic,
+            version=version,
+            tensor_count=tensor_count,
+            metadata_kv_count=metadata_kv_count,
+        )
+
+    def _read_string(self, f: BinaryIO) -> str:
+        """Read a GGUF string (length-prefixed)."""
+        length = struct.unpack("<Q", f.read(8))[0]
+        return f.read(length).decode("utf-8")
+
+    def _read_value(self, f: BinaryIO, value_type: int) -> Any:
+        """Read a metadata value based on its type."""
+        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.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]
+        elif value_type == GGUFValueType.BOOL:
+            return bool(struct.unpack("<B", f.read(1))[0])
+        elif value_type == GGUFValueType.STRING:
+            return self._read_string(f)
+        elif value_type == GGUFValueType.ARRAY:
+            return self._read_array(f)
+        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]
+
+        elements = []
+        for _ in range(length):
+            elements.append(self._read_value(f, element_type))
+        return elements
+
+    def _parse_metadata(self, f: BinaryIO) -> None:
+        """Parse the metadata key-value pairs."""
+        for _ in range(self.header.metadata_kv_count):
+            key = self._read_string(f)
+            value_type = struct.unpack("<I", f.read(4))[0]
+            value = self._read_value(f, value_type)
+            self.metadata[key] = value
+
+        # Get alignment value from metadata if present
+        if "general.alignment" in self.metadata:
+            self.alignment = self.metadata["general.alignment"]
+
+    def _parse_tensor_info(self, f: BinaryIO) -> None:
+        """Parse the tensor information."""
+        for _ in range(self.header.tensor_count):
+            name = self._read_string(f)
+            n_dimensions = struct.unpack("<I", f.read(4))[0]
+
+            dimensions = []
+            for _ in range(n_dimensions):
+                dimensions.append(struct.unpack("<Q", f.read(8))[0])
+
+            ggml_type = struct.unpack("<I", f.read(4))[0]
+            offset = struct.unpack("<Q", f.read(8))[0]
+
+            tensor_info = GGUFTensorInfo(
+                name=name,
+                n_dimensions=n_dimensions,
+                dimensions=dimensions,
+                ggml_type=ggml_type,
+                offset=offset,
+            )
+            self.tensors.append(tensor_info)
+
+        # Calculate tensor data offset
+        current_pos = f.tell()
+        self.tensor_data_offset = self._align_offset(current_pos, self.alignment)
+
+    def _align_offset(self, offset: int, alignment: int) -> int:
+        """Align an offset to the specified alignment."""
+        return (offset + alignment - 1) // alignment * alignment
+
+    def get_architecture(self) -> str:
+        """Get the model architecture from metadata."""
+        return self.metadata.get("general.architecture", "unknown")
+
+    def get_model_name(self) -> str:
+        """Get the model name from metadata."""
+        return self.metadata.get("general.name", "unknown")
+
+    def get_context_length(self) -> int:
+        """Get the context length from metadata."""
+        arch = self.get_architecture()
+        key = f"{arch}.context_length"
+        return self.metadata.get(key, 0)
+
+    def get_embedding_length(self) -> int:
+        """Get the embedding length from metadata."""
+        arch = self.get_architecture()
+        key = f"{arch}.embedding_length"
+        return self.metadata.get(key, 0)
+
+    def get_block_count(self) -> int:
+        """Get the number of blocks from metadata."""
+        arch = self.get_architecture()
+        key = f"{arch}.block_count"
+        return self.metadata.get(key, 0)
+
+    def get_tensor_type_name(self, ggml_type: int) -> str:
+        """Get the human-readable name of a tensor type."""
+        return GGML_TYPE_NAMES.get(ggml_type, f"UNKNOWN_{ggml_type}")
+
+    def parse_headers_only(self) -> None:
+        """Parse only the header and metadata — skip tensor info.
+
+        This is a fast path for quick metadata inspection. It reads:
+        - Magic (4 bytes) + version (4 bytes) + tensor_count (8 bytes) +
+          metadata_kv_count (8 bytes) = 24 bytes fixed header
+        - Metadata KV pairs (bounded by metadata_kv_count)
+        - Returns without computing alignment or reading tensor info
+        """
+        with open(self.file_path, "rb") as f:
+            self._parse_header(f)
+            self._parse_metadata(f)
+
+    def parse_multi_shard(self) -> None:
+        """Parse a multi-shard GGUF model.
+
+        Shard 0 contains the full header and metadata. Each shard has its own
+        tensor infos. Tensor offsets are adjusted to be absolute within each shard.
+        """
+        shards = discover_shards(str(self.file_path))
+
+        if len(shards) <= 1:
+            self.parse()
+            return
+
+        self.file_path = Path(shards[0])
+        with open(self.file_path, "rb") as f:
+            self._parse_header(f)
+            self._parse_metadata(f)
+            self._parse_tensor_info(f)
+
+        self._shard_paths = [Path(s) for s in shards]
+
+    def load_tensor_data(self, tensor_name: str) -> "np.ndarray":
+        """Load and dequantize a single tensor's data using gguf.GGUFReader."""
+        from gguf import GGUFReader as _GGUFReader
+        from gguf import dequantize as _dequantize
+
+        reader = _GGUFReader(str(self.file_path))
+        for rt in reader.tensors:
+            if rt.name == tensor_name:
+                return _dequantize(rt.data, rt.tensor_type)
+        raise KeyError(f"Tensor '{tensor_name}' not found")
+
+    def load_all_tensor_data(self) -> dict:
+        """Load and dequantize all tensor data. Returns {name: float32 array}."""
+        from gguf import GGUFReader as _GGUFReader
+        from gguf import dequantize as _dequantize
+
+        result = {}
+        reader = _GGUFReader(str(self.file_path))
+        for rt in reader.tensors:
+            result[rt.name] = _dequantize(rt.data, rt.tensor_type)
+        self.tensor_data = result
+        return result
+
+
+def discover_shards(model_path: str) -> list:
+    """Find all shard files for a split GGUF model.
+
+    If the model_path matches *.gguf-split.*, finds all shards in the same
+    directory and returns them sorted by shard index. Otherwise returns
+    [model_path] for single-file mode.
+    """
+    path = Path(model_path)
+    if ".gguf-split." not in path.name:
+        return [model_path]
+
+    pattern = str(path.parent / "*.gguf-split.*")
+    shards = sorted(
+        glob.glob(pattern),
+        key=lambda p: int(re.search(r"\.(\d+)-of-\d+$", p).group(1)),
+    )
+    return shards if shards else [model_path]