Upload e8_utils.py with huggingface_hub

e8_utils.py

@@ -0,0 +1,173 @@
+"""
+e8_utils.py — fully self-contained E8 lattice utilities for the LatticeMemory HF Space.
+
+No imports from liora_core. All math is inlined.
+"""
+from __future__ import annotations
+
+import math
+from itertools import combinations
+
+import torch
+import torch.nn.functional as F
+
+
+# ---------------------------------------------------------------------------
+# Core E8 lattice math
+# ---------------------------------------------------------------------------
+
+def _decode_d8(x: torch.Tensor) -> torch.Tensor:
+    """Round x to the nearest D8 lattice point (even-sum integer coordinates)."""
+    z = torch.round(x)
+    parity = z.sum(dim=-1) % 2
+    diff = x - z
+    worst = diff.abs().argmax(dim=-1)
+    adj = torch.sign(diff)
+    adj = torch.where(adj == 0, torch.ones_like(adj), adj)
+    mask = torch.zeros_like(x)
+    mask.scatter_(-1, worst.unsqueeze(-1), 1.0)
+    z_fixed = z + adj * mask
+    return torch.where(parity.unsqueeze(-1) == 1, z_fixed, z)
+
+
+def _e8_nearest(x: torch.Tensor) -> torch.Tensor:
+    """Return the nearest E8 lattice point to x (batched, last dim = 8)."""
+    z0 = _decode_d8(x)
+    z1 = _decode_d8(x - 0.5) + 0.5
+    d0 = (x - z0).pow(2).sum(dim=-1, keepdim=True)
+    d1 = (x - z1).pow(2).sum(dim=-1, keepdim=True)
+    return torch.where(d0 <= d1, z0, z1)
+
+
+# ---------------------------------------------------------------------------
+# Shell-1 codebook  (240 vectors, shape [240, 8])
+# ---------------------------------------------------------------------------
+
+def _build_shell1_codebook(device: torch.device = torch.device("cpu")) -> torch.Tensor:
+    """Build the 240-vector E8 shell-1 codebook.
+
+    E8 shell-1 consists of:
+      - 112 vectors of the form (±1, ±1, 0, 0, 0, 0, 0, 0) in all permutations
+      - 128 vectors of the form (±½, ±½, …, ±½) with an even number of minus signs
+    Total: 112 + 128 = 240 vectors.
+    """
+    vecs: list[list[float]] = []
+    # ±1 in two coordinates, rest 0
+    for i, j in combinations(range(8), 2):
+        for si in (1.0, -1.0):
+            for sj in (1.0, -1.0):
+                v = [0.0] * 8
+                v[i] = si
+                v[j] = sj
+                vecs.append(v)
+    # (±½)^8 with even number of minus signs
+    for mask in range(256):
+        signs = [1.0 if (mask >> bit) & 1 == 0 else -1.0 for bit in range(8)]
+        if signs.count(-1.0) % 2 == 0:
+            vecs.append([s * 0.5 for s in signs])
+    codebook = torch.tensor(vecs, dtype=torch.float32, device=device)
+    if codebook.shape != (240, 8):
+        raise RuntimeError(
+            f"expected E8 shell-1 codebook shape (240, 8), got {tuple(codebook.shape)}"
+        )
+    return codebook
+
+
+# ---------------------------------------------------------------------------
+# RF-Snap: batch snap embeddings to E8
+# ---------------------------------------------------------------------------
+
+@torch.no_grad()
+def nestquant_snap(embeddings: torch.Tensor) -> torch.Tensor:
+    """Snap a batch of embeddings [B, D] (D divisible by 8) to E8 lattice points.
+
+    Each 8-dim block is independently scaled, snapped to the nearest E8 point,
+    and rescaled back.  The operation is a no-op in the limit of small beta.
+
+    Args:
+        embeddings: float32 tensor of shape [B, D].
+
+    Returns:
+        Snapped embeddings of shape [B, D].
+    """
+    if embeddings.dim() != 2:
+        raise ValueError(f"nestquant_snap expects [B, D], got {tuple(embeddings.shape)}")
+    B, D = embeddings.shape
+    if D % 8 != 0:
+        raise ValueError(f"D={D} must be divisible by 8")
+
+    blocks = embeddings.float().reshape(B, D // 8, 8)           # [B, n_blocks, 8]
+    beta = blocks.norm(p=2, dim=-1).clamp_min(1e-8) / math.sqrt(2.0)  # [B, n_blocks]
+    snapped = _e8_nearest(blocks / beta.unsqueeze(-1))          # [B, n_blocks, 8]
+    return (snapped * beta.unsqueeze(-1)).reshape(B, D)
+
+
+# ---------------------------------------------------------------------------
+# E8 address: encode a single embedding as a hex string
+# ---------------------------------------------------------------------------
+
+@torch.no_grad()
+def embedding_to_e8_address(embedding: torch.Tensor, device: torch.device = torch.device("cpu")) -> str:
+    """Convert a single embedding vector to its E8 lattice address (hex string).
+
+    Each 8-dim block is mapped to an index in [0, 239] (the closest shell-1
+    codebook vector after unit-normalisation), then packed as a byte.  The
+    resulting byte string is returned as a lowercase hex string.
+
+    Expected address length: (D // 8) * 2 hex chars.  For D=1024 → 256 chars.
+
+    Args:
+        embedding: 1-D float tensor of shape [D].
+        device: torch device for codebook (CPU by default).
+
+    Returns:
+        Lowercase hex string of length D // 4.
+    """
+    if embedding.dim() != 1:
+        raise ValueError(f"embedding_to_e8_address expects 1-D tensor, got {tuple(embedding.shape)}")
+    D = embedding.numel()
+    if D % 8 != 0:
+        raise ValueError(f"D={D} must be divisible by 8")
+
+    codebook = _build_shell1_codebook(device)  # [240, 8]
+
+    vector = embedding.float().to(device)
+    blocks = vector.reshape(-1, 8)                                       # [n_blocks, 8]
+    beta = blocks.norm(p=2, dim=-1).clamp_min(1e-8) / math.sqrt(2.0)   # [n_blocks]
+    snapped = _e8_nearest(blocks / beta.unsqueeze(-1))                   # [n_blocks, 8]
+
+    # Unit-normalise snapped blocks and project onto codebook
+    snapped_unit = snapped / snapped.norm(dim=-1, keepdim=True).clamp_min(1e-8)
+    dots = snapped_unit @ codebook.T / math.sqrt(2.0)                   # [n_blocks, 240]
+    indices = dots.argmax(dim=-1)                                        # [n_blocks], values in [0,239]
+
+    return bytes(indices.tolist()).hex()
+
+
+# ---------------------------------------------------------------------------
+# Index size comparison
+# ---------------------------------------------------------------------------
+
+def index_size_bytes(n_docs: int, d_model: int) -> dict[str, int]:
+    """Return byte counts for different index formats.
+
+    Keys:
+      float32  — raw fp32 storage (4 bytes/float)
+      int4     — 4-bit quantization (0.5 bytes/float)
+      rfsnap   — RF-Snap E8 (3 bytes per 8-dim block: 1 byte index + 2 bytes fp16 scale)
+
+    Args:
+        n_docs:  number of document vectors.
+        d_model: embedding dimension (must be divisible by 8).
+
+    Returns:
+        Dict with keys 'float32', 'int4', 'rfsnap' and integer byte values.
+    """
+    if d_model % 8 != 0:
+        raise ValueError(f"d_model={d_model} must be divisible by 8")
+    n_blocks = d_model // 8
+    return {
+        "float32": n_docs * d_model * 4,
+        "int4":    n_docs * d_model // 2,
+        "rfsnap":  n_docs * n_blocks * 3,
+    }