k/llm
1
0
Fork 0
Code Issues 6 Pull requests Projects Releases Packages Wiki Activity Actions

Implimented TransformerBlock

Browse source
This commit is contained in:
k 2026-01-06 19:53:37 -05:00
parent 77aa0de0eb
commit 23f62c7e64
1 changed files with 13 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

19
model.py
View file

@ -42,12 +42,19 @@ class FeedForwardNetwork:
return self return self
class Block: class Block:
def __init__(self): def __init__(self,embed_size,n_heads):
pass #TODO self.mha = MultiHeadAttention(embed_size,n_heads)
def __call__(self): self.ffn = FeedForwardNetwork(embed_size)
pass #TODO self.mhaNorm = nn.RMSNorm(embed_size)
def cast(self): self.ffnNorm = nn.RMSNorm(embed_size)
pass #TODO def __call__(self,x):
x = x + self.mha(self.mhaNorm(x))
x = x + self.ffn(self.ffnNorm(x))
return x
def cast(self,dtype):
self.mha = self.mha.cast(dtype)
self.ffn = self.ffn.cast(dtype)
return self
class Transformer(): class Transformer():
def __init__(self): def __init__(self):