Hands‑On Build Guide: Pure‑Python GGUF and Safetensors Loader

# loader.py
from __future__ import annotations
import hashlib
import struct
import zlib
from pathlib import Path
from typing import Dict, Tuple, Iterator

# ------------------------------------------------------------------
# Optional: tiny helper to pretty‑print tensor shapes
# ------------------------------------------------------------------
def _shape_from_bytes(data: bytes) -> Tuple[int, ...]:
    """Interpret the first 4 bytes as a little‑int32 shape (for demo)."""
    if len(data) < 4:
        raise ValueError("Not enough bytes for shape")
    return struct.unpack("<I", data[:4])[0],  # dummy return for brevity

class TensorContainer:
    """Hold validated tensors, provide simple lookup and optional caching."""

    def __init__(self) -> None:
        self._tensors: Dict[str, bytes] = {}

    def add(self, name: str, data: bytes) -> None:
        """Store a tensor after checksum validation (caller already validated)."""
        self._tensors[name] = data

    def get(self, name: str) -> bytes:
        if name not in self._tensors:
            raise KeyError(f"Tensor {name!r} not found")
        return self._tensors[name]

    def keys(self) -> Iterator[str]:
        return iter(self._tensors)

    def __len__(self) -> int:
        return len(self._tensors)

GGUF_MAGIC = b"GGUF"
GGUF_VERSION = 3  # current stable version at time of writing

def parse_gguf_header(path: Path) -> Tuple[int, list]:
    """Return (version, tensor_entries) from a GGUF file.

    tensor_entries = [(name, dtype, offset, size), ...]
    """
    data = path.read_bytes()
    if data[:4] != GGUF_MAGIC:
        raise ValueError("Not a GGUF file")
    version = struct.unpack("<I", data[4:8])[0]
    if version != GGUF_VERSION:
        raise ValueError(f"Unsupported GGUF version {version}")

    # Very simplified parsing: skip to first tensor entry (offset 8 + 2 bytes padding)
    # In a full implementation you would walk the entry count & offsets.
    # Here we just demonstrate the idea.
    return version, []

def validate_gguf_checksum(payload: bytes, expected_crc: int) -> bool:
    """CRC32 check against the tensor payload."""
    computed = zlib.crc32(payload) & 0xFFFFFFFF
    return computed == expected_crc

import json

SAFETENSORS_MAGIC = b"\x00\x00\x00\x53"  # placeholder; real format uses a 4‑byte magic

def parse_safetensors_header(path: Path) -> Tuple[dict, int]:
    """Return (metadata_dict, data_offset) from a safetensors file."""
    data = path.read_bytes()
    # Locate the JSON header (terminated by a newline or null byte)
    # This is a gross simplification; a production parser uses the official spec.
    try:
        # Find the first '{' and parse until the matching '}'
        start = data.index(b"{")
        end = data.index(b"}", start) + 1
        header_json = data[start:end].decode()
        metadata = json.loads(header_json)
    except Exception as exc:
        raise ValueError("Unable to parse safetensors header") from exc
    # Data starts right after the closing '}'
    data_offset = end
    return metadata, data_offset

def validate_safetensors_checksum(metadata: dict, payload: bytes) -> bool:
    """Verify the Blake2s checksum embedded in the metadata."""
    expected = metadata.get("checksum", "")
    if not expected:
        return False
    # Blake2s digest – 32‑byte output, encoded as hex in the metadata
    computed = hashlib.blake2s(payload).hexdigest()
    return computed == expected

def merge_shards(shard_paths: list[Path], container: TensorContainer) -> TensorContainer:
    """Load several GGUF/safetensors files, validate, and merge into one container."""
    for p in shard_paths:
        # Decide format by extension for this demo
        if p.suffix.lower() == ".gguf":
            version, entries = parse_gguf_header(p)
            # Walk entries, validate, add to container
            # (full loop omitted for brevity)
            # Example:
            # for name, dtype, offset, size in entries:
            #     payload = p.read_bytes()[offset:offset+size]
            #     if not validate_gguf_checksum(payload, expected_crc):
            #         raise RuntimeError(f"Checksum mismatch in {p}")
            #     container.add(name, payload)
        elif p.suffix.lower() in {".safetensors", ".safetensor"}:
            metadata, offset = parse_safetensors_header(p)
            raw_data = p.read_bytes()[offset:]
            if not validate_safetensors_checksum(metadata, raw_data):
                raise RuntimeError(f"Checksum mismatch in {p}")
            # Assume a single tensor named "weights" for demo
            container.add("weights", raw_data)
    return container

if __name__ == "__main__":
    # Create a dummy GGUF‑like file for illustration (real files must follow spec)
    from pathlib import Path
    import struct, zlib

    dummy_path = Path("dummy.gguf")
    # magic + version + fake entry
    header = GGUF_MAGIC + struct.pack("<I", 3) + b"\x00" * 8
    # payload with fake tensor
    payload = b"\x01" * 256
    crc = zlib.crc32(payload) & 0xFFFFFFFF
    footer = struct.pack("<I", crc)
    dummy_path.write_bytes(header + payload + footer)

    cnt = TensorContainer()
    # In a real run you would call merge_shards([dummy_path], cnt)
    # Here we just add manually:
    cnt.add("embedding", payload)
    print(f"Loaded {len(cnt)} tensor(s): {list(cnt.keys())}")