File size: 1,379 Bytes
7a0c684 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 |
import numpy as np
from .embedding import embedding_lookup, add_positional_encoding
from .positional_encoding import sinusoidal_positional_encoding
from .decoder import transformer_decoder_block
class TransformerDecoder:
def __init__(self, vocab_size, hidden_dim, num_layers, num_heads, max_seq_len, embedding_weights, block_weights_list, driver=None, scheduler=None):
self.vocab_size = vocab_size
self.hidden_dim = hidden_dim
self.num_layers = num_layers
self.num_heads = num_heads
self.max_seq_len = max_seq_len
self.embedding_weights = embedding_weights # (vocab_size, hidden_dim)
self.block_weights_list = block_weights_list # list of dicts, one per block
self.pos_encoding = sinusoidal_positional_encoding(max_seq_len, hidden_dim)
self.driver = driver
self.scheduler = scheduler
def forward(self, input_ids, enc_out, self_mask=None, enc_dec_mask=None):
# input_ids: (batch, tgt_seq_len)
x = embedding_lookup(input_ids, self.embedding_weights)
x = add_positional_encoding(x, self.pos_encoding[:x.shape[1]])
for block_weights in self.block_weights_list:
x = transformer_decoder_block(x, enc_out, block_weights, self.num_heads, self_mask, enc_dec_mask, driver=self.driver, scheduler=self.scheduler)
return x
|