Create app.py

app.py

@@ -0,0 +1,372 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import gradio as gr
+from transformers import (
+    AutoTokenizer,
+    PretrainedConfig,
+    PreTrainedModel,
+    GenerationMixin,
+)
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+# -------------------------------------------------------------------------
+# Model components (must match the saved checkpoint exactly)
+# -------------------------------------------------------------------------
+
+class FWKVConfig(PretrainedConfig):
+    model_type = "fwkv"
+
+    def __init__(
+        self,
+        d_model: int = 512,
+        d_emb: int = 128,
+        n_layers: int = 14,
+        ffn_mult: int = 4,
+        vocab_size: int = 50257,
+        seq_len: int = 1024,        # trained with 1024
+        wkv_floor: float = 0.1,
+        tie_word_embeddings: bool = True,
+        **kwargs,
+    ):
+        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
+        self.d_model = d_model
+        self.d_emb = d_emb
+        self.n_layers = n_layers
+        self.ffn_mult = ffn_mult
+        self.vocab_size = vocab_size
+        self.seq_len = seq_len
+        self.wkv_floor = wkv_floor
+
+
+# ── ROSA (exact copy from training) ─────────────────────────────────────
+def rosa(x: list[int]) -> list[int]:
+    """Causal copy‑signal predictor; returns y[i] = token after longest
+    repeating suffix ending at i, or -1 if none."""
+    n = len(x)
+    if n == 0:
+        return []
+    y = [-1] * n
+    s = 2 * n + 2
+    trans = [None] * s
+    link = [-1] * s
+    length = [0] * s
+    last_end = [-1] * s
+    trans[0] = {}
+    last = 0
+    size = 1
+
+    for i, t in enumerate(x):
+        cur = size; size += 1
+        trans[cur] = {}
+        length[cur] = length[last] + 1
+        p = last
+        while p != -1 and t not in trans[p]:
+            trans[p][t] = cur
+            p = link[p]
+        if p == -1:
+            link[cur] = 0
+        else:
+            q = trans[p][t]
+            if length[p] + 1 == length[q]:
+                link[cur] = q
+            else:
+                clone = size; size += 1
+                trans[clone] = trans[q].copy()
+                length[clone] = length[p] + 1
+                link[clone] = link[q]
+                last_end[clone] = last_end[q]
+                while p != -1 and trans[p][t] == q:
+                    trans[p][t] = clone
+                    p = link[p]
+                link[q] = clone
+                link[cur] = clone
+        last = cur
+
+        v = cur
+        pred = -1
+        while v != -1:
+            if length[v] > 0 and last_end[v] >= 0:
+                pred = x[last_end[v] + 1]
+                break
+            v = link[v]
+        y[i] = pred
+
+        v = last
+        while v != -1 and last_end[v] < i:
+            last_end[v] = i
+            v = link[v]
+    return y
+
+
+# ── Vectorised parallel scan ─────────────────────────────────────────────
+def parallel_scan_decay(a: torch.Tensor, W: torch.Tensor) -> torch.Tensor:
+    """Hillis–Steele inclusive scan with constant per‑channel decay."""
+    W = W.to(dtype=a.dtype)          # keep precision
+    val = a
+    T = a.shape[1]
+    d = 1
+    while d < T:
+        shifted = F.pad(val[:, :-d, :], (0, 0, d, 0))
+        val = val + (W ** d) * shifted
+        d *= 2
+    return val
+
+
+# ── Factorised tied head ─────────────────────────────────────────────────
+class FactorizedTiedHead(nn.Module):
+    def __init__(self, vocab_size: int, d_model: int, d_emb: int):
+        super().__init__()
+        self.d_model = d_model
+        self.d_emb = d_emb
+        self.weight = nn.Parameter(torch.empty(vocab_size, d_emb))
+        self.proj = nn.Linear(d_emb, d_model, bias=False)
+
+    def embed(self, input_ids):
+        return self.proj(F.embedding(input_ids, self.weight))
+
+    def to_emb_space(self, x):
+        return F.linear(x, self.proj.weight.t())
+
+    def logits(self, x_emb):
+        return F.linear(x_emb, self.weight)
+
+
+# ── FWKV Block (with parallel scan) ──────────────────────────────────────
+class FWKVBlock(nn.Module):
+    def __init__(self, d: int, ffn_mult: int = 4, floor: float = 0.1):
+        super().__init__()
+        self.floor = floor
+        self.proj_k = nn.Linear(d, d, bias=False)
+        self.proj_v = nn.Linear(d, d, bias=False)
+        self.proj_r = nn.Linear(d, d, bias=False)
+        self.proj_out = nn.Linear(d, d, bias=False)
+        self.w = nn.Parameter(torch.ones(d) * 2.0)
+        self.ffn = nn.Sequential(
+            nn.Linear(d, ffn_mult * d, bias=False),
+            nn.GELU(),
+            nn.Linear(ffn_mult * d, d, bias=False),
+        )
+        self.norm_wkv = nn.LayerNorm(d)
+        self.norm_ffn = nn.LayerNorm(d)
+
+    @property
+    def W(self):
+        return torch.clamp(torch.sigmoid(self.w), min=self.floor)
+
+    def forward(self, x, state=None):
+        B, T, d = x.shape
+        W = self.W
+        k = self.proj_k(x)
+        v = self.proj_v(x)
+        r = torch.sigmoid(self.proj_r(x))
+
+        a = k * v
+        if state is not None:
+            a = a.clone()
+            a[:, 0] = a[:, 0] + W * state
+
+        wkv_out = parallel_scan_decay(a, W)
+        new_state = wkv_out[:, -1].detach()
+
+        x = self.norm_wkv(x + self.proj_out(r * wkv_out))
+        x = self.norm_ffn(x + self.ffn(x))
+        return x, new_state
+
+
+# ── Full Language Model ──────────────────────────────────────────────────
+class FWKVLanguageModel(PreTrainedModel, GenerationMixin):
+    config_class = FWKVConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.shared = FactorizedTiedHead(config.vocab_size, config.d_model, config.d_emb)
+        self.rosa_emb = nn.Embedding(config.vocab_size + 1, config.d_emb, padding_idx=0)
+        self.blocks = nn.ModuleList([
+            FWKVBlock(config.d_model, config.ffn_mult, config.wkv_floor)
+            for _ in range(config.n_layers)
+        ])
+        self.norm = nn.LayerNorm(config.d_model)
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.shared.weight
+
+    def forward(
+        self,
+        input_ids,
+        rosa_ids=None,
+        past_key_values=None,
+        labels=None,
+        use_cache=True,
+        **kwargs,
+    ):
+        if rosa_ids is None:
+            # fallback: compute on the fly (expensive, but safe)
+            rows = [rosa(row.tolist()) for row in input_ids.detach().cpu()]
+            rosa_ids = torch.tensor(rows, device=input_ids.device, dtype=torch.long)
+
+        x = self.shared.embed(input_ids)
+        rosa_idx = (rosa_ids + 1).clamp(min=0)
+        x = x + self.shared.proj(self.rosa_emb(rosa_idx))
+
+        states_in = past_key_values or [None] * len(self.blocks)
+        states_out = []
+        for block, state in zip(self.blocks, states_in):
+            x, new_state = block(x, state)
+            states_out.append(new_state)
+
+        x = self.norm(x)
+        x_emb = self.shared.to_emb_space(x)
+        logits = self.shared.logits(x_emb)
+
+        return CausalLMOutputWithPast(
+            loss=None,
+            logits=logits,
+            past_key_values=states_out if use_cache else None,
+        )
+
+    def prepare_inputs_for_generation(self, input_ids, past_key_values=None,
+                                      rosa_ids=None, **kwargs):
+        if past_key_values is not None:
+            input_ids = input_ids[:, -1:]
+            if rosa_ids is not None:
+                rosa_ids = rosa_ids[:, -1:]
+        return {"input_ids": input_ids, "rosa_ids": rosa_ids,
+                "past_key_values": past_key_values, "use_cache": True}
+
+
+# -------------------------------------------------------------------------
+# Loading & Chat helpers
+# -------------------------------------------------------------------------
+
+USER_TOKEN = "<|user|>"
+ASSISTANT_TOKEN = "<|assistant|>"
+
+def load_model():
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    print(f"Loading FWKV-ROSA from Hub on {device} ...")
+    try:
+        model = FWKVLanguageModel.from_pretrained("FlameF0X/FWKV-ROSA")
+        model = model.to(device)
+        model.eval()
+        tokenizer = AutoTokenizer.from_pretrained("FlameF0X/FWKV-ROSA")
+        status = "FWKV-ROSA chat model ready!"
+    except Exception as e:
+        model, tokenizer = None, None
+        status = f"Error loading model: {e}"
+        print(status)
+    return model, tokenizer, status
+
+
+model, tokenizer, load_status = load_model()
+
+@torch.no_grad()
+def generate_reply(ids: list[int], max_new_tokens=150, temperature=0.8, top_k=50):
+    """Autoregressive generation with ROSA updates, returning full sequence."""
+    device = next(model.parameters()).device
+    eos_id = tokenizer.eos_token_id
+
+    # initial forward over the whole prompt
+    inp = torch.tensor([ids], device=device)
+    rosa_ids = torch.tensor([rosa(ids)], device=device)
+    out = model(input_ids=inp, rosa_ids=rosa_ids, use_cache=True)
+    states = out.past_key_values
+    logits = out.logits[0, -1]
+    generated = list(ids)
+
+    for _ in range(max_new_tokens):
+        scaled = logits / max(temperature, 1e-5)
+        if top_k and top_k < scaled.size(-1):
+            kth = torch.topk(scaled, top_k).values[-1]
+            scaled[scaled < kth] = float('-inf')
+        probs = torch.softmax(scaled, dim=-1)
+        next_token = torch.multinomial(probs, 1).item()
+        generated.append(next_token)
+        if next_token == eos_id:
+            break
+
+        # ROSA for the extended sequence, use the last prediction
+        next_rosa = rosa(generated)[-1]
+        step_inp = torch.tensor([[next_token]], device=device)
+        step_rosa = torch.tensor([[next_rosa]], device=device)
+        out = model(input_ids=step_inp, rosa_ids=step_rosa,
+                    past_key_values=states, use_cache=True)
+        states = out.past_key_values
+        logits = out.logits[0, -1]
+
+    return generated
+
+
+def chat_function(message, history):
+    """Gradio ChatInterface callback. history is a list of (user, assistant) pairs."""
+    # Build the full conversation in the chat template
+    messages = []
+    for user_msg, asst_msg in history:
+        messages.append({"role": "user", "content": user_msg})
+        messages.append({"role": "assistant", "content": asst_msg})
+    messages.append({"role": "user", "content": message})
+
+    # Convert to token IDs
+    user_id = tokenizer.convert_tokens_to_ids(USER_TOKEN)
+    asst_id = tokenizer.convert_tokens_to_ids(ASSISTANT_TOKEN)
+    eos_id = tokenizer.eos_token_id
+    ids = []
+    for turn in messages:
+        role = turn["role"]
+        content = turn["content"]
+        content_ids = tokenizer.encode(" " + content)
+        if role == "user":
+            ids += [user_id] + content_ids
+        elif role == "assistant":
+            ids += [asst_id] + content_ids + [eos_id]
+
+    # Truncate from the left if needed
+    max_len = model.config.seq_len
+    if len(ids) > max_len:
+        ids = ids[-max_len:]
+
+    # Cue the assistant to start replying
+    ids.append(asst_id)
+
+    # Generate the assistant's reply
+    full_gen = generate_reply(ids, max_new_tokens=150, temperature=0.8, top_k=50)
+
+    # Extract only the new assistant tokens (after the last asst_id)
+    # Find the position of the last asst_id and take everything after it
+    assistant_start = len(ids) - 1  # index of the asst_id we just appended
+    reply_ids = full_gen[assistant_start + 1:]  # skip the asst_id itself
+    reply = tokenizer.decode(reply_ids, skip_special_tokens=True).strip()
+    return reply
+
+
+# ── Gradio UI ────────────────────────────────────────────────────────────
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown(f"""
+    # ⚡ FWKV-ROSA Chat
+    **Model:** [FlameF0X/FWKV-ROSA](https://huggingface.co/FlameF0X/FWKV-ROSA)
+    *{load_status}*
+
+    This is a 56M‑parameter recurrent LM trained with the RWKV‑8 ROSA
+    copy‑signal mechanism. It uses the chat template:
+
+        `<|user|> message <|assistant|> reply <eos>`
+
+    You can chat naturally; the model will remember recent context up to
+    {model.config.seq_len if model else 1024} tokens.
+    """)
+
+    chatbot = gr.ChatInterface(
+        fn=chat_function,
+        title="",
+        description="",
+        examples=[
+            "Explain how a linear recurrent network can still copy long‑range patterns.",
+            "Write a short poem about a fox discovering a hidden library.",
+        ],
+        theme="soft",
+    )
+
+if __name__ == "__main__":
+    demo.launch()