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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cogforge-engine
-Version: 1.0.7
+Version: 1.1.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-1.0.7 → cogforge_engine-1.1.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():
@@ -899,6 +995,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 +1023,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
     

cogforge_engine-1.1.0/cogforge/models.py

@@ -0,0 +1,114 @@
+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)   
+    
+    
+    
+                              

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

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "cogforge-engine" 
-version = "1.0.7"
+version = "1.1.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