import torch import torch.nn as nn import torch.nn.functional as F import json import os # ====================================================================== # RWKV-Mamba Hybrid Recurrence # ====================================================================== class RWKVMambaHybrid(nn.Module): def __init__(self, d_model, d_state=64): super().__init__() self.d_model = d_model self.d_state = d_state self.w_mix = nn.Parameter(torch.ones(d_model) * 0.5) self.A = nn.Parameter(torch.randn(d_state, d_state) * 0.01) self.B = nn.Parameter(torch.randn(d_state, d_model) * 0.01) self.C = nn.Parameter(torch.randn(d_model, d_state) * 0.01) self.D = nn.Parameter(torch.ones(d_model) * 0.1) def forward(self, x): B, T, C = x.shape h = torch.zeros(B, C, device=x.device) s = torch.zeros(B, self.d_state, device=x.device) outputs = [] for t in range(T): x_t = x[:, t, :] h = self.w_mix * h + (1 - self.w_mix) * x_t s = s @ self.A.T + x_t @ self.B.T y_t = s @ self.C.T + h * self.D outputs.append(y_t) return torch.stack(outputs, dim=1) # ====================================================================== # Full Multi-Head Attention # ====================================================================== class FullAttention(nn.Module): def __init__(self, d_model, n_heads=16): super().__init__() self.d_model = d_model self.n_heads = n_heads self.head_dim = d_model // n_heads assert d_model % n_heads == 0, "d_model must be divisible by n_heads" self.qkv = nn.Linear(d_model, d_model * 3) self.out_proj = nn.Linear(d_model, d_model) def forward(self, x, mask=None): B, T, C = x.shape qkv = self.qkv(x) q, k, v = qkv.chunk(3, dim=-1) q = q.view(B, T, self.n_heads, self.head_dim).transpose(1, 2) k = k.view(B, T, self.n_heads, self.head_dim).transpose(1, 2) v = v.view(B, T, self.n_heads, self.head_dim).transpose(1, 2) attn = (q @ k.transpose(-2, -1)) / (self.head_dim ** 0.5) if mask is not None: mask = mask.expand(B, self.n_heads, T, T).bool() attn = attn.masked_fill(mask == 0, float('-inf')) attn = F.softmax(attn, dim=-1) out = attn @ v out = out.transpose(1, 2).contiguous().view(B, T, C) return self.out_proj(out) # ====================================================================== # i3 Hybrid Block # ====================================================================== class i3HybridBlock(nn.Module): def __init__(self, d_model, d_state=64, ffn_mult=4): super().__init__() self.ln1 = nn.LayerNorm(d_model) self.hybrid = RWKVMambaHybrid(d_model, d_state) self.ln2 = nn.LayerNorm(d_model) d_ff = d_model * ffn_mult self.ffn = nn.Sequential( nn.Linear(d_model, d_ff), nn.GELU(), nn.Linear(d_ff, d_model) ) def forward(self, x, mask=None): x = x + self.hybrid(self.ln1(x)) x = x + self.ffn(self.ln2(x)) return x # ====================================================================== # i3 Attention Block # ====================================================================== class i3AttentionBlock(nn.Module): def __init__(self, d_model, n_heads=16, ffn_mult=4): super().__init__() self.ln1 = nn.LayerNorm(d_model) self.attn = FullAttention(d_model, n_heads) self.ln2 = nn.LayerNorm(d_model) d_ff = d_model * ffn_mult self.ffn = nn.Sequential( nn.Linear(d_model, d_ff), nn.GELU(), nn.Linear(d_ff, d_model) ) def forward(self, x, mask=None): x = x + self.attn(self.ln1(x), mask) x = x + self.ffn(self.ln2(x)) return x # ====================================================================== # Full i3 Model # ====================================================================== class i3Model(nn.Module): def __init__(self, vocab_size, d_model=512, n_heads=16, max_seq_len=256, d_state=32): super().__init__() self.vocab_size = vocab_size self.d_model = d_model self.max_seq_len = max_seq_len self.embed = nn.Embedding(vocab_size, d_model) self.pos_embed = nn.Embedding(max_seq_len, d_model) hybrid_layers = [i3HybridBlock(d_model, d_state=d_state) for _ in range(10)] attention_layers = [i3AttentionBlock(d_model, n_heads=n_heads) for _ in range(6)] self.layers = nn.ModuleList(hybrid_layers + attention_layers) self.ln_f = nn.LayerNorm(d_model) self.head = nn.Linear(d_model, vocab_size) self.apply(self._init_weights) def _init_weights(self, module): if isinstance(module, (nn.Linear, nn.Embedding)): module.weight.data.normal_(0.0, 0.02) if isinstance(module, nn.Linear) and module.bias is not None: module.bias.data.zero_() def forward(self, idx, targets=None): B, T = idx.shape pos = torch.arange(T, device=idx.device).unsqueeze(0) x = self.embed(idx) + self.pos_embed(pos) mask = torch.tril(torch.ones(T, T, device=idx.device)).view(1, 1, T, T) for layer in self.layers: x = layer(x, mask) x = self.ln_f(x) logits = self.head(x) loss = None if targets is not None: loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1)) return logits, loss @torch.no_grad() def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None): for _ in range(max_new_tokens): idx_cond = idx if idx.size(1) <= self.max_seq_len else idx[:, -self.max_seq_len:] logits, _ = self(idx_cond) logits = logits[:, -1, :] / temperature if top_k is not None: v, _ = torch.topk(logits, min(top_k, logits.size(-1))) logits[logits < v[:, [-1]]] = -float('Inf') probs = F.softmax(logits, dim=-1) idx_next = torch.multinomial(probs, 1) idx = torch.cat((idx, idx_next), dim=1) return idx # ====================================================================== # ChunkTokenizer # ====================================================================== class ChunkTokenizer: def __init__(self, vocab_only=False): self.chunk_to_idx = {} self.idx_to_chunk = {} self.vocab_size = 0 self.unk_token = '' self.unk_idx = 0 self.common_trigrams = frozenset({ 'the', 'and', 'ing', 'ion', 'tio', 'for', 'tha', 'ter', 'hat', 'his', 'ere', 'ent', 'her', 'was', 'you', 'are', 'not', 'but', 'can', 'all', 'whi', 'one', 'our', 'out', 'whe', 'hav', 'thi', 'wit' }) def _should_use_3char(self, pos, text): if pos + 3 > len(text): return False chunk_3 = text[pos:pos+3] if chunk_3 in self.common_trigrams: return True if pos > 0 and text[pos-1] == ' ': return True if pos + 3 < len(text) and text[pos+3] == ' ': return True return False def encode(self, text): text = text.lower() pos, indices = 0, [] while pos < len(text): chunk_len = 3 if self._should_use_3char(pos, text) else 2 chunk_len = min(chunk_len, len(text) - pos) chunk = text[pos:pos+chunk_len] if chunk in self.chunk_to_idx: indices.append(self.chunk_to_idx[chunk]) else: indices.append(self.unk_idx) pos += chunk_len return indices def decode(self, indices): return ''.join([self.idx_to_chunk.get(int(i), self.unk_token) for i in indices]) def save(self, path): data = { 'chunk_to_idx': self.chunk_to_idx, 'idx_to_chunk': {int(k): v for k, v in self.idx_to_chunk.items()}, 'vocab_size': self.vocab_size, 'unk_token': self.unk_token, 'unk_idx': self.unk_idx } with open(path, 'w') as f: json.dump(data, f) def load(self, path): with open(path, 'r') as f: data = json.load(f) self.chunk_to_idx = data['chunk_to_idx'] self.idx_to_chunk = {int(k): v for k, v in data['idx_to_chunk'].items()} self.vocab_size = data['vocab_size'] self.unk_token = data.get('unk_token', '') self.unk_idx = data.get('unk_idx', 0)