build(agent): jabba#56a767 iteration

nebulaforge/qdiff.py

@@ -0,0 +1,137 @@
+from __future__ import annotations
+
+"""QDiff: compact quantized delta/checkpoint utilities.
+
+This module provides a minimal, well-tested implementation of a
+quantized residual delta suitable for toy workloads and unit tests.
+
+Features implemented:
+- 8-bit symmetric quantization of residuals
+- optional top-k sparsification (keep largest magnitude residuals)
+- chunking and Merkle-style manifest (SHA256 per-chunk)
+
+This is intentionally small: it demonstrates the QDiff concept and is
+meant to be extended by other agents (e.g. block-sparse, file-backed,
+and resumable manifests).
+"""
+
+from dataclasses import dataclass, asdict
+import hashlib
+import json
+from typing import Dict, List, Optional, Tuple
+
+import numpy as np
+
+
+@dataclass
+class QDiffBundle:
+    header: Dict
+    chunks: Dict[str, bytes]
+
+
+def _quantize_residuals(base: np.ndarray, new: np.ndarray) -> Tuple[np.ndarray, float]:
+    # compute residuals and a symmetric scale to fit in int8
+    resid = new - base
+    max_abs = float(np.max(np.abs(resid)))
+    if max_abs == 0.0:
+        return np.zeros_like(resid, dtype=np.int8), 1.0
+    scale = max_abs / 127.0
+    q = np.round(resid / scale).astype(np.int8)
+    return q, scale
+
+
+def _dequantize(q: np.ndarray, scale: float) -> np.ndarray:
+    return q.astype(np.float32) * scale
+
+
+def _sparsify(q: np.ndarray, top_k: Optional[int]) -> Tuple[np.ndarray, Optional[List[int]]]:
+    if top_k is None or top_k >= q.size:
+        return q, None
+    # keep top_k by absolute magnitude
+    idx = np.argsort(np.abs(q))[-top_k:]
+    mask = np.zeros(q.size, dtype=bool)
+    mask[idx] = True
+    sparse_q = np.zeros_like(q)
+    sparse_q[mask] = q[mask]
+    return sparse_q, idx.tolist()
+
+
+def _chunk_bytes(b: bytes, chunk_size: int = 1024) -> List[bytes]:
+    return [b[i : i + chunk_size] for i in range(0, len(b), chunk_size)]
+
+
+def _sha256_hex(b: bytes) -> str:
+    return hashlib.sha256(b).hexdigest()
+
+
+def create_qdiff(
+    base: List[float],
+    new: List[float],
+    top_k: Optional[int] = None,
+    chunk_size: int = 1024,
+) -> QDiffBundle:
+    """Create a QDiff bundle describing new relative to base.
+
+    base and new are numeric arrays (lists or numpy-compatible). Returns
+    a QDiffBundle with a header and chunk map (sha256->bytes).
+    """
+    a = np.asarray(base, dtype=np.float32)
+    b = np.asarray(new, dtype=np.float32)
+    if a.shape != b.shape:
+        raise ValueError("base and new must have same shape")
+
+    q, scale = _quantize_residuals(a, b)
+    sparse_q, indices = _sparsify(q, top_k)
+
+    # serialize: header JSON, and the quantized bytes as raw int8
+    q_bytes = sparse_q.tobytes()
+
+    chunks = {}
+    chunk_list = _chunk_bytes(q_bytes, chunk_size=chunk_size)
+    manifest = []
+    for c in chunk_list:
+        h = _sha256_hex(c)
+        chunks[h] = c
+        manifest.append(h)
+
+    header = {
+        "shape": list(a.shape),
+        "dtype": "float32",
+        "quant": "int8",
+        "scale": float(scale),
+        "top_k_indices": indices,
+        "manifest": manifest,
+        "chunk_size": chunk_size,
+    }
+
+    return QDiffBundle(header=header, chunks=chunks)
+
+
+def apply_qdiff(base: List[float], bundle: QDiffBundle) -> List[float]:
+    """Apply QDiff bundle to a base array and return reconstructed new array."""
+    a = np.asarray(base, dtype=np.float32)
+    shape = tuple(bundle.header["shape"])
+    if a.shape != shape:
+        raise ValueError("base shape does not match bundle header")
+
+    # reconstruct bytes from manifest
+    manifest = bundle.header["manifest"]
+    parts = [bundle.chunks[h] for h in manifest]
+    q_bytes = b"".join(parts)
+
+    # ensure length matches
+    expected_elems = int(np.prod(shape))
+    q = np.frombuffer(q_bytes, dtype=np.int8, count=expected_elems)
+
+    scale = float(bundle.header["scale"])
+    deq = _dequantize(q, scale)
+    new = a + deq.reshape(shape)
+    return new.tolist()
+
+
+def manifest_proofs(bundle: QDiffBundle) -> Dict[str, str]:
+    """Return a mapping of chunk hash -> hex digest (Merkle leaf hashes).
+
+    This is a thin helper used by tests and auditors.
+    """
+    return {h: _sha256_hex(bundle.chunks[h]) for h in bundle.header["manifest"]}

tests/test_qdiff.py

@@ -0,0 +1,32 @@
+from nebulaforge.qdiff import create_qdiff, apply_qdiff, manifest_proofs
+
+
+def test_qdiff_create_and_apply_identity():
+    base = [0.0, 1.0, 2.0, -1.0]
+    new = [0.0, 1.0, 2.0, -1.0]
+    bundle = create_qdiff(base, new)
+    out = apply_qdiff(base, bundle)
+    # exact equality for identical arrays
+    assert out == new
+
+
+def test_qdiff_quantized_residuals_and_topk():
+    base = [0.0, 0.0, 0.0, 0.0]
+    new = [0.0, 10.0, -5.0, 2.0]
+    # keep only top-2 residuals
+    bundle = create_qdiff(base, new, top_k=2, chunk_size=8)
+    out = apply_qdiff(base, bundle)
+    # Reconstructed values should be close to original for top-k kept; others may be zero
+    # Check that at least two elements match closely
+    matches = sum(1 for a, b in zip(out, new) if abs(a - b) < 1e-2)
+    assert matches >= 2
+
+
+def test_manifest_hashes():
+    base = [0.0] * 16
+    new = [i * 0.1 for i in range(16)]
+    bundle = create_qdiff(base, new, chunk_size=7)
+    proofs = manifest_proofs(bundle)
+    # every manifest entry should have a corresponding proof equal to the key
+    for k, v in proofs.items():
+        assert k == v