Implimented Transformer(decode only)

model.py

@@ -57,9 +57,18 @@ class Block:
         return self
 
 class Transformer():
-    def __init__(self):
+    def __init__(self,vocab_size,embed_size,n_heads,n_blocks):
-        pass #TODO
+        self.tok_embed = nn.Embedding(vocab_size,embed_size)
-    def __call__(self):
+        self.blocks = [Block(embed_size,n_heads) for _ in range(n_blocks)]
-        pass #TODO
+        self.norm = nn.RMSNorm(embed_size)
-    def cast(self):
+        self.output = nn.Linear(embed_size,vocab_size,bias=False)
-        pass #TODO
+    def __call__(self,x):
+        x = self.tok_embed(x)
+        x = x.sequential(self.blocks)
+        x = self.norm(x)
+        return self.output(x)
+    def cast(self,dtype):
+        self.tok_embed.weight = self.tok_embed.weight.cast(dtype)
+        self.blocks = [b.cast(dtype) for b in self.blocks]
+        self.output.weight = self.output.weight.cast(dtype)
+        return self