Create Model_Active.py

Model_Active.py

@@ -0,0 +1,163 @@
+import math
+import logging
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+logger = logging.getLogger(__name__)
+
+
+class RWKV_TimeMix(nn.Module):
+    def __init__(self, config, layer_id):
+        super().__init__()
+        assert config.n_attn % config.n_head == 0
+        self.layer_id = layer_id
+        self.ctx_len = config.ctx_len
+        self.n_head = config.n_head
+        self.head_size = config.n_attn // config.n_head
+
+        self.time_ww = nn.Parameter(
+            torch.ones(config.n_head, config.ctx_len, config.ctx_len))
+        self.time_gamma = nn.Parameter(torch.ones(config.ctx_len, 1))
+
+        self.time_shift = nn.ZeroPad2d((0, 0, 1, -1))
+
+        self.key = nn.Linear(config.n_embd, config.n_attn)
+        self.value = nn.Linear(config.n_embd, config.n_attn)
+        self.receptance = nn.Linear(config.n_embd, config.n_attn)
+
+        self.output = nn.Linear(config.n_attn, config.n_embd)
+
+        self.key.scale_init = 0
+        self.receptance.scale_init = 0
+        self.output.scale_init = 0
+
+    def forward(self, x):
+        B, T, C = x.size()
+
+        x = torch.cat(
+            [self.time_shift(x[:, :, :C//2]), x[:, :, C//2:]], dim=-1)
+
+        k = self.key(x)
+        v = self.value(x)
+        r = self.receptance(x)
+
+        k = torch.clamp(k, max=30, min=-60)
+        k = torch.exp(k)
+        sum_k = torch.cumsum(k, dim=1)
+
+        kv = (k * v).view(B, T, self.n_head, self.head_size)
+
+        wkv = (torch.einsum('htu,buhc->bthc', self.time_ww[:,:T,:T], kv)
+               ).contiguous().view(B, T, -1)
+
+        rwkv = torch.sigmoid(r) * wkv / sum_k
+
+        rwkv = self.output(rwkv)
+        return rwkv * self.time_gamma[:T, :]
+
+class RWKV_ChannelMix(nn.Module):
+    def __init__(self, config, layer_id):
+        super().__init__()
+        self.layer_id = layer_id
+        self.time_shift = nn.ZeroPad2d((0, 0, 1, -1))
+
+        hidden_sz = 5 * config.n_ffn // 2
+        self.key = nn.Linear(config.n_embd, hidden_sz)
+        self.value = nn.Linear(config.n_embd, hidden_sz)
+        self.weight = nn.Linear(hidden_sz, config.n_embd)
+        self.receptance = nn.Linear(config.n_embd, config.n_embd)
+
+        self.receptance.scale_init = 0
+        self.weight.scale_init = 0
+
+    def forward(self, x):
+        B, T, C = x.size()
+
+        x = torch.cat(
+            [self.time_shift(x[:, :, :C//2]), x[:, :, C//2:]], dim=-1)
+        k = self.key(x)
+        v = self.value(x)
+        r = self.receptance(x)
+
+        wkv = self.weight(F.mish(k) * v)
+
+        rwkv = torch.sigmoid(r) * wkv
+
+        return rwkv
+
+
+class GPTConfig:
+    def __init__(self, vocab_size, ctx_len, **kwargs):
+        self.vocab_size = vocab_size
+        self.ctx_len = ctx_len
+        for k, v in kwargs.items():
+            setattr(self, k, v)
+
+
+class Block(nn.Module):
+    def __init__(self, config, layer_id):
+        super().__init__()
+        self.config = config
+
+        self.ln1 = nn.LayerNorm(config.n_embd)
+        self.ln2 = nn.LayerNorm(config.n_embd)
+
+        self.attn = RWKV_TimeMix(config, layer_id)
+        self.mlp = RWKV_ChannelMix(config, layer_id)
+
+    def forward(self, x):
+
+        x = x + self.attn(self.ln1(x))
+        x = x + self.mlp(self.ln2(x))
+
+        return x
+
+
+class GPT(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+
+        self.tok_emb = nn.Embedding(config.vocab_size, config.n_embd)
+
+        self.blocks = nn.Sequential(*[Block(config, i)
+                                    for i in range(config.n_layer)])
+
+        self.ln_f = nn.LayerNorm(config.n_embd)
+        self.time_out = nn.Parameter(torch.ones(1, config.ctx_len, 1))
+        self.head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+
+        self.head_q = nn.Linear(config.n_embd, 256)
+        self.head_k = nn.Linear(config.n_embd, 256)
+        self.register_buffer("copy_mask", torch.tril(torch.ones(config.ctx_len, config.ctx_len)))
+
+        self.ctx_len = config.ctx_len
+
+        logger.info("number of parameters: %e", sum(p.numel()
+                    for p in self.parameters()))
+
+    def get_ctx_len(self):
+        return self.ctx_len
+
+    def forward(self, idx, targets=None):
+        B, T = idx.size()
+        assert T <= self.ctx_len, "Cannot forward, because len(input) > model ctx_len."
+
+        x = self.tok_emb(idx)
+
+        x = self.blocks(x)
+
+        x = self.ln_f(x)
+        q = self.head_q(x)[:,:T,:]
+        k = self.head_k(x)[:,:T,:]
+        c = (q @ k.transpose(-2, -1)) * (1.0 / 256)
+        c = c.masked_fill(self.copy_mask[:T,:T] == 0, 0)
+        c = c @ F.one_hot(idx, num_classes = self.config.vocab_size).float()       
+        x = x * self.time_out[:, :T, :]
+        x = self.head(x) + c
+
+        loss = None
+        if targets is not None:
+            loss = F.cross_entropy(x.view(-1, x.size(-1)), targets.view(-1))
+
+        return x, loss