Create i3_model.py

i3_model.py

@@ -0,0 +1,230 @@
+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 = '<UNK>'
+        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', '<UNK>')
+        self.unk_idx = data.get('unk_idx', 0)