wop/FWKV-ROSA
2
1import math2import time3import torch4import torch.nn as nn5import torch.nn.functional as F6import gradio as gr7import spaces8from transformers import (9 AutoTokenizer,10 PretrainedConfig,11 PreTrainedModel,12 GenerationMixin,13)14from transformers.modeling_outputs import CausalLMOutputWithPast15 16class FWKVConfig(PretrainedConfig):17 """Configuration class for FWKV-ROSA model architecture."""18 model_type = "fwkv"19 20 def __init__(21 self,22 d_model: int = 512,23 d_emb: int = 128,24 n_layers: int = 14,25 ffn_mult: int = 4,26 vocab_size: int = 50257,27 seq_len: int = 1024, # trained with 102428 wkv_floor: float = 0.1,29 tie_word_embeddings: bool = True,30 **kwargs,31 ):32 super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)33 self.d_model = d_model34 self.d_emb = d_emb35 self.n_layers = n_layers36 self.ffn_mult = ffn_mult37 self.vocab_size = vocab_size38 self.seq_len = seq_len39 self.wkv_floor = wkv_floor40 41def rosa(x: list[int]) -> list[int]:42 """Causal copy‑signal predictor; returns y[i] = token after longest43 repeating suffix ending at i, or -1 if none."""44 n = len(x)45 if n == 0:46 return []47 y = [-1] * n48 s = 2 * n + 249 trans = [None] * s50 link = [-1] * s51 length = [0] * s52 last_end = [-1] * s53 trans[0] = {}54 last = 055 size = 156 57 for i, t in enumerate(x):58 cur = size; size += 159 trans[cur] = {}60 length[cur] = length[last] + 161 p = last62 while p != -1 and t not in trans[p]:63 trans[p][t] = cur64 p = link[p]65 if p == -1:66 link[cur] = 067 else:68 q = trans[p][t]69 if length[p] + 1 == length[q]:70 link[cur] = q71 else:72 clone = size; size += 173 trans[clone] = trans[q].copy()74 length[clone] = length[p] + 175 link[clone] = link[q]76 last_end[clone] = last_end[q]77 while p != -1 and trans[p][t] == q:78 trans[p][t] = clone79 p = link[p]80 link[q] = clone81 link[cur] = clone82 last = cur83 84 v = cur85 pred = -186 while v != -1:87 if length[v] > 0 and last_end[v] >= 0:88 pred = x[last_end[v] + 1]89 break90 v = link[v]91 y[i] = pred92 93 v = last94 while v != -1 and last_end[v] < i:95 last_end[v] = i96 v = link[v]97 return y98 99def parallel_scan_decay(a: torch.Tensor, W: torch.Tensor) -> torch.Tensor:100 """Hillis–Steele inclusive scan with constant per‑channel decay."""101 W = W.to(dtype=a.dtype) # keep precision102 val = a103 T = a.shape[1]104 d = 1105 while d < T:106 shifted = F.pad(val[:, :-d, :], (0, 0, d, 0))107 val = val + (W ** d) * shifted108 d *= 2109 return val110 111class FactorizedTiedHead(nn.Module):112 """Factorized embedding projection and tied output head."""113 def __init__(self, vocab_size: int, d_model: int, d_emb: int):114 super().__init__()115 self.d_model = d_model116 self.d_emb = d_emb117 self.weight = nn.Parameter(torch.empty(vocab_size, d_emb))118 self.proj = nn.Linear(d_emb, d_model, bias=False)119 120 def embed(self, input_ids):121 return self.proj(F.embedding(input_ids, self.weight))122 123 def to_emb_space(self, x):124 return F.linear(x, self.proj.weight.t())125 126 def logits(self, x_emb):127 return F.linear(x_emb, self.weight)128 129class FWKVBlock(nn.Module):130 """FWKV layer with linear time attention-style recurrent mechanism."""131 def __init__(self, d: int, ffn_mult: int = 4, floor: float = 0.1):132 super().__init__()133 self.floor = floor134 self.proj_k = nn.Linear(d, d, bias=False)135 self.proj_v = nn.Linear(d, d, bias=False)136 self.proj_r = nn.Linear(d, d, bias=False)137 self.proj_out = nn.Linear(d, d, bias=False)138 self.w = nn.Parameter(torch.ones(d) * 2.0)139 self.ffn = nn.Sequential(140 nn.Linear(d, ffn_mult * d, bias=False),141 nn.GELU(),142 nn.Linear(ffn_mult * d, d, bias=False),143 )144 self.norm_wkv = nn.LayerNorm(d)145 self.norm_ffn = nn.LayerNorm(d)146 147 @property148 def W(self):149 return torch.clamp(torch.sigmoid(self.w), min=self.floor)150 151 def forward(self, x, state=None):152 B, T, d = x.shape153 W = self.W154 k = self.proj_k(x)155 v = self.proj_v(x)156 r = torch.sigmoid(self.proj_r(x))157 158 a = k * v159 if state is not None:160 a = a.clone()161 a[:, 0] = a[:, 0] + W * state162 163 wkv_out = parallel_scan_decay(a, W)164 new_state = wkv_out[:, -1].detach()165 166 x = self.norm_wkv(x + self.proj_out(r * wkv_out))167 x = self.norm_ffn(x + self.ffn(x))168 return x, new_state169 170class FWKVLanguageModel(PreTrainedModel, GenerationMixin):171 """Full causal language model utilizing FWKV recurrent layers and ROSA embeddings."""172 config_class = FWKVConfig173 174 def __init__(self, config):175 super().__init__(config)176 self.shared = FactorizedTiedHead(config.vocab_size, config.d_model, config.d_emb)177 self.rosa_emb = nn.Embedding(config.vocab_size + 1, config.d_emb, padding_idx=0)178 self.blocks = nn.ModuleList([179 FWKVBlock(config.d_model, config.ffn_mult, config.wkv_floor)180 for _ in range(config.n_layers)181 ])182 self.norm = nn.LayerNorm(config.d_model)183 self.post_init()184 185 def get_input_embeddings(self):186 return self.shared.weight187 188 def forward(189 self,190 input_ids,191 rosa_ids=None,192 past_key_values=None,193 labels=None,194 use_cache=True,195 **kwargs,196 ):197 if rosa_ids is None:198 rows = [rosa(row.tolist()) for row in input_ids.detach().cpu()]199 rosa_ids = torch.tensor(rows, device=input_ids.device, dtype=torch.long)200 201 x = self.shared.embed(input_ids)202 rosa_idx = (rosa_ids + 1).clamp(min=0)203 x = x + self.shared.proj(self.rosa_emb(rosa_idx))204 205 states_in = past_key_values or [None] * len(self.blocks)206 states_out = []207 for block, state in zip(self.blocks, states_in):208 x, new_state = block(x, state)209 states_out.append(new_state)210 211 x = self.norm(x)212 x_emb = self.shared.to_emb_space(x)213 logits = self.shared.logits(x_emb)214 215 return CausalLMOutputWithPast(216 loss=None,217 logits=logits,218 past_key_values=states_out if use_cache else None,219 )220 221 def prepare_inputs_for_generation(self, input_ids, past_key_values=None,222 rosa_ids=None, **kwargs):223 if past_key_values is not None:224 input_ids = input_ids[:, -1:]225 if rosa_ids is not None:226 rosa_ids = rosa_ids[:, -1:]227 return {"input_ids": input_ids, "rosa_ids": rosa_ids,228 "past_key_values": past_key_values, "use_cache": True}229 230USER_TOKEN = "<|user|>"231ASSISTANT_TOKEN = "<|assistant|>"232 233def load_model():234 device = "cuda" if torch.cuda.is_available() else "cpu"235 print(f"Loading FWKV-ROSA from Hub on {device} ...")236 try:237 model = FWKVLanguageModel.from_pretrained("FWKV/FWKV-ROSA")238 model = model.to(device)239 model.eval()240 tokenizer = AutoTokenizer.from_pretrained("FWKV/FWKV-ROSA")241 status = "FWKV-ROSA chat model ready!"242 except Exception as e:243 model, tokenizer = None, None244 status = f"Error loading model: {e}"245 print(status)246 return model, tokenizer, status247 248 249model, tokenizer, load_status = load_model()250 251@spaces.GPU252@torch.no_grad()253def generate_reply_stream(ids: list[int], max_new_tokens=150, temperature=0.8, top_k=50):254 """Autoregressive generation with ROSA updates, yielding token lists and comprehensive throughput metrics."""255 device = next(model.parameters()).device256 eos_id = tokenizer.eos_token_id257 258 # Initial forward pass over the prompt259 inp = torch.tensor([ids], device=device)260 rosa_ids = torch.tensor([rosa(ids)], device=device)261 out = model(input_ids=inp, rosa_ids=rosa_ids, use_cache=True)262 states = out.past_key_values263 logits = out.logits[0, -1]264 generated = list(ids)265 reply_tokens = []266 267 start_time = time.perf_counter()268 prev_step_time = start_time269 instant_tps_list = []270 271 for _ in range(max_new_tokens):272 scaled = logits / max(temperature, 1e-5)273 if top_k and top_k < scaled.size(-1):274 kth = torch.topk(scaled, top_k).values[-1]275 scaled[scaled < kth] = float('-inf')276 probs = torch.softmax(scaled, dim=-1)277 next_token = torch.multinomial(probs, 1).item()278 generated.append(next_token)279 if next_token == eos_id:280 break281 282 reply_tokens.append(next_token)283 now = time.perf_counter()284 285 # Calculate per-step instant duration and speed286 step_duration = now - prev_step_time287 prev_step_time = now288 289 if step_duration > 0:290 instant_tps = 1.0 / step_duration291 instant_tps_list.append(instant_tps)292 293 # Compute aggregate throughput metrics294 total_elapsed = now - start_time295 avg_tps = len(reply_tokens) / total_elapsed if total_elapsed > 0 else 0.0296 current_tps = instant_tps_list[-1] if instant_tps_list else avg_tps297 min_tps = min(instant_tps_list) if instant_tps_list else avg_tps298 max_tps = max(instant_tps_list) if instant_tps_list else avg_tps299 300 stats = {301 "current": current_tps,302 "avg": avg_tps,303 "min": min_tps,304 "max": max_tps,305 }306 307 yield reply_tokens, stats308 309 # ROSA prediction for the next step310 next_rosa = rosa(generated)[-1]311 step_inp = torch.tensor([[next_token]], device=device)312 step_rosa = torch.tensor([[next_rosa]], device=device)313 out = model(input_ids=step_inp, rosa_ids=step_rosa,314 past_key_values=states, use_cache=True)315 states = out.past_key_values316 logits = out.logits[0, -1]317 318def extract_text_content(content) -> str:319 """Safely extract plain text from string, list, or dict content structures returned by Gradio."""320 if isinstance(content, str):321 return content322 if isinstance(content, list):323 parts = []324 for item in content:325 if isinstance(item, str):326 parts.append(item)327 elif isinstance(item, dict):328 if "text" in item:329 parts.append(str(item["text"]))330 elif "content" in item:331 parts.append(extract_text_content(item["content"]))332 else:333 parts.append(str(item))334 return " ".join(parts)335 if isinstance(content, dict):336 if "text" in content:337 return str(content["text"])338 return str(content)339 return str(content) if content is not None else ""340 341def chat_function(message, history):342 """Gradio ChatInterface streaming handler formatted with speed stats (Live, Avg, Min, Max)."""343 messages = []344 for turn in history:345 if isinstance(turn, (list, tuple)):346 user_msg, asst_msg = turn347 messages.append({"role": "user", "content": extract_text_content(user_msg)})348 if asst_msg:349 messages.append({"role": "assistant", "content": extract_text_content(asst_msg)})350 elif isinstance(turn, dict):351 messages.append({352 "role": turn.get("role", "user"),353 "content": extract_text_content(turn.get("content", ""))354 })355 messages.append({"role": "user", "content": extract_text_content(message)})356 357 # Encode token sequence according to model chat template358 user_id = tokenizer.convert_tokens_to_ids(USER_TOKEN)359 asst_id = tokenizer.convert_tokens_to_ids(ASSISTANT_TOKEN)360 eos_id = tokenizer.eos_token_id361 ids = []362 for turn in messages:363 role = turn["role"]364 content = turn["content"]365 if not content.strip():366 continue367 content_ids = tokenizer.encode(" " + content)368 if role == "user":369 ids += [user_id] + content_ids370 elif role == "assistant":371 ids += [asst_id] + content_ids + [eos_id]372 373 # Truncate left if context exceeds model max sequence length374 max_len = model.config.seq_len if model else 1024375 if len(ids) > max_len:376 ids = ids[-max_len:]377 378 # Prompt assistant response379 ids.append(asst_id)380 381 # Stream generated output with full throughput statistics382 for reply_tokens, stats in generate_reply_stream(ids, max_new_tokens=150, temperature=0.8, top_k=50):383 reply = tokenizer.decode(reply_tokens, skip_special_tokens=True).strip()384 metrics_bar = (385 f"⚡ **{stats['current']:.1f} tok/s** "386 f"*(Avg: **{stats['avg']:.1f}** | Min: **{stats['min']:.1f}** | Max: **{stats['max']:.1f}** tok/s)*"387 )388 yield f"{reply}\n\n{metrics_bar}"389 390with gr.Blocks(theme=gr.themes.Soft()) as demo:391 gr.Markdown(f"""392 # ⚡ FWKV-ROSA Chat393 **Model:** [FWKV/FWKV-ROSA](https://huggingface.co/FWKV/FWKV-ROSA)394 *{load_status}*395 396 This is a 56M‑parameter recurrent LM trained with the RWKV‑8 ROSA397 copy‑signal mechanism. It uses the chat template:398 399 `<|user|> message <|assistant|> reply <eos>`400 401 You can chat naturally; the model will remember recent context up to402 {model.config.seq_len if model else 1024} tokens.403 """)404 405 chatbot = gr.ChatInterface(406 fn=chat_function,407 title="",408 description="",409 examples=[410 "Explain how a linear recurrent network can still copy long‑range patterns.",411 "Write a short poem about a fox discovering a hidden library.",412 ],413 )414 415if __name__ == "__main__":416 demo.launch()