FlameF0X/FWKV-ROSA
1
1import math2import time3import torch4import torch.nn as nn5import torch.nn.functional as F6import gradio as gr7from transformers import (8 AutoTokenizer,9 PretrainedConfig,10 PreTrainedModel,11 GenerationMixin,12)13from transformers.modeling_outputs import CausalLMOutputWithPast14 15class FWKVConfig(PretrainedConfig):16 """Configuration class for FWKV-ROSA model architecture."""17 model_type = "fwkv"18 19 def __init__(20 self,21 d_model: int = 512,22 d_emb: int = 128,23 n_layers: int = 14,24 ffn_mult: int = 4,25 vocab_size: int = 50257,26 seq_len: int = 1024, # trained with 102427 wkv_floor: float = 0.1,28 tie_word_embeddings: bool = True,29 **kwargs,30 ):31 super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)32 self.d_model = d_model33 self.d_emb = d_emb34 self.n_layers = n_layers35 self.ffn_mult = ffn_mult36 self.vocab_size = vocab_size37 self.seq_len = seq_len38 self.wkv_floor = wkv_floor39 40def rosa(x: list[int]) -> list[int]:41 """Causal copy‑signal predictor; returns y[i] = token after longest42 repeating suffix ending at i, or -1 if none."""43 n = len(x)44 if n == 0:45 return []46 y = [-1] * n47 s = 2 * n + 248 trans = [None] * s49 link = [-1] * s50 length = [0] * s51 last_end = [-1] * s52 trans[0] = {}53 last = 054 size = 155 56 for i, t in enumerate(x):57 cur = size; size += 158 trans[cur] = {}59 length[cur] = length[last] + 160 p = last61 while p != -1 and t not in trans[p]:62 trans[p][t] = cur63 p = link[p]64 if p == -1:65 link[cur] = 066 else:67 q = trans[p][t]68 if length[p] + 1 == length[q]:69 link[cur] = q70 else:71 clone = size; size += 172 trans[clone] = trans[q].copy()73 length[clone] = length[p] + 174 link[clone] = link[q]75 last_end[clone] = last_end[q]76 while p != -1 and trans[p][t] == q:77 trans[p][t] = clone78 p = link[p]79 link[q] = clone80 link[cur] = clone81 last = cur82 83 v = cur84 pred = -185 while v != -1:86 if length[v] > 0 and last_end[v] >= 0:87 pred = x[last_end[v] + 1]88 break89 v = link[v]90 y[i] = pred91 92 v = last93 while v != -1 and last_end[v] < i:94 last_end[v] = i95 v = link[v]96 return y97 98def parallel_scan_decay(a: torch.Tensor, W: torch.Tensor) -> torch.Tensor:99 """Hillis–Steele inclusive scan with constant per‑channel decay."""100 W = W.to(dtype=a.dtype) # keep precision101 val = a102 T = a.shape[1]103 d = 1104 while d < T:105 shifted = F.pad(val[:, :-d, :], (0, 0, d, 0))106 val = val + (W ** d) * shifted107 d *= 2108 return val109 110class FactorizedTiedHead(nn.Module):111 """Factorized embedding projection and tied output head."""112 def __init__(self, vocab_size: int, d_model: int, d_emb: int):113 super().__init__()114 self.d_model = d_model115 self.d_emb = d_emb116 self.weight = nn.Parameter(torch.empty(vocab_size, d_emb))117 self.proj = nn.Linear(d_emb, d_model, bias=False)118 119 def embed(self, input_ids):120 return self.proj(F.embedding(input_ids, self.weight))121 122 def to_emb_space(self, x):123 return F.linear(x, self.proj.weight.t())124 125 def logits(self, x_emb):126 return F.linear(x_emb, self.weight)127 128class FWKVBlock(nn.Module):129 """FWKV layer with linear time attention-style recurrent mechanism."""130 def __init__(self, d: int, ffn_mult: int = 4, floor: float = 0.1):131 super().__init__()132 self.floor = floor133 self.proj_k = nn.Linear(d, d, bias=False)134 self.proj_v = nn.Linear(d, d, bias=False)135 self.proj_r = nn.Linear(d, d, bias=False)136 self.proj_out = nn.Linear(d, d, bias=False)137 self.w = nn.Parameter(torch.ones(d) * 2.0)138 self.ffn = nn.Sequential(139 nn.Linear(d, ffn_mult * d, bias=False),140 nn.GELU(),141 nn.Linear(ffn_mult * d, d, bias=False),142 )143 self.norm_wkv = nn.LayerNorm(d)144 self.norm_ffn = nn.LayerNorm(d)145 146 @property147 def W(self):148 return torch.clamp(torch.sigmoid(self.w), min=self.floor)149 150 def forward(self, x, state=None):151 B, T, d = x.shape152 W = self.W153 k = self.proj_k(x)154 v = self.proj_v(x)155 r = torch.sigmoid(self.proj_r(x))156 157 a = k * v158 if state is not None:159 a = a.clone()160 a[:, 0] = a[:, 0] + W * state161 162 wkv_out = parallel_scan_decay(a, W)163 new_state = wkv_out[:, -1].detach()164 165 x = self.norm_wkv(x + self.proj_out(r * wkv_out))166 x = self.norm_ffn(x + self.ffn(x))167 return x, new_state168 169class FWKVLanguageModel(PreTrainedModel, GenerationMixin):170 """Full causal language model utilizing FWKV recurrent layers and ROSA embeddings."""171 config_class = FWKVConfig172 173 def __init__(self, config):174 super().__init__(config)175 self.shared = FactorizedTiedHead(config.vocab_size, config.d_model, config.d_emb)176 self.rosa_emb = nn.Embedding(config.vocab_size + 1, config.d_emb, padding_idx=0)177 self.blocks = nn.ModuleList([178 FWKVBlock(config.d_model, config.ffn_mult, config.wkv_floor)179 for _ in range(config.n_layers)180 ])181 self.norm = nn.LayerNorm(config.d_model)182 self.post_init()183 184 def get_input_embeddings(self):185 return self.shared.weight186 187 def forward(188 self,189 input_ids,190 rosa_ids=None,191 past_key_values=None,192 labels=None,193 use_cache=True,194 **kwargs,195 ):196 if rosa_ids is None:197 rows = [rosa(row.tolist()) for row in input_ids.detach().cpu()]198 rosa_ids = torch.tensor(rows, device=input_ids.device, dtype=torch.long)199 200 x = self.shared.embed(input_ids)201 rosa_idx = (rosa_ids + 1).clamp(min=0)202 x = x + self.shared.proj(self.rosa_emb(rosa_idx))203 204 states_in = past_key_values or [None] * len(self.blocks)205 states_out = []206 for block, state in zip(self.blocks, states_in):207 x, new_state = block(x, state)208 states_out.append(new_state)209 210 x = self.norm(x)211 x_emb = self.shared.to_emb_space(x)212 logits = self.shared.logits(x_emb)213 214 return CausalLMOutputWithPast(215 loss=None,216 logits=logits,217 past_key_values=states_out if use_cache else None,218 )219 220 def prepare_inputs_for_generation(self, input_ids, past_key_values=None,221 rosa_ids=None, **kwargs):222 if past_key_values is not None:223 input_ids = input_ids[:, -1:]224 if rosa_ids is not None:225 rosa_ids = rosa_ids[:, -1:]226 return {"input_ids": input_ids, "rosa_ids": rosa_ids,227 "past_key_values": past_key_values, "use_cache": True}228 229USER_TOKEN = "<|user|>"230ASSISTANT_TOKEN = "<|assistant|>"231 232def load_model():233 device = "cuda" if torch.cuda.is_available() else "cpu"234 print(f"Loading FWKV-ROSA from Hub on {device} ...")235 try:236 model = FWKVLanguageModel.from_pretrained("FWKV/FWKV-ROSA")237 model = model.to(device)238 model.eval()239 tokenizer = AutoTokenizer.from_pretrained("FWKV/FWKV-ROSA")240 status = "FWKV-ROSA chat model ready!"241 except Exception as e:242 model, tokenizer = None, None243 status = f"Error loading model: {e}"244 print(status)245 return model, tokenizer, status246 247 248model, tokenizer, load_status = load_model()249 250@torch.no_grad()251def generate_reply_stream(ids: list[int], max_new_tokens=150, temperature=0.8, top_k=50):252 """Autoregressive generation with ROSA updates, yielding token lists and comprehensive throughput metrics."""253 device = next(model.parameters()).device254 eos_id = tokenizer.eos_token_id255 256 # Initial forward pass over the prompt257 inp = torch.tensor([ids], device=device)258 rosa_ids = torch.tensor([rosa(ids)], device=device)259 out = model(input_ids=inp, rosa_ids=rosa_ids, use_cache=True)260 states = out.past_key_values261 logits = out.logits[0, -1]262 generated = list(ids)263 reply_tokens = []264 265 start_time = time.perf_counter()266 prev_step_time = start_time267 instant_tps_list = []268 269 for _ in range(max_new_tokens):270 scaled = logits / max(temperature, 1e-5)271 if top_k and top_k < scaled.size(-1):272 kth = torch.topk(scaled, top_k).values[-1]273 scaled[scaled < kth] = float('-inf')274 probs = torch.softmax(scaled, dim=-1)275 next_token = torch.multinomial(probs, 1).item()276 generated.append(next_token)277 if next_token == eos_id:278 break279 280 reply_tokens.append(next_token)281 now = time.perf_counter()282 283 # Calculate per-step instant duration and speed284 step_duration = now - prev_step_time285 prev_step_time = now286 287 if step_duration > 0:288 instant_tps = 1.0 / step_duration289 instant_tps_list.append(instant_tps)290 291 # Compute aggregate throughput metrics292 total_elapsed = now - start_time293 avg_tps = len(reply_tokens) / total_elapsed if total_elapsed > 0 else 0.0294 current_tps = instant_tps_list[-1] if instant_tps_list else avg_tps295 min_tps = min(instant_tps_list) if instant_tps_list else avg_tps296 max_tps = max(instant_tps_list) if instant_tps_list else avg_tps297 298 stats = {299 "current": current_tps,300 "avg": avg_tps,301 "min": min_tps,302 "max": max_tps,303 }304 305 yield reply_tokens, stats306 307 # ROSA prediction for the next step308 next_rosa = rosa(generated)[-1]309 step_inp = torch.tensor([[next_token]], device=device)310 step_rosa = torch.tensor([[next_rosa]], device=device)311 out = model(input_ids=step_inp, rosa_ids=step_rosa,312 past_key_values=states, use_cache=True)313 states = out.past_key_values314 logits = out.logits[0, -1]315 316def extract_text_content(content) -> str:317 """Safely extract plain text from string, list, or dict content structures returned by Gradio."""318 if isinstance(content, str):319 return content320 if isinstance(content, list):321 parts = []322 for item in content:323 if isinstance(item, str):324 parts.append(item)325 elif isinstance(item, dict):326 if "text" in item:327 parts.append(str(item["text"]))328 elif "content" in item:329 parts.append(extract_text_content(item["content"]))330 else:331 parts.append(str(item))332 return " ".join(parts)333 if isinstance(content, dict):334 if "text" in content:335 return str(content["text"])336 return str(content)337 return str(content) if content is not None else ""338 339def chat_function(message, history):340 """Gradio ChatInterface streaming handler formatted with speed stats (Live, Avg, Min, Max)."""341 messages = []342 for turn in history:343 if isinstance(turn, (list, tuple)):344 user_msg, asst_msg = turn345 messages.append({"role": "user", "content": extract_text_content(user_msg)})346 if asst_msg:347 messages.append({"role": "assistant", "content": extract_text_content(asst_msg)})348 elif isinstance(turn, dict):349 messages.append({350 "role": turn.get("role", "user"),351 "content": extract_text_content(turn.get("content", ""))352 })353 messages.append({"role": "user", "content": extract_text_content(message)})354 355 # Encode token sequence according to model chat template356 user_id = tokenizer.convert_tokens_to_ids(USER_TOKEN)357 asst_id = tokenizer.convert_tokens_to_ids(ASSISTANT_TOKEN)358 eos_id = tokenizer.eos_token_id359 ids = []360 for turn in messages:361 role = turn["role"]362 content = turn["content"]363 if not content.strip():364 continue365 content_ids = tokenizer.encode(" " + content)366 if role == "user":367 ids += [user_id] + content_ids368 elif role == "assistant":369 ids += [asst_id] + content_ids + [eos_id]370 371 # Truncate left if context exceeds model max sequence length372 max_len = model.config.seq_len if model else 1024373 if len(ids) > max_len:374 ids = ids[-max_len:]375 376 # Prompt assistant response377 ids.append(asst_id)378 379 # Stream generated output with full throughput statistics380 for reply_tokens, stats in generate_reply_stream(ids, max_new_tokens=150, temperature=0.8, top_k=50):381 reply = tokenizer.decode(reply_tokens, skip_special_tokens=True).strip()382 metrics_bar = (383 f"⚡ **{stats['current']:.1f} tok/s** "384 f"*(Avg: **{stats['avg']:.1f}** | Min: **{stats['min']:.1f}** | Max: **{stats['max']:.1f}** tok/s)*"385 )386 yield f"{reply}\n\n{metrics_bar}"387 388with gr.Blocks(theme=gr.themes.Soft()) as demo:389 gr.Markdown(f"""390 # ⚡ FWKV-ROSA Chat391 **Model:** [FKWV/FWKV-ROSA](https://huggingface.co/FWKV/FWKV-ROSA)392 *{load_status}*393 394 This is a 56M‑parameter recurrent LM trained with the RWKV‑8 ROSA395 copy‑signal mechanism. It uses the chat template:396 397 `<|user|> message <|assistant|> reply <eos>`398 399 You can chat naturally; the model will remember recent context up to400 {model.config.seq_len if model else 1024} tokens.401 """)402 403 chatbot = gr.ChatInterface(404 fn=chat_function,405 title="",406 description="",407 examples=[408 "Explain how a linear recurrent network can still copy long‑range patterns.",409 "Write a short poem about a fox discovering a hidden library.",410 ],411 )412 413if __name__ == "__main__":414 demo.launch()