Spaces:
Paused
Paused
| import math | |
| import time | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import gradio as gr | |
| import spaces | |
| from transformers import ( | |
| AutoTokenizer, | |
| PretrainedConfig, | |
| PreTrainedModel, | |
| GenerationMixin, | |
| ) | |
| from transformers.modeling_outputs import CausalLMOutputWithPast | |
| class FWKVConfig(PretrainedConfig): | |
| """Configuration class for FWKV-ROSA model architecture.""" | |
| 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 | |
| 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 | |
| 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 | |
| class FactorizedTiedHead(nn.Module): | |
| """Factorized embedding projection and tied output head.""" | |
| 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) | |
| class FWKVBlock(nn.Module): | |
| """FWKV layer with linear time attention-style recurrent mechanism.""" | |
| 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) | |
| 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 | |
| class FWKVLanguageModel(PreTrainedModel, GenerationMixin): | |
| """Full causal language model utilizing FWKV recurrent layers and ROSA embeddings.""" | |
| 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: | |
| 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} | |
| 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("FWKV/FWKV-ROSA") | |
| model = model.to(device) | |
| model.eval() | |
| tokenizer = AutoTokenizer.from_pretrained("FWKV/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() | |
| def generate_reply_stream(ids: list[int], max_new_tokens=150, temperature=0.8, top_k=50): | |
| """Autoregressive generation with ROSA updates, yielding token lists and comprehensive throughput metrics.""" | |
| device = next(model.parameters()).device | |
| eos_id = tokenizer.eos_token_id | |
| # Initial forward pass over the 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) | |
| reply_tokens = [] | |
| start_time = time.perf_counter() | |
| prev_step_time = start_time | |
| instant_tps_list = [] | |
| 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 | |
| reply_tokens.append(next_token) | |
| now = time.perf_counter() | |
| # Calculate per-step instant duration and speed | |
| step_duration = now - prev_step_time | |
| prev_step_time = now | |
| if step_duration > 0: | |
| instant_tps = 1.0 / step_duration | |
| instant_tps_list.append(instant_tps) | |
| # Compute aggregate throughput metrics | |
| total_elapsed = now - start_time | |
| avg_tps = len(reply_tokens) / total_elapsed if total_elapsed > 0 else 0.0 | |
| current_tps = instant_tps_list[-1] if instant_tps_list else avg_tps | |
| min_tps = min(instant_tps_list) if instant_tps_list else avg_tps | |
| max_tps = max(instant_tps_list) if instant_tps_list else avg_tps | |
| stats = { | |
| "current": current_tps, | |
| "avg": avg_tps, | |
| "min": min_tps, | |
| "max": max_tps, | |
| } | |
| yield reply_tokens, stats | |
| # ROSA prediction for the next step | |
| 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] | |
| def extract_text_content(content) -> str: | |
| """Safely extract plain text from string, list, or dict content structures returned by Gradio.""" | |
| if isinstance(content, str): | |
| return content | |
| if isinstance(content, list): | |
| parts = [] | |
| for item in content: | |
| if isinstance(item, str): | |
| parts.append(item) | |
| elif isinstance(item, dict): | |
| if "text" in item: | |
| parts.append(str(item["text"])) | |
| elif "content" in item: | |
| parts.append(extract_text_content(item["content"])) | |
| else: | |
| parts.append(str(item)) | |
| return " ".join(parts) | |
| if isinstance(content, dict): | |
| if "text" in content: | |
| return str(content["text"]) | |
| return str(content) | |
| return str(content) if content is not None else "" | |
| def chat_function(message, history): | |
| """Gradio ChatInterface streaming handler formatted with speed stats (Live, Avg, Min, Max).""" | |
| messages = [] | |
| for turn in history: | |
| if isinstance(turn, (list, tuple)): | |
| user_msg, asst_msg = turn | |
| messages.append({"role": "user", "content": extract_text_content(user_msg)}) | |
| if asst_msg: | |
| messages.append({"role": "assistant", "content": extract_text_content(asst_msg)}) | |
| elif isinstance(turn, dict): | |
| messages.append({ | |
| "role": turn.get("role", "user"), | |
| "content": extract_text_content(turn.get("content", "")) | |
| }) | |
| messages.append({"role": "user", "content": extract_text_content(message)}) | |
| # Encode token sequence according to model chat template | |
| 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"] | |
| if not content.strip(): | |
| continue | |
| content_ids = tokenizer.encode(" " + content) | |
| if role == "user": | |
| ids += [user_id] + content_ids | |
| elif role == "assistant": | |
| ids += [asst_id] + content_ids + [eos_id] | |
| # Truncate left if context exceeds model max sequence length | |
| max_len = model.config.seq_len if model else 1024 | |
| if len(ids) > max_len: | |
| ids = ids[-max_len:] | |
| # Prompt assistant response | |
| ids.append(asst_id) | |
| # Stream generated output with full throughput statistics | |
| for reply_tokens, stats in generate_reply_stream(ids, max_new_tokens=150, temperature=0.8, top_k=50): | |
| reply = tokenizer.decode(reply_tokens, skip_special_tokens=True).strip() | |
| metrics_bar = ( | |
| f"⚡ **{stats['current']:.1f} tok/s** " | |
| f"*(Avg: **{stats['avg']:.1f}** | Min: **{stats['min']:.1f}** | Max: **{stats['max']:.1f}** tok/s)*" | |
| ) | |
| yield f"{reply}\n\n{metrics_bar}" | |
| with gr.Blocks(theme=gr.themes.Soft()) as demo: | |
| gr.Markdown(f""" | |
| # ⚡ FWKV-ROSA Chat | |
| **Model:** [FWKV/FWKV-ROSA](https://huggingface.co/FWKV/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.", | |
| ], | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() |