Create app.py

app.py

@@ -0,0 +1,923 @@
+#!/usr/bin/env python3
+
+import random
+import time
+import urllib.request
+
+import gradio as gr
+import spaces
+import torch
+import torch.nn.functional as F
+import triton
+import triton.language as tl
+from transformers import AutoModel, AutoTokenizer
+
+
+MODEL_ID = "SixOpen/HARE"
+
+model = AutoModel.from_pretrained(MODEL_ID, trust_remote_code=True).eval()
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+
+
+@triton.jit
+def _wkv7_fwd_kernel(
+    R, K, V, DECAY, A, O,
+    STATE_OUT, STATE_IN,
+    sab_scale, T,
+    stride_b, stride_t, stride_h,
+    H: tl.constexpr, D: tl.constexpr, BLOCK_D: tl.constexpr,
+    RETURN_STATE: tl.constexpr, HAS_INIT_STATE: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    b_idx = pid // H
+    h_idx = pid % H
+    base = b_idx * stride_b + h_idx * stride_h
+
+    di = tl.arange(0, BLOCK_D)
+    dj = tl.arange(0, BLOCK_D)
+    mask_i = di < D
+    mask_j = dj < D
+
+    if HAS_INIT_STATE:
+        s_off = b_idx * (H * D * D) + h_idx * (D * D)
+        state_ptrs = STATE_IN + s_off + di[:, None] * D + dj[None, :]
+        state_mask = mask_i[:, None] & mask_j[None, :]
+        state = tl.load(state_ptrs, mask=state_mask, other=0.0).to(tl.float32)
+    else:
+        state = tl.zeros((BLOCK_D, BLOCK_D), dtype=tl.float32)
+
+    for t in range(T):
+        t_off = base + t * stride_t
+        kt = tl.load(K + t_off + dj, mask=mask_j, other=0.0).to(tl.float32)
+        vt = tl.load(V + t_off + di, mask=mask_i, other=0.0).to(tl.float32)
+        rt = tl.load(R + t_off + dj, mask=mask_j, other=0.0).to(tl.float32)
+        dt = tl.load(DECAY + t_off + dj, mask=mask_j, other=1.0).to(tl.float32)
+        at = tl.load(A + t_off + dj, mask=mask_j, other=0.0).to(tl.float32)
+
+        sa = tl.sum(state * (-kt)[None, :], axis=1)
+        ka = kt * at
+        sab = sa[:, None] * ka[None, :]
+        state = state * dt[None, :] + sab_scale * sab + vt[:, None] * kt[None, :]
+        state = tl.minimum(tl.maximum(state, -10.0), 10.0)
+
+        out_t = tl.sum(state * rt[None, :], axis=1)
+        tl.store(O + t_off + di, out_t, mask=mask_i)
+
+    if RETURN_STATE:
+        s_off = b_idx * (H * D * D) + h_idx * (D * D)
+        state_ptrs = STATE_OUT + s_off + di[:, None] * D + dj[None, :]
+        state_mask = mask_i[:, None] & mask_j[None, :]
+        tl.store(state_ptrs, state, mask=state_mask)
+
+
+def wkv7_scan_triton(r, decay, k, v, a, sab_scale, return_state=False, init_state=None):
+    B, T, H, D = r.shape
+    r, k, v, decay, a = [x.contiguous() for x in (r, k, v, decay, a)]
+    o = torch.empty_like(r)
+    state_out = None
+    if return_state:
+        state_out = torch.empty(B, H, D, D, dtype=torch.float32, device=r.device)
+    has_init = init_state is not None
+    if has_init:
+        init_state = init_state.contiguous().float()
+    stride_b = T * H * D
+    stride_t = H * D
+    stride_h = D
+    BLOCK_D = triton.next_power_of_2(D)
+    _wkv7_fwd_kernel[(B * H,)](
+        r, k, v, decay, a, o,
+        state_out, init_state,
+        float(sab_scale), T,
+        stride_b, stride_t, stride_h,
+        H=H, D=D, BLOCK_D=BLOCK_D,
+        RETURN_STATE=return_state,
+        HAS_INIT_STATE=has_init,
+    )
+    if return_state:
+        return o, state_out
+    return o
+
+
+def find_birwkv_layers(model):
+    layers = []
+    ids = {}
+    for m in model.modules():
+        if type(m).__name__ == 'BiRWKV7Layer':
+            ids[id(m)] = len(layers)
+            layers.append(m)
+    return layers, ids
+
+
+class SpanEncoder:
+
+    def __init__(self, model, tokenizer, chunk_size=512):
+        self.model = model
+        self.tokenizer = tokenizer
+        self.device = next(model.parameters()).device
+        self.chunk_size = chunk_size
+
+        self.birwkv_layers, self.birwkv_ids = find_birwkv_layers(model)
+        self._originals = {}
+        self._hooked = False
+        self._active_states = [None] * len(self.birwkv_layers)
+        self.span_data = {}
+
+    def _hook(self):
+        if self._hooked:
+            return
+        for layer in self.birwkv_layers:
+            self._originals[id(layer)] = layer.forward
+            layer.forward = self._make_fwd(layer)
+        self._hooked = True
+
+    def _unhook(self):
+        if not self._hooked:
+            return
+        for layer in self.birwkv_layers:
+            layer.forward = self._originals[id(layer)]
+        self._originals.clear()
+        self._hooked = False
+
+    def _make_fwd(self, layer):
+        enc = self
+        idx = self.birwkv_ids[id(layer)]
+
+        def fwd(x, attention_mask=None, **kwargs):
+            B, T, C_ = x.shape
+            H, D = layer.num_heads, layer.head_size
+            prev = enc._active_states[idx]
+            if prev is not None:
+                x_prev = torch.cat([prev['last_x'], x[:, :-1]], dim=1)
+            else:
+                x_prev = F.pad(x[:, :-1], (0, 0, 1, 0))
+
+            def mix(mu):
+                return x + (x_prev - x) * torch.sigmoid(mu)
+
+            r = layer.W_r(mix(layer.mu_r)).view(B, T, H, D)
+            w = layer.W_w(mix(layer.mu_w)).view(B, T, H, D)
+            k = layer.W_k(mix(layer.mu_k)).view(B, T, H, D)
+            v = layer.W_v(mix(layer.mu_v)).view(B, T, H, D)
+            a = layer.W_a(mix(layer.mu_a)).view(B, T, H, D)
+            g = torch.sigmoid(layer.W_g(mix(layer.mu_g)))
+            sab_scale = torch.sigmoid(layer.sab_gate)
+            init_st = prev['wkv_state'] if prev else None
+
+            r_f, k_f, v_f = r.float(), k.float() * (D ** -0.5), v.float()
+            a_f = torch.sigmoid(a.float())
+            decay = torch.exp(-0.6065306597633104 * torch.sigmoid(w.float()))
+            out_fwd, wkv_state = wkv7_scan_triton(
+                r_f, decay, k_f, v_f, a_f, sab_scale,
+                return_state=True, init_state=init_st)
+            out_bwd = wkv7_scan_triton(
+                r_f.flip(1), decay.flip(1), k_f.flip(1),
+                v_f.flip(1), a_f.flip(1), sab_scale,
+                return_state=False).flip(1)
+
+            enc._active_states[idx] = {
+                'wkv_state': wkv_state,
+                'last_x': x[:, -1:].detach().clone(),
+            }
+            out = ((out_fwd + out_bwd) * 0.5).reshape(B, T, C_)
+            out = layer.group_norm(out.transpose(1, 2)).transpose(1, 2)
+            out = layer.W_o(out * g)
+            return out, None
+        return fwd
+
+    @torch.no_grad()
+    def _forward_encode_raw(self, text, init_states=None, max_length=8192):
+        self._hook()
+        if init_states is not None:
+            self._active_states = [
+                {k: v.clone() for k, v in s.items()} if s else None
+                for s in init_states
+            ]
+        else:
+            self._active_states = [None] * len(self.birwkv_layers)
+
+        enc = self.tokenizer(text, return_tensors='pt', truncation=True,
+                             max_length=max_length)
+        ids = enc['input_ids'].to(self.device)
+        mask = enc['attention_mask'].to(self.device)
+
+        h = self.model(input_ids=ids, attention_mask=mask).last_hidden_state
+        content = h[0, 1:-1, :].cpu()
+        n_content = content.shape[0]
+
+        final_states = [
+            {k: v.clone() for k, v in s.items()} if s else None
+            for s in self._active_states
+        ]
+        self._unhook()
+        return content, n_content, final_states
+
+    def _chunk_hidden(self, content, return_residual=False):
+        T = content.shape[0]
+        chunks = []
+        last_end = 0
+        for start in range(0, T, self.chunk_size):
+            end = min(start + self.chunk_size, T)
+            if end - start < 32:
+                break
+            emb = F.normalize(content[start:end].mean(0, keepdim=True),
+                              p=2, dim=-1)
+            chunks.append(emb)
+            last_end = end
+        if not chunks and T > 0:
+            chunks.append(F.normalize(content.mean(0, keepdim=True),
+                                      p=2, dim=-1))
+            last_end = T
+        if return_residual:
+            residual = content[last_end:] if last_end < T else None
+            return chunks, residual
+        return chunks
+
+    @torch.no_grad()
+    def encode_query(self, query):
+        assert not self._hooked
+        enc = self.tokenizer(query, return_tensors='pt', truncation=True,
+                             max_length=512)
+        ids = enc['input_ids'].to(self.device)
+        mask = enc['attention_mask'].to(self.device)
+        h = self.model(input_ids=ids, attention_mask=mask).last_hidden_state
+        m = mask.unsqueeze(-1).float()
+        emb = (h * m).sum(1) / m.sum(1).clamp(min=1e-9)
+        return F.normalize(emb, p=2, dim=-1).cpu()
+
+    def encode_span(self, text, key):
+        content, n_tok, states = self._forward_encode_raw(text)
+        chunks, residual = self._chunk_hidden(content, return_residual=True)
+        self.span_data[key] = {
+            'layer_states': states,
+            'chunk_embs': chunks,
+            'n_tokens': n_tok,
+            'residual_hidden': residual,
+        }
+        return n_tok
+
+    def extend_right(self, piece_text, old_key, new_key):
+        old = self.span_data.pop(old_key)
+        content, n_new, states = self._forward_encode_raw(
+            piece_text, init_states=old['layer_states'])
+        if old.get('residual_hidden') is not None:
+            content = torch.cat([old['residual_hidden'], content], dim=0)
+        new_chunks, residual = self._chunk_hidden(
+            content, return_residual=True)
+        self.span_data[new_key] = {
+            'layer_states': states,
+            'chunk_embs': old['chunk_embs'] + new_chunks,
+            'n_tokens': old['n_tokens'] + n_new,
+            'residual_hidden': residual,
+        }
+        return n_new
+
+    def smart_merge(self, new_text, left_key, new_key):
+        left = self.span_data.pop(left_key)
+        self.remove_old(new_key)
+        content, n_new, states = self._forward_encode_raw(
+            new_text, init_states=left['layer_states'])
+        if left.get('residual_hidden') is not None:
+            content = torch.cat([left['residual_hidden'], content], dim=0)
+        new_chunks, residual = self._chunk_hidden(
+            content, return_residual=True)
+        self.span_data[new_key] = {
+            'layer_states': states,
+            'chunk_embs': left['chunk_embs'] + new_chunks,
+            'n_tokens': left['n_tokens'] + n_new,
+            'residual_hidden': residual,
+        }
+        return n_new
+
+    def remove_old(self, new_key):
+        s, e = new_key
+        for old in list(self.span_data.keys()):
+            if old[0] >= s and old[1] <= e:
+                del self.span_data[old]
+
+    def search(self, q_emb, spans, top_k=5):
+        results = []
+        for s, e, text in spans:
+            key = (s, e)
+            data = self.span_data.get(key)
+            if not data or not data['chunk_embs']:
+                continue
+            chunk_mat = torch.cat(data['chunk_embs'], dim=0)
+            sims = (q_emb @ chunk_mat.T).squeeze(0)
+            if sims.dim() == 0:
+                sims = sims.unsqueeze(0)
+            max_sim = sims.max().item()
+            best_idx = sims.argmax().item()
+            n_chunks = len(data['chunk_embs'])
+            chars_per_chunk = len(text) // max(n_chunks, 1)
+            offset = min(best_idx * chars_per_chunk, len(text) - 1)
+            while offset > 0 and text[offset - 1] not in ' \n\t':
+                offset -= 1
+            preview = text[offset:offset + 300].replace('\n', ' ').strip()
+            results.append((s, e, max_sim, preview, data['n_tokens'], n_chunks))
+        results.sort(key=lambda x: x[2], reverse=True)
+        return results[:top_k]
+
+
+class TextProvider:
+
+    def __init__(self, text, piece_size=4096, seed=42):
+        self.text = text
+        self.piece_size = piece_size
+        self.n_pieces = (len(text) + piece_size - 1) // piece_size
+        self.received = [False] * self.n_pieces
+        rng = random.Random(seed)
+        self.arrival_order = list(range(self.n_pieces))
+        rng.shuffle(self.arrival_order)
+        self.next_idx = 0
+
+    def poll_pieces(self):
+        if self.next_idx >= self.n_pieces:
+            return []
+        idx = self.arrival_order[self.next_idx]
+        self.received[idx] = True
+        self.next_idx += 1
+        return [idx]
+
+    def get_spans(self):
+        spans = []
+        i = 0
+        while i < self.n_pieces:
+            if self.received[i]:
+                j = i
+                while j < self.n_pieces and self.received[j]:
+                    j += 1
+                s_byte = i * self.piece_size
+                e_byte = min(j * self.piece_size, len(self.text))
+                spans.append((i, j, self.text[s_byte:e_byte]))
+                i = j
+            else:
+                i += 1
+        return spans
+
+    def piece_text(self, idx):
+        s = idx * self.piece_size
+        return self.text[s:min(s + self.piece_size, len(self.text))]
+
+    def span_text(self, start_piece, end_piece):
+        s = start_piece * self.piece_size
+        e = min(end_piece * self.piece_size, len(self.text))
+        return self.text[s:e]
+
+    def progress(self):
+        return self.next_idx / self.n_pieces
+
+    def is_complete(self):
+        return self.next_idx >= self.n_pieces
+
+
+FRANKENSTEIN_EXCERPT = """\
+I am by birth a Genevese; and my family is one of the most distinguished \
+of that republic. My ancestors had been for many years counsellors and \
+syndics; and my father had filled several public situations with honour \
+and reputation.
+
+When I was thirteen years of age, we all went on a party of pleasure to \
+the baths near Thonon: the inclemency of the weather obliged us to remain \
+a day confined to the inn. In this house I found a volume of the works of \
+Cornelius Agrippa. I opened it with apathy; the theory which he attempts \
+to demonstrate, and the wonderful facts which he relates, soon changed \
+this feeling into enthusiasm. A new light seemed to dawn upon my mind.
+
+When I returned home, my first care was to procure the whole works of \
+this author. My father was not scientific, and I was left to struggle \
+with a child's blindness, added to a student's thirst for knowledge. \
+Under the guidance of my new preceptors, I entered with the greatest \
+diligence into the search of the philosopher's stone and the elixir \
+of life. What glory would attend the discovery, if I could banish \
+disease from the human frame, and render man invulnerable to any but \
+a violent death!
+
+It was on a dreary night of November that I beheld the accomplishment \
+of my toils. With an anxiety that almost amounted to agony, I collected \
+the instruments of life around me, that I might infuse a spark of being \
+into the lifeless thing that lay at my feet. It was already one in the \
+morning; the rain pattered dismally against the panes, and my candle was \
+nearly burnt out, when, by the glimmer of the half-extinguished light, \
+I saw the dull yellow eye of the creature open; it breathed hard, and \
+a convulsive motion agitated its limbs.
+
+How can I describe my emotions at this catastrophe, or how delineate the \
+wretch whom with such infinite pains and care I had endeavoured to form? \
+I had selected his features as beautiful. Beautiful!--Great God! His \
+yellow skin scarcely covered the work of muscles and arteries beneath; \
+his hair was of a lustrous black, and flowing; his teeth of a pearly \
+whiteness; but these luxuriances only formed a more horrid contrast with \
+his watery eyes, that seemed almost of the same colour as the dun white \
+sockets in which they were set, his shrivelled complexion, and straight \
+black lips.
+
+I had worked hard for nearly two years, for the sole purpose of infusing \
+life into an inanimate body. For this I had deprived myself of rest and \
+health. I had desired it with an ardour that far exceeded moderation; but \
+now that I had finished, the beauty of the dream vanished, and breathless \
+horror and disgust filled my heart.
+
+I did not dare return to the apartment which I inhabited, but felt \
+impelled to hurry on, although drenched by the rain which poured from a \
+black and comfortless sky. I passed the night wretchedly. Morning, \
+dismal and wet, at length dawned, and discovered to my sleepless and \
+aching eyes the church of Ingolstadt, its white steeple and clock, \
+which indicated the sixth hour.
+
+"I shall satiate my ardour for destruction," the creature said, "and \
+make you so wretched that the light of day will be hateful to you. I \
+will be with you on your wedding-night." I started forward, and \
+exclaimed, "Villain! before you sign my death-warrant, be sure that \
+you are yourself safe." My rage was without bounds; I would have seized \
+him; but he eluded me, and quitted the house with precipitation.
+
+Great God! why did I not then expire! But I am a wretch, and none ever \
+conceived of the horrors of my secret toil, whilst I dabbled among the \
+unhallowed damps of the grave, or tortured the living animal to animate \
+the lifeless clay.
+
+I was soon borne away by the waves, and lost in darkness and distance. \
+Immense and rugged mountains of ice often barred up my passage, and I \
+heard the thunder of the ground sea beneath. The cold is excessive, and \
+many of my unfortunate comrades have already found a grave amidst this \
+scene of desolation. Frankenstein! he is not here: I will not rest; I \
+pursue him still over the untrodden snow and frozen ocean.
+"""
+
+QUICK_DEMOS = {
+    "Frankenstein (excerpt)": {
+        "text": FRANKENSTEIN_EXCERPT,
+        "queries": [
+            "the creature opens its eyes for the first time",
+            "playing god with science",
+            "a threat on the wedding night",
+            "a frozen arctic wasteland",
+        ],
+        "piece_size": 512,
+        "sleep": 0.3,
+    },
+}
+
+
+def render_grid(received, n_pieces, highlight=None):
+    max_width = 60
+    if n_pieces <= max_width:
+        cells = []
+        for i in range(n_pieces):
+            if i == highlight:
+                bg = '#00ff41'
+            elif received[i]:
+                bg = '#28a745'
+            else:
+                bg = '#3a3a3a'
+            cells.append(
+                f'<span style="display:inline-block;width:14px;height:22px;'
+                f'background:{bg};margin:1px;border-radius:2px"></span>'
+            )
+    else:
+        cells = []
+        for col in range(max_width):
+            s = col * n_pieces // max_width
+            e = (col + 1) * n_pieces // max_width
+            ratio = sum(received[s:e]) / max(1, e - s)
+            hl = highlight is not None and s <= highlight < e
+            if hl:
+                bg = '#00ff41'
+            elif ratio > 0.8:
+                bg = '#28a745'
+            elif ratio > 0.3:
+                bg = '#17a2b8'
+            elif ratio > 0:
+                bg = '#6c757d'
+            else:
+                bg = '#3a3a3a'
+            cells.append(
+                f'<span style="display:inline-block;width:14px;height:22px;'
+                f'background:{bg};margin:1px;border-radius:2px"></span>'
+            )
+
+    n_recv = sum(received)
+    pct = n_recv / max(n_pieces, 1) * 100
+    grid = ''.join(cells)
+    return (
+        f'<div style="font-family:monospace;line-height:1.4;padding:8px 0">'
+        f'<div style="display:flex;flex-wrap:wrap;gap:0">{grid}</div>'
+        f'<div style="margin-top:8px;color:#aaa">'
+        f'Piece {n_recv}/{n_pieces} ({pct:.0f}%)</div></div>'
+    )
+
+
+def render_search(results_dict, peak_scores=None):
+    if not results_dict:
+        return '<p style="color:#888">Waiting for data...</p>'
+
+    def _score_color(score):
+        if score > 0.5:
+            return '#28a745'
+        elif score > 0.4:
+            return '#ffc107'
+        return '#aaa'
+
+    parts = []
+    for query, results in results_dict.items():
+        peak = peak_scores.get(query) if peak_scores else None
+        header = f'&quot;{query}&quot;'
+        if peak:
+            header += (f' <span style="color:#888;font-size:0.85em">'
+                       f'(peak: {peak["score"]:.3f})</span>')
+        parts.append(
+            f'<div style="margin-bottom:16px">'
+            f'<div style="font-weight:bold;color:#58a6ff;margin-bottom:6px">'
+            f'{header}</div>'
+        )
+
+        cur_best = results[0]['score'] if results else 0
+        if peak and peak['score'] > cur_best + 0.01:
+            psc = _score_color(peak['score'])
+            pp = peak['preview'][:300].replace('<', '&lt;').replace('>', '&gt;')
+            parts.append(
+                f'<div style="padding:4px 0 4px 12px;border-left:3px solid {psc};'
+                f'background:rgba(40,167,69,0.08);margin-bottom:2px">'
+                f'<span style="color:{psc};font-weight:bold">{peak["score"]:.3f}</span> '
+                f'<span style="color:#888;font-size:0.85em">peak</span><br>'
+                f'<span style="color:#ccc;font-size:0.9em">{pp}...</span>'
+                f'</div>'
+            )
+
+        if not results:
+            parts.append('<div style="color:#888;padding-left:12px">No results yet</div>')
+        else:
+            for rank, r in enumerate(results[:3], 1):
+                sc = _score_color(r['score'])
+                preview = r['preview'][:300].replace('<', '&lt;').replace('>', '&gt;')
+                parts.append(
+                    f'<div style="padding:4px 0 4px 12px;border-left:3px solid {sc}">'
+                    f'<span style="color:{sc};font-weight:bold">{r["score"]:.3f}</span> '
+                    f'<span style="color:#888">[{r["span"][0]}-{r["span"][1]}]'
+                    f' ({r["n_chunks"]}ch)</span><br>'
+                    f'<span style="color:#ccc;font-size:0.9em">{preview}...</span>'
+                    f'</div>'
+                )
+        parts.append('</div>')
+    return ''.join(parts)
+
+
+def _state_color(intensity):
+    h = int(220 - intensity * 170)
+    s = int(20 + intensity * 70)
+    light = int(12 + intensity * 38)
+    return f'hsl({h},{s}%,{light}%)'
+
+
+def render_state_viz(state_history, n_layers=14):
+    if not state_history:
+        return ('<p style="color:#888">Recurrent state evolution will appear '
+                'as pieces are processed...</p>')
+
+    n_steps = len(state_history)
+    cell_w = max(4, min(14, 600 // max(n_steps, 1)))
+
+    layer_maxes = []
+    for li in range(n_layers):
+        vals = [state_history[t][li] for t in range(n_steps)
+                if li < len(state_history[t])]
+        layer_maxes.append(max(vals) if vals else 1.0)
+
+    rows = []
+    for li in range(n_layers):
+        cells = []
+        for t in range(n_steps):
+            if li < len(state_history[t]):
+                norm = state_history[t][li]
+                intensity = min(norm / max(layer_maxes[li], 1e-6), 1.0)
+                cells.append(
+                    f'<span style="display:inline-block;width:{cell_w}px;'
+                    f'height:12px;background:{_state_color(intensity)};'
+                    f'margin:0 1px"></span>')
+        rows.append(
+            f'<div style="display:flex;align-items:center;margin:0">'
+            f'<span style="width:24px;color:#666;font-size:9px;'
+            f'text-align:right;margin-right:3px;flex-shrink:0">R{li+1}</span>'
+            f'<div style="display:flex">{"".join(cells)}</div>'
+            f'</div>')
+
+    latest = state_history[-1]
+    avg_norm = sum(latest) / len(latest) if latest else 0
+
+    most_active = 0
+    max_delta = 0
+    if len(state_history) >= 2:
+        prev = state_history[-2]
+        for li in range(min(len(latest), len(prev))):
+            d = abs(latest[li] - prev[li])
+            if d > max_delta:
+                max_delta = d
+                most_active = li
+
+    legend = ''.join(
+        f'<span style="display:inline-block;width:16px;height:8px;'
+        f'background:{_state_color(i / 4)};margin:0 1px"></span>'
+        for i in range(5))
+
+    return (
+        f'<div style="font-family:monospace;line-height:1.1">'
+        f'{"".join(rows)}'
+        f'<div style="color:#777;font-size:10px;margin-top:6px">'
+        f'{n_layers} RWKV layers \u00d7 {n_steps} pieces | '
+        f'Avg state magnitude: {avg_norm:.1f}'
+        f'{f" | Most active: R{most_active+1}" if len(state_history) >= 2 else ""}'
+        f'</div>'
+        f'<div style="color:#666;font-size:9px;margin-top:2px">'
+        f'{legend} low \u2192 high state magnitude'
+        f'</div></div>')
+
+
+def load_text(url):
+    resp = urllib.request.urlopen(url, timeout=30)
+    text = resp.read().decode('utf-8', errors='replace')
+    start = text.find('*** START OF')
+    if start != -1:
+        text = text[text.find('\n', start) + 1:]
+    end = text.find('*** END OF')
+    if end != -1:
+        text = text[:end]
+    return text
+
+
+def streaming_loop(provider, encoder, queries, q_embs, sleep_time=0):
+    prev_span_keys = set()
+    hare_tokens = 0
+    baseline_tokens = 0
+    right_extends = 0
+    smart_merges = 0
+    full_reencodes = 0
+    merge_events = 0
+    pieces_processed = 0
+    piece_queue = []
+    peak_scores = {}
+    state_history = []
+    n_rwkv_layers = len(encoder.birwkv_layers)
+
+    while not provider.is_complete():
+        new_pieces = provider.poll_pieces()
+        if new_pieces:
+            piece_queue.extend(new_pieces)
+            random.shuffle(piece_queue)
+
+        if not piece_queue:
+            continue
+
+        idx = piece_queue.pop(0)
+        provider.received[idx] = True
+        pieces_processed += 1
+
+        new_spans = provider.get_spans()
+        new_keys = {(s, e) for s, e, _ in new_spans}
+
+        for s, e, span_text_val in new_spans:
+            key = (s, e)
+            if key in prev_span_keys:
+                continue
+
+            right_key = (s, e - 1)
+            if right_key in encoder.span_data:
+                n = encoder.extend_right(provider.piece_text(e - 1), right_key, key)
+                hare_tokens += n
+                right_extends += 1
+                baseline_tokens += encoder.span_data[key]['n_tokens']
+                continue
+
+            best_left = None
+            for (os_, oe) in list(encoder.span_data.keys()):
+                if os_ == s and oe < e:
+                    if best_left is None or oe > best_left[1]:
+                        best_left = (os_, oe)
+
+            if best_left:
+                new_portion = provider.span_text(best_left[1], e)
+                n = encoder.smart_merge(new_portion, best_left, key)
+                hare_tokens += n
+                smart_merges += 1
+                baseline_tokens += encoder.span_data[key]['n_tokens']
+                continue
+
+            encoder.remove_old(key)
+            n = encoder.encode_span(span_text_val, key)
+            hare_tokens += n
+            full_reencodes += 1
+            baseline_tokens += n
+
+        if len(new_keys) < len(prev_span_keys) and pieces_processed > 1:
+            merge_events += 1
+        prev_span_keys = new_keys
+
+        total_chunks = sum(len(d['chunk_embs']) for d in encoder.span_data.values())
+        eff = baseline_tokens / max(hare_tokens, 1)
+
+        if encoder.span_data:
+            largest_key = max(encoder.span_data.keys(),
+                              key=lambda k: k[1] - k[0])
+            states = encoder.span_data[largest_key].get('layer_states', [])
+            norms = []
+            for st in states:
+                if st is not None and 'wkv_state' in st:
+                    norms.append(st['wkv_state'].norm().item())
+                else:
+                    norms.append(0.0)
+            state_history.append(norms)
+
+        search_results = {}
+        for q in queries:
+            results = encoder.search(q_embs[q], new_spans, top_k=3)
+            search_results[q] = [
+                {'span': (s, e), 'score': sc, 'preview': pv,
+                 'n_chunks': nc, 'n_tokens': nt}
+                for s, e, sc, pv, nt, nc in results
+            ]
+            if results:
+                top = results[0]
+                sc_top = top[2]
+                if q not in peak_scores or sc_top > peak_scores[q]['score']:
+                    peak_scores[q] = {'score': sc_top, 'preview': top[3]}
+
+        grid_html = render_grid(provider.received, provider.n_pieces, highlight=idx)
+        saved = baseline_tokens - hare_tokens
+        eff_md = f"**Efficiency: {eff:.1f}x** | {total_chunks} chunks"
+        tok_md = f"Tokens: {hare_tokens:,} processed | {saved:,} saved via state carry"
+        strat_md = (f"Right-ext: {right_extends} | Smart-merge: {smart_merges} | "
+                    f"Full: {full_reencodes} | Merges: {merge_events}")
+        search_html = render_search(search_results, peak_scores)
+        state_html = render_state_viz(state_history, n_rwkv_layers)
+
+        yield grid_html, eff_md, tok_md, strat_md, search_html, state_html
+
+        if sleep_time > 0:
+            time.sleep(sleep_time)
+
+    eff = baseline_tokens / max(hare_tokens, 1)
+    total_chunks = sum(len(d['chunk_embs']) for d in encoder.span_data.values())
+    saved = baseline_tokens - hare_tokens
+    grid_html = render_grid(provider.received, provider.n_pieces)
+    eff_md = f"**Efficiency: {eff:.1f}x** | {total_chunks} chunks | COMPLETE"
+    tok_md = f"Tokens: {hare_tokens:,} processed | {saved:,} saved via state carry"
+    strat_md = (f"Right-ext: {right_extends} | Smart-merge: {smart_merges} | "
+                f"Full: {full_reencodes} | Merges: {merge_events}")
+
+    final_spans = provider.get_spans()
+    search_results = {}
+    for q in queries:
+        results = encoder.search(q_embs[q], final_spans, top_k=3)
+        search_results[q] = [
+            {'span': (s, e), 'score': sc, 'preview': pv,
+             'n_chunks': nc, 'n_tokens': nt}
+            for s, e, sc, pv, nt, nc in results
+        ]
+    search_html = render_search(search_results, peak_scores)
+    state_html = render_state_viz(state_history, n_rwkv_layers)
+    yield grid_html, eff_md, tok_md, strat_md, search_html, state_html
+
+
+@spaces.GPU
+def start_demo(source_mode, demo_choice, url_input, queries_text, chunk_size):
+    model.cuda()
+    encoder = SpanEncoder(model, tokenizer, chunk_size=chunk_size)
+
+    if source_mode == "Quick Demo":
+        config = QUICK_DEMOS[demo_choice]
+        provider = TextProvider(config['text'],
+                                piece_size=config['piece_size'], seed=42)
+        queries = config['queries']
+        sleep_time = config['sleep']
+    elif source_mode == "URL":
+        if not url_input:
+            yield ('<p style="color:#ffc107">Enter a URL to a text file.</p>',
+                   '', '', '', '', '')
+            return
+        text = load_text(url=url_input)
+        provider = TextProvider(text, piece_size=4096, seed=42)
+        queries = [q.strip() for q in queries_text.split(',') if q.strip()]
+        sleep_time = 0
+    else:
+        return
+
+    if not queries:
+        queries = ["search query"]
+
+    q_embs = {q: encoder.encode_query(q) for q in queries}
+
+    yield from streaming_loop(provider, encoder, queries, q_embs, sleep_time)
+
+
+def toggle_inputs(source_mode):
+    frankenstein_q = "on a dreary night the creature first opened its eyes, an innocent woman is wrongly executed, playing god with science"
+    return (
+        gr.update(visible=(source_mode == "Quick Demo")),
+        gr.update(visible=(source_mode == "URL")),
+        gr.update(visible=(source_mode != "Quick Demo"),
+                  value=frankenstein_q),
+    )
+
+
+def update_queries(demo_choice):
+    config = QUICK_DEMOS.get(demo_choice, {})
+    queries = config.get('queries', [])
+    return ', '.join(queries)
+
+
+def build_demo():
+    with gr.Blocks(title="HARE Streaming Demo") as demo:
+        gr.Markdown(
+            "# HARE: Streaming Semantic Search",
+        )
+        gr.Markdown(
+            "Watch [HARE](https://huggingface.co/SixOpen/HARE) build a "
+            "semantic search index in real-time as content streams in "
+            "piece by piece. Unlike standard embedding models, HARE's "
+            "recurrent state carries forward full context without "
+            "re-encoding, allowing for search over live transcripts, "
+            "distributed content, and streaming files without "
+            "needing to download them in full.",
+        )
+
+        with gr.Row():
+            with gr.Column(scale=1, min_width=280):
+                source_mode = gr.Radio(
+                    ["URL", "Quick Demo"],
+                    value="URL",
+                    label="Source",
+                )
+                demo_choice = gr.Dropdown(
+                    list(QUICK_DEMOS.keys()),
+                    value=list(QUICK_DEMOS.keys())[0],
+                    label="Demo Content",
+                    visible=False,
+                )
+                url_input = gr.Textbox(
+                    label="Text URL",
+                    value="https://gutenberg.org/files/84/84-0.txt",
+                    placeholder="https://gutenberg.org/files/84/84-0.txt",
+                    visible=True,
+                )
+                queries_input = gr.Textbox(
+                    label="Search Queries (comma-separated)",
+                    value="on a dreary night the creature first opened its eyes, an innocent woman is wrongly executed, playing god with science",
+                    visible=True,
+                )
+
+                with gr.Accordion("Settings", open=False):
+                    chunk_size = gr.Slider(
+                        128, 1024, value=512, step=64,
+                        label="Chunk Size (tokens)",
+                    )
+
+                start_btn = gr.Button("Start Demo", variant="primary", size="lg")
+
+            with gr.Column(scale=2):
+                gr.Markdown("### Download Progress")
+                piece_grid = gr.HTML(
+                    '<div style="padding:20px;color:#666;text-align:center">'
+                    'Click "Start Demo" to begin</div>'
+                )
+
+                gr.Markdown("### Encoding Efficiency")
+                with gr.Row():
+                    efficiency_md = gr.Markdown("**Efficiency: --**")
+                with gr.Row():
+                    tokens_md = gr.Markdown("Tokens: --")
+                    strategy_md = gr.Markdown("Right-ext: -- | Smart-merge: -- | Full: --")
+
+                gr.Markdown("### Search Results")
+                search_html = gr.HTML(
+                    '<p style="color:#888">Results will appear here as '
+                    'pieces are processed...</p>'
+                )
+
+                gr.Markdown("### Recurrent State Evolution")
+                state_viz = gr.HTML(
+                    '<p style="color:#888">State heatmap will appear as '
+                    'pieces are processed...</p>'
+                )
+
+        source_mode.change(
+            toggle_inputs,
+            inputs=[source_mode],
+            outputs=[demo_choice, url_input, queries_input],
+        )
+        demo_choice.change(
+            update_queries,
+            inputs=[demo_choice],
+            outputs=[queries_input],
+        )
+        start_btn.click(
+            start_demo,
+            inputs=[source_mode, demo_choice, url_input, queries_input,
+                    chunk_size],
+            outputs=[piece_grid, efficiency_md, tokens_md, strategy_md,
+                     search_html, state_viz],
+        )
+
+    return demo
+
+
+demo = build_demo()
+demo.queue().launch()