65 lines
2.9 KiB
Python
65 lines
2.9 KiB
Python
import numpy as np
|
|
from tinygrad.tensor import Tensor
|
|
|
|
class TransformerBlock:
|
|
def __init__(self, embed_dim, num_heads, ff_dim, prenorm=False, act=lambda x: x.relu(), dropout=0.1):
|
|
assert embed_dim % num_heads == 0, "embed_dim must be divisible by num_heads"
|
|
|
|
self.num_heads = num_heads
|
|
self.head_size = embed_dim // num_heads
|
|
self.prenorm, self.act = prenorm, act
|
|
self.dropout = dropout
|
|
|
|
self.query = (Tensor.scaled_uniform(embed_dim, embed_dim), Tensor.zeros(embed_dim))
|
|
self.key = (Tensor.scaled_uniform(embed_dim, embed_dim), Tensor.zeros(embed_dim))
|
|
self.value = (Tensor.scaled_uniform(embed_dim, embed_dim), Tensor.zeros(embed_dim))
|
|
|
|
self.out = (Tensor.scaled_uniform(embed_dim, embed_dim), Tensor.zeros(embed_dim))
|
|
|
|
self.ff1 = (Tensor.scaled_uniform(embed_dim, ff_dim), Tensor.zeros(ff_dim))
|
|
self.ff2 = (Tensor.scaled_uniform(ff_dim, embed_dim), Tensor.zeros(embed_dim))
|
|
|
|
self.ln1 = (Tensor.ones(embed_dim), Tensor.zeros(embed_dim))
|
|
self.ln2 = (Tensor.ones(embed_dim), Tensor.zeros(embed_dim))
|
|
|
|
def attn(self, x):
|
|
# x: (bs, time, embed_dim) -> (bs, time, embed_dim)
|
|
query, key, value = [x.linear(*y).reshape(shape=(x.shape[0], -1, self.num_heads, self.head_size)).transpose(1,2) for y in [self.query, self.key, self.value]]
|
|
attention = Tensor.scaled_dot_product_attention(query, key, value).transpose(1,2)
|
|
return attention.reshape(shape=(x.shape[0], -1, self.num_heads * self.head_size)).linear(*self.out)
|
|
|
|
def __call__(self, x):
|
|
if self.prenorm:
|
|
x = x + self.attn(x.layernorm().linear(*self.ln1)).dropout(self.dropout)
|
|
x = x + self.act(x.layernorm().linear(*self.ln2).linear(*self.ff1)).linear(*self.ff2).dropout(self.dropout)
|
|
else:
|
|
x = x + self.attn(x).dropout(self.dropout)
|
|
x = x.layernorm().linear(*self.ln1)
|
|
x = x + self.act(x.linear(*self.ff1)).linear(*self.ff2).dropout(self.dropout)
|
|
x = x.layernorm().linear(*self.ln2)
|
|
return x
|
|
|
|
class Transformer:
|
|
def __init__(self, syms, maxlen, layers, embed_dim, num_heads, ff_dim):
|
|
self.maxlen, self.syms = maxlen, syms
|
|
self.embed = Tensor.scaled_uniform(maxlen+syms, embed_dim, requires_grad=False)
|
|
self.tbs = []
|
|
for i in range(layers):
|
|
self.tbs.append(TransformerBlock(embed_dim, num_heads, ff_dim))
|
|
self.final = Tensor.scaled_uniform(embed_dim, syms)
|
|
|
|
def forward(self, x):
|
|
bs = x.shape[0]
|
|
xnp = x.numpy().astype(np.int32)
|
|
onehot = np.zeros((bs, x.shape[1], self.maxlen+self.syms), dtype=np.float32)
|
|
for i in range(x.shape[1]):
|
|
onehot[range(bs), i, i] = 1
|
|
onehot[range(bs), i, self.maxlen + xnp[:, i]] = 1
|
|
onehot = onehot.reshape(bs*x.shape[1], self.maxlen+self.syms)
|
|
|
|
x = Tensor(onehot, device=x.device).dot(self.embed).reshape(shape=(bs, x.shape[1], -1))
|
|
x = x.sequential(self.tbs)
|
|
x = x.reshape(shape=(-1, x.shape[-1])).dot(self.final).log_softmax()
|
|
return x.reshape(shape=(bs, -1, x.shape[-1]))
|
|
|