Spaces:
Sleeping
Sleeping
| """LatticeMemory HF Space — interactive demo.""" | |
| from __future__ import annotations | |
| import sys | |
| import time | |
| from pathlib import Path | |
| sys.path.insert(0, str(Path(__file__).parent)) | |
| import gradio as gr | |
| import torch | |
| import torch.nn.functional as F | |
| from sentence_transformers import SentenceTransformer | |
| from e8_utils import nestquant_snap, embedding_to_e8_address, index_size_bytes | |
| import json | |
| import plotly.graph_objects as go | |
| # --------------------------------------------------------------------------- | |
| # Load model and index at startup | |
| # --------------------------------------------------------------------------- | |
| MODEL_ID = "dfrokido/bge-large-e8-snap" | |
| INDEX_PATH = Path(__file__).parent / "data" / "index.pt" | |
| print("Loading model...") | |
| _device = "cuda" if torch.cuda.is_available() else "cpu" | |
| _model = SentenceTransformer(MODEL_ID, device=_device) | |
| print("Loading index...") | |
| _idx = torch.load(INDEX_PATH, weights_only=False) | |
| _float_norm = _idx["float_norm"] # [N, D] | |
| _snap_norm = _idx["snap_norm"] # [N, D] | |
| _doc_ids = _idx["doc_ids"] | |
| _doc_texts = _idx["doc_texts"] | |
| _d_model = _idx["d_model"] | |
| _n_docs = _idx["n_docs"] | |
| _sizes = _idx["sizes"] | |
| _sz = index_size_bytes(_n_docs, _d_model) | |
| INDEX_STATS = ( | |
| f"**Index:** {_n_docs} docs · {_d_model}-dim\n\n" | |
| f"| Method | Size | Compression |\n" | |
| f"|---|---|---|\n" | |
| f"| float32 | {_sz['float32']/1e6:.2f} MB | 1x |\n" | |
| f"| int4 | {_sz['int4']/1e6:.2f} MB | 8x |\n" | |
| f"| **RF-Snap** | **{_sz['rfsnap']/1e6:.2f} MB** | **10.7x** |" | |
| ) | |
| # --------------------------------------------------------------------------- | |
| # Retrieval logic | |
| # --------------------------------------------------------------------------- | |
| def _encode(text: str) -> torch.Tensor: | |
| """Encode a single query string to a normalized float32 tensor [1, D].""" | |
| emb = _model.encode([text], convert_to_tensor=True) | |
| return F.normalize(emb.float().cpu(), p=2, dim=1) | |
| def retrieve(query: str) -> tuple[str, str, str]: | |
| """Encode query, retrieve top-3 with float32 and RF-Snap, return results.""" | |
| if not query.strip(): | |
| return "", "", INDEX_STATS | |
| q_emb = _encode(query) # [1, D] | |
| # Float32 retrieval | |
| t0 = time.perf_counter() | |
| float_scores = (_float_norm @ q_emb.T).squeeze() | |
| float_top = torch.topk(float_scores, k=3).indices.tolist() | |
| float_ms = (time.perf_counter() - t0) * 1000 | |
| # RF-Snap retrieval | |
| q_snap = F.normalize(nestquant_snap(q_emb), p=2, dim=1) | |
| t0 = time.perf_counter() | |
| snap_scores = (_snap_norm @ q_snap.T).squeeze() | |
| snap_top = torch.topk(snap_scores, k=3).indices.tolist() | |
| snap_ms = (time.perf_counter() - t0) * 1000 | |
| # E8 address of query | |
| address_hex = embedding_to_e8_address(q_emb.squeeze(0)) | |
| address_display = " ".join( | |
| address_hex[i:i+8] for i in range(0, len(address_hex), 8) | |
| ) | |
| # Results table | |
| rows = [] | |
| for rank, (fi, si) in enumerate(zip(float_top, snap_top), 1): | |
| f_text = _doc_texts[fi][:120].replace("|", "/") | |
| s_text = _doc_texts[si][:120].replace("|", "/") | |
| match = "yes" if fi == si else "~" | |
| rows.append(f"| {rank} | {f_text}... | {s_text}... | {match} |") | |
| results_md = ( | |
| f"**float32:** {float_ms:.2f} ms " | |
| f"**RF-Snap:** {snap_ms:.2f} ms " | |
| f"**Speedup:** {float_ms/max(snap_ms,0.01):.1f}x\n\n" | |
| f"| Rank | Float32 result | RF-Snap result | Match |\n" | |
| f"|---|---|---|---|\n" | |
| + "\n".join(rows) | |
| ) | |
| return address_display, results_md, INDEX_STATS | |
| # --------------------------------------------------------------------------- | |
| # Benchmark figure | |
| # --------------------------------------------------------------------------- | |
| def _build_benchmark_figure() -> go.Figure: | |
| """Load benchmark.json and build a Plotly latency-scaling chart.""" | |
| data_path = Path(__file__).parent / "data" / "benchmark.json" | |
| if not data_path.exists(): | |
| fig = go.Figure() | |
| fig.add_annotation(text="Benchmark data not found", showarrow=False) | |
| return fig | |
| with open(data_path) as f: | |
| raw = json.load(f)["results"] | |
| sizes = sorted(int(k) for k in raw.keys()) | |
| methods = { | |
| "float32": ("Float32 scan", "#ef4444", "dash"), | |
| "int4": ("Int4 scan", "#f97316", "dot"), | |
| "rfsnap_hit": ("RF-Snap hit", "#22c55e", "solid"), | |
| "rfsnap_miss": ("RF-Snap miss", "#86efac", "dashdot"), | |
| } | |
| fig = go.Figure() | |
| for key, (label, color, dash) in methods.items(): | |
| if key not in raw[str(sizes[0])]: | |
| continue | |
| y = [raw[str(n)][key]["p50"] for n in sizes] | |
| fig.add_trace(go.Scatter( | |
| x=sizes, y=y, name=label, | |
| line=dict(color=color, dash=dash, width=2), | |
| mode="lines+markers", | |
| )) | |
| fig.update_layout( | |
| title="Retrieval latency p50 vs corpus size (bge-large 1024-dim)", | |
| xaxis_title="Corpus size (docs)", | |
| yaxis_title="Latency p50 (ms)", | |
| xaxis_type="log", | |
| template="plotly_dark", | |
| legend=dict(orientation="h", yanchor="bottom", y=1.02), | |
| ) | |
| return fig | |
| _benchmark_fig = _build_benchmark_figure() | |
| # --------------------------------------------------------------------------- | |
| # UI | |
| # --------------------------------------------------------------------------- | |
| with gr.Blocks(title="LatticeMemory") as demo: | |
| gr.Markdown( | |
| "# LatticeMemory\n" | |
| "### 10.7x smaller index. Same retrieval quality. Every concept has a permanent address.\n" | |
| "Type anything — see its E8 lattice address and retrieve from a 500-doc corpus." | |
| ) | |
| with gr.Tabs(): | |
| with gr.Tab("Live Demo"): | |
| with gr.Row(): | |
| query_box = gr.Textbox( | |
| label="Query", | |
| placeholder="What is the capital of France?", | |
| lines=2, | |
| scale=3, | |
| ) | |
| submit_btn = gr.Button("Retrieve", variant="primary", scale=1) | |
| address_out = gr.Textbox( | |
| label="E8 Address (128 blocks - permanent coordinate for this concept)", | |
| lines=3, | |
| interactive=False, | |
| ) | |
| results_out = gr.Markdown(label="Top-3 Results") | |
| stats_out = gr.Markdown(value=INDEX_STATS) | |
| submit_btn.click( | |
| fn=retrieve, | |
| inputs=query_box, | |
| outputs=[address_out, results_out, stats_out], | |
| ) | |
| query_box.submit( | |
| fn=retrieve, | |
| inputs=query_box, | |
| outputs=[address_out, results_out, stats_out], | |
| ) | |
| with gr.Tab("Benchmark"): | |
| gr.Markdown( | |
| "## O(1) vs O(N) — Retrieval Latency Scaling\n" | |
| "RF-Snap hit latency stays **flat** regardless of corpus size. " | |
| "All scan methods (float32, int4) grow linearly with N. " | |
| "At 100K docs: **RF-Snap is 17x faster** than float32 scan.\n\n" | |
| "bge-large 1024-dim · 100 warmup queries per corpus size" | |
| ) | |
| gr.Plot(value=_benchmark_fig) | |
| gr.Markdown( | |
| "| Method | Compression | 100K docs p50 | vs RF-Snap |\n" | |
| "|---|---|---|---|\n" | |
| "| Float32 | 1x | 20.8 ms | 17x slower |\n" | |
| "| Int8 | 4x | 17.8 ms | 14x slower |\n" | |
| "| Int4 | 8x | 18.5 ms | 15x slower |\n" | |
| "| **RF-Snap hit** | **10.7x** | **1.2 ms** | **baseline** |\n" | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() | |