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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cogforge-engine
-Version: 1.0.7
+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.0.7 → cogforge_engine-2.0.0}/cogforge/app.py

@@ -51,7 +51,7 @@ class Tensor:
         # if self.requires_grad:
         self.grad = None if backend.NO_GRAD else backend.np.zeros(self.shape,dtype=backend.np.float32 if typed == "compressed" else backend.np.float64)
         self._backwards = lambda:None
-        self._children = set(children)
+        self._children = set() if backend.NO_GRAD else set(children)
         self.typed = typed
         
     def __getitem__(self, key):
@@ -59,6 +59,7 @@ class Tensor:
         
         out = Tensor(out_data,children=(self,))
         
+        if backend.NO_GRAD: return out_data
         def _backward():
             self.grad[key] += out.grad
             # backend.np.add.at(self.grad, key, out.grad)
@@ -73,7 +74,11 @@ class Tensor:
         Concatenates a sequence of Tensors along a specified axis.
         """
         arrays = tuple(t.data for t in tensors)
+        if backend.NO_GRAD: return Tensor(backend.np.concatenate(arrays, axis=axis))
+        
         out = Tensor(backend.np.concatenate(arrays,axis=axis),children=tensors)
+        
+        
         sizes = [a.shape[axis] for a in arrays[:-1]]
         split_indices = numpy.cumsum(sizes).tolist()
         
@@ -95,6 +100,8 @@ class Tensor:
         other = other if isinstance(other, Tensor) else Tensor(other)
         out_data = self.data + other.data
         out = Tensor(out_data, (self, other))
+        
+        if backend.NO_GRAD: return Tensor(out_data)
 
         def _backward():
             if self.shape != out_data.shape:
@@ -129,6 +136,8 @@ class Tensor:
         other = other if isinstance(other, Tensor) else Tensor(other)
         out_data = self.data * other.data
         
+        if backend.NO_GRAD: return Tensor(out_data)
+        
         out = Tensor(out_data, children=(self, other))
         
         def _backward():
@@ -177,6 +186,9 @@ class Tensor:
         return grad
 
     def __matmul__(self, other):    
+        
+        if backend.NO_GRAD: return Tensor(self.data@other.data)
+        
         out = Tensor(self.data @ other.data, children=(self, other))
         def _backward():
             gs = out.grad @ backend.np.swapaxes(other.data, -1, -2)
@@ -189,6 +201,8 @@ class Tensor:
     def relu(self):
         out_data = backend.np.maximum(0,self.data)
         
+        if backend.NO_GRAD: return Tensor(out_data)
+        
         out = Tensor(out_data, children=(self,))
         
         def _backward():
@@ -208,6 +222,10 @@ class Tensor:
             arr = self.data
         
         out_data = arr.reshape(shape)
+        
+        if backend.NO_GRAD: return Tensor(out_data)
+        
+        
         out = Tensor(out_data, children=(self,))
         
         def _backward():
@@ -222,6 +240,9 @@ class Tensor:
         
     def sigmoid(self):
         out_data = 1 / (1 + backend.np.exp(-self.data))
+        
+        if backend.NO_GRAD: return Tensor(out_data)
+        
         out = Tensor(out_data, children=(self,))
         def _backward():
             local_derivative = out_data * (1 - out_data)
@@ -233,6 +254,7 @@ class Tensor:
     
     def tanh(self):
         out_data = backend.np.tanh(self.data)
+        if backend.NO_GRAD: return Tensor(out_data)
         out = Tensor(out_data,children=(self,))
         def _backward():
             dervitive = 1-(out_data*out_data)
@@ -246,6 +268,7 @@ class Tensor:
         expVal = backend.np.exp(self.data-maxVal)
         
         out_data = expVal/backend.np.sum(expVal,axis=axis,keepdims=True)
+        if backend.NO_GRAD: return Tensor(out_data)
         out=Tensor(out_data,(self,))
         
         def _backward():
@@ -263,6 +286,7 @@ class Tensor:
         batch_size = predictions.shape[0]
         
         loss_data = - backend.np.sum(targets * backend.np.log(probabilities)) / batch_size
+        if backend.NO_GRAD: return cls(loss_data)
         loss = Tensor(loss_data, children=(predictions,))
         
         def _backward():
@@ -325,6 +349,8 @@ class Tensor:
         batch_size = self.shape[0]
         out_data = self.data.reshape(batch_size, -1)
         
+        if backend.NO_GRAD: return Tensor(out_data)
+        
         out = Tensor(out_data, children=(self,))
         def _backward():
             self.grad += out.grad.reshape(self.shape)
@@ -346,6 +372,8 @@ class Tensor:
         
         out_data = self.data.reshape(B,T//num,C*num)
         
+        if backend.NO_GRAD: return Tensor(out_data)
+        
         out = Tensor(out_data, children=(self,))
         
         def _backward():
@@ -369,6 +397,8 @@ class Tensor:
 
         # only real rows contribute to the loss value
         loss_data = -backend.np.sum(mask * targets * backend.np.log(probs)) / n_real
+        
+        if backend.NO_GRAD: return cls(loss_data)
         loss = cls(loss_data, children=(predictions,))
 
         def _backward():
@@ -388,7 +418,14 @@ class Tensor:
         p = e / e.sum(axis=-1, keepdims=True)
         N = int(numpy.prod(scores.shape[:-1]))
         flat, idx, rows = p.reshape(N, -1), target_ids.reshape(N), backend.np.arange(N)
-        loss = cls(-backend.np.log(backend.np.clip(flat[rows, idx], 1e-15, 1.0)).sum() / N, children=(scores,))
+        # loss = cls(-backend.np.log(backend.np.clip(flat[rows, idx], 1e-15, 1.0)).sum() / N, children=(scores,))
+        
+        loss_data = -backend.np.log(backend.np.clip(flat[rows, idx], 1e-15, 1.0)).sum() / N
+        
+        if backend.NO_GRAD: return cls(loss_data)
+        
+        loss = cls(loss_data, children=(scores,))
+        
         def _backward():
             # g = p.reshape(N, -1).copy()
             g = p.reshape(N, -1)
@@ -402,7 +439,14 @@ class Tensor:
         p = _stable_softmax(scores.data)
         N = int(numpy.prod(scores.shape[:-1]))
         flat, idx, rows = p.reshape(N, -1), target_ids.reshape(N), backend.np.arange(N)
-        loss = cls(-backend.np.log(backend.np.clip(flat[rows, idx], 1e-15, 1.0)).sum() / N, children=(scores,))
+        # loss = cls(-backend.np.log(backend.np.clip(flat[rows, idx], 1e-15, 1.0)).sum() / N, children=(scores,))
+        
+        loss_data = -backend.np.log(backend.np.clip(flat[rows, idx], 1e-15, 1.0)).sum() / N
+        
+        if backend.NO_GRAD: return cls(loss_data)
+        
+        loss = cls(loss_data, children=(scores,))
+        
         def _backward():
             g = p.reshape(N, -1)
             g[rows, idx] -= 1.0
@@ -420,8 +464,14 @@ class Tensor:
         p = e / e.sum(axis=1, keepdims=True)          # softmax, computed privately
         B = scores.shape[0]
 
-        loss = cls(-backend.np.sum(targets * backend.np.log(backend.np.clip(p, 1e-15, 1.0))) / B, children=(scores,))
+        # loss = cls(-backend.np.sum(targets * backend.np.log(backend.np.clip(p, 1e-15, 1.0))) / B, children=(scores,))
 
+        loss_data = -backend.np.sum(targets * backend.np.log(backend.np.clip(p, 1e-15, 1.0))) / B
+        
+        if backend.NO_GRAD: return cls(loss_data)
+        
+        loss = cls(loss_data, children=(scores,))
+        
         def _backward():
             scores.grad += (p - targets) / B * loss.grad
         loss._backwards = _backward
@@ -434,7 +484,10 @@ class Tensor:
         e = backend.np.exp(z)
         p = e / e.sum(axis=-1, keepdims=True)
         N = numpy.prod(scores.shape[:-1])                  # B for 2D, B*T for 3D
-        loss = cls(-backend.np.sum(targets * backend.np.log(backend.np.clip(p, 1e-15, 1.0))) / N, children=(scores,))
+        # loss = cls(-backend.np.sum(targets * backend.np.log(backend.np.clip(p, 1e-15, 1.0))) / N, children=(scores,))
+        loss_data = -backend.np.sum(targets * backend.np.log(backend.np.clip(p, 1e-15, 1.0))) / N        
+        if backend.NO_GRAD: return cls(loss_data)        
+        loss = cls(loss_data, children=(scores,))
         def _backward():
             scores.grad += (p - targets) / N * loss.grad
         loss._backwards = _backward
@@ -452,8 +505,13 @@ class Tensor:
         n_real = mask.sum()
         n_real = n_real if n_real > 0 else 1.0
 
-        loss = cls(-backend.np.sum(mask * targets * backend.np.log(backend.np.clip(p, 1e-15, 1.0))) / n_real, children=(scores,))
-
+        # loss = cls(-backend.np.sum(mask * targets * backend.np.log(backend.np.clip(p, 1e-15, 1.0))) / n_real, children=(scores,))
+        loss_data = -backend.np.sum(mask * targets * backend.np.log(backend.np.clip(p, 1e-15, 1.0))) / n_real
+        
+        if backend.NO_GRAD: return cls(loss_data)
+        
+        loss = cls(loss_data, children=(scores,))
+        
         def _backward():
             grad = (p - targets) / n_real
             grad = grad * mask                         
@@ -462,6 +520,7 @@ class Tensor:
         return loss
     
     def transpose(self,axes):
+        if backend.NO_GRAD: return Tensor(backend.np.transpose(self.data,axes=axes))
         out = Tensor(backend.np.transpose(self.data,axes=axes),children=(self,))
         # inv = backend.np.argsort(axes)
         inv = tuple(int(i) for i in numpy.argsort(axes)) 
@@ -473,12 +532,49 @@ class Tensor:
     
     def masked_fill(self, mask, value):
         out_data = backend.np.where(mask, value, self.data)
+        if backend.NO_GRAD: return Tensor(out_data)
         out = Tensor(out_data, children=(self,))
         def _backward():
             self.grad += backend.np.where(mask, 0.0, 1.0) * out.grad
         out._backwards = _backward
         return out
-                            
+    
+    def dropout(self,p=0.1,training=True):
+        if not training or p == 0.0:
+            return self
+        
+        mask = (backend.np.random.rand(*self.shape)>p).astype(self.data.dtype)
+        scale = 1.0 / (1.0 - p)
+        out_data = self.data * mask * scale
+        if backend.NO_GRAD: return Tensor(out_data)
+        
+        out = Tensor(out_data, children=(self,))
+        
+        def _backward():
+            self.grad += (mask * scale) * out.grad
+            
+        out._backwards = _backward
+        return out
+    
+    def dropTheWholeNeuron(self,p=0.1,training=True,axis=-1,batch_ind = 0):
+        if not training or p == 0.0:
+            return self
+        
+        N = self.shape[axis]
+        B = self.shape[batch_ind]
+        R = (1 for _ in range(len(self.shape)-2))
+
+        mask = (backend.np.random.rand(B,*R,N)>p).astype(self.self.data.dtype)
+        scale = 1.0 / (1.0 - p)
+        out_data = self.data * mask *scale
+        if backend.NO_GRAD: return Tensor(out_data)
+        out = Tensor(out_data, children=(self,))
+
+        def _backward():
+            self.grad += (mask * scale) * out.grad
+            
+        out._backwards = _backward
+        return out
         
             
             
@@ -539,7 +635,7 @@ class Embedding:
     def __call__(self, input_indices):
         idx = backend.np.asarray(input_indices)   
         out_data = self.weights.data[idx]
-        
+        if backend.NO_GRAD: return Tensor(out_data)
         out = Tensor(out_data, children=(self.weights,))
         
         def _backward():
@@ -722,7 +818,7 @@ class BatchNorm1D:
             self.x_hat = (x.data-self.running_mean)/backend.np.sqrt(self.running_var + self.eps)
         
         out_data = self.gamma.data * self.x_hat + self.beta.data
-        
+        if backend.NO_GRAD: return Tensor(out_data)
         out = Tensor(out_data, children=(x, self.gamma, self.beta))
         
         def _backward():
@@ -885,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:
@@ -899,6 +1036,7 @@ class LayerNorm:
         std_inv = 1.0 / backend.np.sqrt(var + self.eps)
         x_hat = (x.data - mu) * std_inv
         out_data = self.gamma.data * x_hat + self.beta.data
+        if backend.NO_GRAD: return Tensor(out_data)
         out = Tensor(out_data, children=(x, self.gamma, self.beta))
 
         D = x.data.shape[-1]
@@ -926,25 +1064,33 @@ class FeedForward:
         self.fc1 = Linear(dmodel,dff)
         self.fc2 = Linear(dff,dmodel)
         
-    def __call__(self, x):
-        return self.fc2(self.fc1(x).relu())
+    def __call__(self, x,is_training = False):
+        hidden:Tensor = self.fc1(x).relu()
+        hidden = hidden.dropout(p=0.15,training=is_training)
+        return self.fc2(hidden)
     def parameters(self):
         return self.fc1.parameters() + self.fc2.parameters()
 
 class Transformer:
     """The transformer block has no bridge. So demb = dmodel"""
     
-    def __init__(self, dmodel, n, dff = None,rope=None):
+    def __init__(self, dmodel, n, dff = None,rope=None, is_training = False):
         self.ln1 = LayerNorm(dmodel)
         self.attn = MultiHeadAttention(dinp=dmodel,dmodel=dmodel,dout=dmodel,n=n,rope=rope)
         self.ln2 = LayerNorm(dmodel)
         self.ff  = FeedForward(dmodel=dmodel,dff=dff)
+        self.is_training = is_training
+        
+    def train(self, enabled = True):
+        self.is_training = enabled
+    def infer(self, enabled = True):
+        self.is_training = not enabled
         
     def __call__(self,x,mask=None):
         a = self.ln1(x)
         x = x + self.attn(a,a,a,mask=mask)
         f = self.ln2(x)
-        x = x+ self.ff(f)
+        x = x+ self.ff(f,is_training=self.is_training)
         
         return x
     
@@ -952,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.0.7 → cogforge_engine-2.0.0}/pyproject.toml

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

cogforge_engine-1.0.7/cogforge/models.py

@@ -1,85 +0,0 @@
-from . import backend
-from .app import Embedding, PositionalEncoding,Transformer,LayerNorm,Linear,RotatoryPositionalEncoding,to_cpu
-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)
-        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                                          # <-- was a bare `return` 
-        
-        
-        
-        
-class GPT2:
-    def __init__(self, vocab, d_model, n_heads, n_layers, max_len, d_ff=None, base=10000.0):
-        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)
-                       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]
-        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)                                 # token ids live on host
-        for _ in range(n_new):
-            cond = idx[:, -self.max_len:]
-            logits = to_cpu(self(cond).data)[:, -1, :] / temperature   # forward on GPU, logits to host
-            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