streaming.py

import torch
import torch.nn.functional as F

from birwkv7 import BiRWKV7Layer


def wkv7_forward_scan(r, w, k, v, a, sab_scale, init_state=None):
    B, T, H, D = r.shape
    r, w, k, v, a = [x.float() for x in (r, w, k, v, a)]
    k = k * (D ** -0.5)
    decay = torch.exp(-0.6065306597633104 * torch.sigmoid(w))
    a = torch.sigmoid(a)
    sab_s = float(sab_scale)
    state = init_state.float().clone() if init_state is not None else \
        torch.zeros(B, H, D, D, device=r.device, dtype=torch.float32)
    outputs = []
    for t in range(T):
        kt, vt, rt, at, dt = k[:, t], v[:, t], r[:, t], a[:, t], decay[:, t]
        sa = torch.einsum('bhij,bhj->bhi', state, -kt)
        sab = torch.einsum('bhi,bhj->bhij', sa, kt * at)
        state = state * dt.unsqueeze(-2) + sab_s * sab + \
            torch.einsum('bhi,bhj->bhij', vt, kt)
        state = state.clamp(-10.0, 10.0)
        outputs.append(torch.einsum('bhij,bhj->bhi', state, rt))
    return torch.stack(outputs, dim=1), state.detach()


class SpanEncoder:

    def __init__(self, model, tokenizer, device, chunk_size=512):
        self.model = model
        self.tokenizer = tokenizer
        self.device = device
        self.chunk_size = chunk_size

        self.birwkv_layers = []
        self.birwkv_ids = {}
        for m in model.modules():
            if isinstance(m, BiRWKV7Layer):
                self.birwkv_ids[id(m)] = len(self.birwkv_layers)
                self.birwkv_layers.append(m)

        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

            try:
                from birwkv7_triton import wkv7_scan_triton
                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)
            except (ImportError, Exception):
                out_fwd, wkv_state = wkv7_forward_scan(
                    r, w, k, v, a, sab_scale, init_st)
                out_bwd = wkv7_forward_scan(
                    r.flip(1), w.flip(1), k.flip(1),
                    v.flip(1), a.flip(1), sab_scale, None)[0].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