cogforge-engine 1.1.0__tar.gz → 2.0.0__tar.gz

{cogforge_engine-1.1.0 → cogforge_engine-2.0.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cogforge-engine
-Version: 1.1.0
+Version: 2.0.0
 Summary: A custom autograd engine and Transformer block built from scratch.
 Project-URL: Homepage, https://github.com/avikmjd2/cogforge
 Author-email: Avik Majumder <avikmjd2@gmail.com>

cogforge_engine-2.0.0/TODO

@@ -0,0 +1,7 @@
+GPT-2 Residual Initialization Scaling
+
+Weight decay (decay weights not biases and layer norm)
+
+cosine LR function
+
+Batch acummalator normalization

{cogforge_engine-1.1.0 → cogforge_engine-2.0.0}/cogforge/app.py

@@ -981,6 +981,47 @@ class MultiHeadAttention:
     def parameters(self):
         return (self.q.parameters() + self.k.parameters()
                 + self.v.parameters() + self.o.parameters())
+        
+
+
+class CrossAttention:
+    def __init__(self, dim_dec, dim_enc,d_out, dec_rope = None, enc_rope = None,d_k=None, h=None, d_model = None):
+        assert (d_k and h) or (d_model and h), "Model Parameters not defined"
+        self.d_model = d_model if d_model is not None else h*d_k
+        self.q = Linear(nin=dim_dec, nout=self.d_model)
+        self.k = Linear(nin=dim_enc, nout=self.d_model)
+        self.v = Linear(nin=dim_enc, nout=self.d_model) 
+        self.d_out = d_out
+        self.o = Linear(nin=self.d_model, nout=d_out)
+        self.dk = d_k if d_k is not None else d_model//h
+        self.h = h  
+        self.dec_rope = dec_rope
+        self.enc_rope = enc_rope
+        self.attention = Attention(self.dk)
+        
+    def split(self,x:Tensor):
+        B,T,_ = x.data.shape
+        return x.view((B,T,self.h,self.dk)).transpose((0,2,1,3)) #(B, h, T, d_k)
+    
+    def merge(self,x:Tensor):
+        B, h, T, d_k = x.data.shape
+        return x.transpose((0, 2, 1, 3)).view((B, T, h * d_k))
+        
+    def __call__(self, x_decod,x_encod, mask=None):
+        Q = self.split(self.q(x_decod))
+        K = self.split(self.k(x_encod))
+        V = self.split(self.v(x_encod))
+        if self.dec_rope is not None:
+            Q = self.dec_rope(Q)
+        if self.enc_rope is not None:
+            K = self.enc_rope(K)
+        
+        out = self.attention(Q=Q,K=K,V=V,mask=mask)
+        return self.o(self.merge(out))
+    
+    def parameters(self):
+        return (self.q.parameters() + self.k.parameters()
+                + self.v.parameters() + self.o.parameters())
     
     
 class LayerNorm:
@@ -1057,6 +1098,44 @@ class Transformer:
         return (self.ln1.parameters() + self.attn.parameters()
                 + self.ln2.parameters() + self.ff.parameters())
         
+ 
+class Decoder:
+    """The decoder block has no bridge. So demb = dmodel"""
+    def __init__(self, d_model, n_heads, d_ff=None, is_training=False,dec_in_rope=None,enc_rope=None,dec_rope=None):
+        self.ln1 = LayerNorm(d_model)
+        self.self_attn = MultiHeadAttention(dinp=d_model,dmodel=d_model,dout=d_model,n=n_heads,rope=dec_in_rope)
+        
+        self.ln2 = LayerNorm(d_model)
+        self.cross_attn = CrossAttention(dim_dec=d_model,dim_enc=d_model,d_out=d_model,dec_rope=dec_rope,enc_rope=enc_rope,h=n_heads,d_k=d_model//n_heads)
+        
+        self.ln3 = LayerNorm(d_model)
+        self.ff = FeedForward(dmodel=d_model, dff=d_ff)
+        
+        self.is_training = is_training
+        
+    def __call__(self, x_dec, x_enc, mask=None,cross_mask=None):
+        a = self.ln1(x_dec)
+        x_dec = x_dec + self.self_attn(a,a,a,mask=mask)
+        
+        c = self.ln2(x_dec)
+        x_dec = x_dec + self.cross_attn(x_decod=c,x_encod=x_enc,mask=cross_mask)
+        
+        f = self.ln3(x_dec)
+        x_dec = x_dec + self.ff(f,self.is_training) 
+        
+        return x_dec
+         
+    def train(self, enabled = True):
+        self.is_training = enabled
+    def infer(self, enabled = True):
+        self.is_training = not enabled
+        
+    def parameters(self):
+        return (self.ln1.parameters() + self.self_attn.parameters() +
+                self.ln2.parameters() + self.cross_attn.parameters() +
+                self.ln3.parameters() + self.ff.parameters())
+            
+ 
         
 class PositionalEncoding:
     def __init__(self,max_len,dmodel):

cogforge_engine-2.0.0/cogforge/models.py

@@ -0,0 +1,254 @@
+from . import backend
+from .app import Embedding, PositionalEncoding,Transformer,LayerNorm,Linear,RotatoryPositionalEncoding,to_cpu,needGradientHence,Decoder,Tensor
+import numpy
+
+
+class GPTV1:
+    def __init__(self, vocab, d_model, n_heads, n_layers, max_len, d_ff=None):
+        self.tok = Embedding(vocab, d_model)
+        self.pos    = PositionalEncoding(max_len, d_model)
+        self.blocks = [Transformer(dmodel=d_model,n=n_heads,dff=d_ff) for _ in range(n_layers)]
+        self.ln_f   = LayerNorm(d_model) 
+        self.head   = Linear(d_model, vocab)
+        self.max_len = max_len
+        
+    def __call__(self, idx):
+        T = idx.shape[1]
+        casual = backend.np.triu((backend.np.ones((T,T))),k=1).astype(bool)
+        x = self.pos(self.tok(idx))
+        for block in self.blocks:
+            x = block(x,mask=casual)
+            
+        x = self.ln_f(x)
+        return self.head(x)
+    
+    def parameters(self):
+        ps = self.tok.parameters() + self.ln_f.parameters() + self.head.parameters()
+        for b in self.blocks: ps += b.parameters()
+        return ps
+    
+    def generate(self, idx, n_new, temperature=1.0, top_k=None):
+        idx = numpy.asarray(idx)
+        original_grad_state = not backend.NO_GRAD
+        needGradientHence(False)
+        try:
+            for _ in range(n_new):
+                cond = idx[:, -self.max_len:]
+                logits = to_cpu(self(cond).data)[:, -1, :] / temperature
+                if top_k is not None:
+                    kth = numpy.sort(logits, axis=-1)[:, -top_k][:, None]
+                    logits = numpy.where(logits < kth, -1e9, logits)
+                z = logits - logits.max(-1, keepdims=True)
+                p = numpy.exp(z).astype(numpy.float64); p /= p.sum(-1, keepdims=True)
+                nxt = numpy.array([[numpy.random.choice(len(pr), p=pr)] for pr in p])
+                idx = numpy.concatenate([idx, nxt], axis=1)
+            return idx 
+        finally:
+            needGradientHence(original_grad_state)
+        
+        
+        
+        
+class GPT2:
+    def __init__(self, vocab, d_model, n_heads, n_layers, max_len, d_ff=None, base=10000.0,training = False):
+        assert d_model % n_heads == 0
+        self.tok    = Embedding(vocab, d_model)
+        self.rope   = RotatoryPositionalEncoding(max_len, d_model // n_heads, base)  # dim = d_k
+        self.blocks = [Transformer(d_model, n_heads, dff=d_ff, rope=self.rope,is_training=training)
+                       for _ in range(n_layers)]
+        self.ln_f   = LayerNorm(d_model)
+        self.head   = Linear(d_model, vocab)
+        self.max_len = max_len
+        self.is_training = training
+
+    def train(self, enabled = True):
+        self.is_training = enabled
+        for block in self.blocks:
+            block.train(enabled=enabled)
+    def infer(self, enabled = True):
+        self.is_training = not enabled
+        for block in self.blocks:
+            block.infer(enabled=enabled)
+    
+
+    def __call__(self, idx):
+        T = idx.shape[1]
+        causal = backend.np.triu(backend.np.ones((T, T)), k=1).astype(bool)
+        x = self.tok(idx)                          # straight from token embedding
+        for block in self.blocks:
+            x = block(x, mask=causal)
+        return self.head(self.ln_f(x))
+
+    def parameters(self):
+        ps = self.tok.parameters() + self.ln_f.parameters() + self.head.parameters()
+        for b in self.blocks: ps += b.parameters()
+        return ps   
+    
+    def generate(self, idx, n_new, temperature=1.0, top_k=None):
+        idx = numpy.asarray(idx)                                 
+        was_training = self.is_training
+        self.infer()
+        original_grad_state = not backend.NO_GRAD
+        needGradientHence(False)
+        
+        try:
+            for _ in range(n_new):
+                cond = idx[:, -self.max_len:]
+                logits = to_cpu(self(cond).data)[:, -1, :] / temperature   
+                if top_k is not None:
+                    kth = numpy.sort(logits, axis=-1)[:, -top_k][:, None]
+                    logits = numpy.where(logits < kth, -1e9, logits)
+                z = logits - logits.max(-1, keepdims=True)
+                p = numpy.exp(z).astype(numpy.float64); p /= p.sum(-1, keepdims=True)
+                nxt = numpy.array([[numpy.random.choice(len(pr), p=pr)] for pr in p])
+                idx = numpy.concatenate([idx, nxt], axis=1)
+                # backend._cp.get_default_memory_pool().free_all_blocks() if backend.USE_GPU else None
+            return idx 
+        finally:
+            if was_training:
+                self.train()
+            needGradientHence(original_grad_state)   
+    
+    
+
+
+class Seq2Seq:
+    def __init__(self, enc_vocab, dec_vocab, d_model, n_heads, num_enc_layers, num_dec_layers, max_len, d_ff=None, training=False,shared_tok = False,pad_id=0,encoder_rope = None,dec_in_rope = None, dec_rope = None, enc_rope=None):
+        if enc_vocab==dec_vocab and shared_tok: 
+            self.shared_tok = Embedding(enc_vocab, d_model)
+            self.shared_tok.weights.data *= (1.0 / backend.np.sqrt(d_model))
+            self.emb_scale = backend.np.sqrt(d_model)
+            
+            self.enc_tok = self.dec_tok = self.shared_tok
+            self.tie_weights = True
+            self.head_b = Tensor(backend.np.zeros((dec_vocab,)))
+        else:
+            self.enc_tok = Embedding(enc_vocab, d_model)
+            self.dec_tok = Embedding(dec_vocab, d_model)
+            
+            self.head = Linear(d_model, dec_vocab)
+            self.tie_weights = False
+            self.emb_scale = 1.0        
+        
+        self.pad_id = pad_id
+        
+        self.rope = RotatoryPositionalEncoding(max_len, d_model // n_heads) if (dec_in_rope or dec_rope or enc_rope or encoder_rope) is not None else None
+        self.dec_in_rope = self.rope if dec_in_rope is not None else None
+        self.dec_rope = self.rope if dec_rope is not None else None
+        self.enc_rope = self.rope if enc_rope is not None else None
+        self.encoder_rope = self.rope if encoder_rope is not None else None
+        
+        self.encoders = [Transformer(dmodel=d_model, n=n_heads, dff=d_ff, is_training=training, rope=self.encoder_rope) for _ in range(num_enc_layers)]
+        self.decoders = [
+            Decoder(d_model=d_model, n_heads=n_heads, d_ff=d_ff, is_training=training, dec_in_rope=self.dec_in_rope,dec_rope=self.dec_rope,enc_rope=self.enc_rope) 
+            for _ in range(num_dec_layers)
+        ]
+        
+        self.ln_f = LayerNorm(d_model)
+        self.enc_ln = LayerNorm(d_model)
+        
+        self.max_len = max_len
+        self.is_training = training
+        
+    def train(self, enabled=True):
+        self.is_training = enabled
+        for block in self.encoders: block.train(enabled)
+        for block in self.decoders: block.train(enabled)
+
+    def infer(self, enabled=True):
+        self.is_training = not enabled
+        for block in self.encoders: block.infer(enabled)
+        for block in self.decoders: block.infer(enabled)
+    
+    @staticmethod   
+    def make_pad_mask(idx, pad_id):
+        # idx: (B, T) int  ->  (B, 1, 1, T) bool, True where padding
+        m = (backend.np.asarray(idx) == pad_id)
+        return m[:, None, None, :]
+        
+    def __call__(self, enc_idx, dec_idx):
+        """
+        enc_idx: Scrambled words, shape (B, T_enc)
+        dec_idx: Shifted target sentence starting with <SOS>, shape (B, T_dec)
+        """
+        enc_pad = self.make_pad_mask(enc_idx, self.pad_id) 
+        x_enc = self.enc_tok(enc_idx) * self.emb_scale
+        for block in self.encoders:
+            x_enc = block(x_enc, mask=enc_pad) 
+            
+        x_enc = self.enc_ln(x_enc)  
+            
+        T_dec = dec_idx.shape[1]
+        causal_mask = backend.np.triu(backend.np.ones((T_dec, T_dec)), k=1).astype(bool)
+        
+        x_dec = self.dec_tok(dec_idx) * self.emb_scale
+        for block in self.decoders:
+            x_dec = block(x_dec=x_dec, x_enc=x_enc, mask=causal_mask,cross_mask=enc_pad)
+            
+        x = self.ln_f(x_dec)
+        if self.tie_weights:
+            logits = x @ self.shared_tok.weights.transpose((1, 0)) + self.head_b
+        else:
+            logits = self.head(x)
+        return logits
+
+    def parameters(self):
+        seen, ps = set(), []
+        def add(lst):
+            for p in lst:
+                if id(p) not in seen:
+                    seen.add(id(p)); ps.append(p)
+        add(self.enc_tok.parameters()); add(self.dec_tok.parameters())
+        add(self.ln_f.parameters()); add(self.enc_ln.parameters())
+        add([self.head_b] if self.tie_weights else self.head.parameters())
+        for b in self.encoders: add(b.parameters())
+        for b in self.decoders: add(b.parameters())
+        return ps
+    
+    
+    def encode(self, enc_idx):
+        enc_pad = self.make_pad_mask(enc_idx, self.pad_id)
+        x_enc = self.enc_tok(enc_idx) * self.emb_scale
+        for block in self.encoders:
+            x_enc = block(x_enc, mask=enc_pad)
+        x_enc = self.enc_ln(x_enc)
+        return x_enc, enc_pad
+
+    def decode_step(self, dec_idx, x_enc, enc_pad):
+        T_dec = dec_idx.shape[1]
+        causal = backend.np.triu(backend.np.ones((T_dec, T_dec)), k=1).astype(bool)
+        x_dec = self.dec_tok(dec_idx) * self.emb_scale
+        for block in self.decoders:
+            x_dec = block(x_dec=x_dec, x_enc=x_enc, mask=causal, cross_mask=enc_pad)
+        x = self.ln_f(x_dec)
+        if self.tie_weights:
+            return x @ self.shared_tok.weights.transpose((1, 0)) + self.head_b
+        return self.head(x)
+
+    def generate(self, enc_idx, sos_id, eos_id=None, max_new=50, temperature=1.0, top_k=None):
+        enc_idx = numpy.asarray(enc_idx)
+        B = enc_idx.shape[0]
+        was_training = self.is_training
+        self.infer()
+        original_grad_state = not backend.NO_GRAD
+        needGradientHence(False)
+        try:
+            x_enc, enc_pad = self.encode(enc_idx)          # encode once
+            dec_idx = numpy.full((B, 1), sos_id, dtype=numpy.int64)
+            for _ in range(max_new):
+                logits = to_cpu(self.decode_step(dec_idx, x_enc, enc_pad).data)[:, -1, :] / temperature
+                if top_k is not None:
+                    kth = numpy.sort(logits, axis=-1)[:, -top_k][:, None]
+                    logits = numpy.where(logits < kth, -1e9, logits)
+                z = logits - logits.max(-1, keepdims=True)
+                p = numpy.exp(z).astype(numpy.float64); p /= p.sum(-1, keepdims=True)
+                nxt = numpy.array([[numpy.random.choice(len(pr), p=pr)] for pr in p])
+                dec_idx = numpy.concatenate([dec_idx, nxt], axis=1)
+                if eos_id is not None and bool((nxt[:, 0] == eos_id).all()):
+                    break
+            return dec_idx
+        finally:
+            if was_training:
+                self.train()
+            needGradientHence(original_grad_state)
+            

{cogforge_engine-1.1.0 → cogforge_engine-2.0.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "cogforge-engine" 
-version = "1.1.0"
+version = "2.0.0"
 authors = [
   { name="Avik Majumder", email="avikmjd2@gmail.com" },
 ]

cogforge_engine-1.1.0/cogforge/models.py

@@ -1,114 +0,0 @@
-from . import backend
-from .app import Embedding, PositionalEncoding,Transformer,LayerNorm,Linear,RotatoryPositionalEncoding,to_cpu,needGradientHence
-import numpy
-
-
-class GPTV1:
-    def __init__(self, vocab, d_model, n_heads, n_layers, max_len, d_ff=None):
-        self.tok = Embedding(vocab, d_model)
-        self.pos    = PositionalEncoding(max_len, d_model)
-        self.blocks = [Transformer(dmodel=d_model,n=n_heads,dff=d_ff) for _ in range(n_layers)]
-        self.ln_f   = LayerNorm(d_model) 
-        self.head   = Linear(d_model, vocab)
-        self.max_len = max_len
-        
-    def __call__(self, idx):
-        T = idx.shape[1]
-        casual = backend.np.triu((backend.np.ones((T,T))),k=1).astype(bool)
-        x = self.pos(self.tok(idx))
-        for block in self.blocks:
-            x = block(x,mask=casual)
-            
-        x = self.ln_f(x)
-        return self.head(x)
-    
-    def parameters(self):
-        ps = self.tok.parameters() + self.ln_f.parameters() + self.head.parameters()
-        for b in self.blocks: ps += b.parameters()
-        return ps
-    
-    def generate(self, idx, n_new, temperature=1.0, top_k=None):
-        idx = numpy.asarray(idx)
-        original_grad_state = not backend.NO_GRAD
-        needGradientHence(False)
-        try:
-            for _ in range(n_new):
-                cond = idx[:, -self.max_len:]
-                logits = to_cpu(self(cond).data)[:, -1, :] / temperature
-                if top_k is not None:
-                    kth = numpy.sort(logits, axis=-1)[:, -top_k][:, None]
-                    logits = numpy.where(logits < kth, -1e9, logits)
-                z = logits - logits.max(-1, keepdims=True)
-                p = numpy.exp(z).astype(numpy.float64); p /= p.sum(-1, keepdims=True)
-                nxt = numpy.array([[numpy.random.choice(len(pr), p=pr)] for pr in p])
-                idx = numpy.concatenate([idx, nxt], axis=1)
-            return idx 
-        finally:
-            needGradientHence(original_grad_state)
-        
-        
-        
-        
-class GPT2:
-    def __init__(self, vocab, d_model, n_heads, n_layers, max_len, d_ff=None, base=10000.0,training = False):
-        assert d_model % n_heads == 0
-        self.tok    = Embedding(vocab, d_model)
-        self.rope   = RotatoryPositionalEncoding(max_len, d_model // n_heads, base)  # dim = d_k
-        self.blocks = [Transformer(d_model, n_heads, dff=d_ff, rope=self.rope,is_training=training)
-                       for _ in range(n_layers)]
-        self.ln_f   = LayerNorm(d_model)
-        self.head   = Linear(d_model, vocab)
-        self.max_len = max_len
-        self.is_training = training
-
-    def train(self, enabled = True):
-        self.is_training = enabled
-        for block in self.blocks:
-            block.train(enabled=enabled)
-    def infer(self, enabled = True):
-        self.is_training = not enabled
-        for block in self.blocks:
-            block.infer(enabled=enabled)
-    
-
-    def __call__(self, idx):
-        T = idx.shape[1]
-        causal = backend.np.triu(backend.np.ones((T, T)), k=1).astype(bool)
-        x = self.tok(idx)                          # straight from token embedding
-        for block in self.blocks:
-            x = block(x, mask=causal)
-        return self.head(self.ln_f(x))
-
-    def parameters(self):
-        ps = self.tok.parameters() + self.ln_f.parameters() + self.head.parameters()
-        for b in self.blocks: ps += b.parameters()
-        return ps   
-    
-    def generate(self, idx, n_new, temperature=1.0, top_k=None):
-        idx = numpy.asarray(idx)                                 
-        was_training = self.is_training
-        self.infer()
-        original_grad_state = not backend.NO_GRAD
-        needGradientHence(False)
-        
-        try:
-            for _ in range(n_new):
-                cond = idx[:, -self.max_len:]
-                logits = to_cpu(self(cond).data)[:, -1, :] / temperature   
-                if top_k is not None:
-                    kth = numpy.sort(logits, axis=-1)[:, -top_k][:, None]
-                    logits = numpy.where(logits < kth, -1e9, logits)
-                z = logits - logits.max(-1, keepdims=True)
-                p = numpy.exp(z).astype(numpy.float64); p /= p.sum(-1, keepdims=True)
-                nxt = numpy.array([[numpy.random.choice(len(pr), p=pr)] for pr in p])
-                idx = numpy.concatenate([idx, nxt], axis=1)
-                # backend._cp.get_default_memory_pool().free_all_blocks() if backend.USE_GPU else None
-            return idx 
-        finally:
-            if was_training:
-                self.train()
-            needGradientHence(original_grad_state)   
-    
-    
-    
-