52 lines
1.5 KiB
Python
52 lines
1.5 KiB
Python
import numpy as np
|
|
from transformer import SimpleTransformer, create_padding_mask, create_look_ahead_mask
|
|
|
|
|
|
def main():
|
|
vocab_size = 1000
|
|
d_model = 512
|
|
num_heads = 8
|
|
num_layers = 6
|
|
d_ff = 2048
|
|
max_seq_len = 100
|
|
|
|
model = SimpleTransformer(vocab_size, d_model, num_heads, num_layers, d_ff, max_seq_len)
|
|
|
|
batch_size = 2
|
|
seq_len = 10
|
|
x = np.random.randint(0, vocab_size, (batch_size, seq_len))
|
|
|
|
print("=== Simple Transformer Example ===")
|
|
print(f"Vocabulary size: {vocab_size}")
|
|
print(f"Model dimension: {d_model}")
|
|
print(f"Number of heads: {num_heads}")
|
|
print(f"Number of layers: {num_layers}")
|
|
print(f"Feed-forward dimension: {d_ff}")
|
|
print(f"Max sequence length: {max_seq_len}")
|
|
print()
|
|
|
|
print(f"Input shape: {x.shape}")
|
|
print(f"Input sample: {x[0]}")
|
|
print()
|
|
|
|
output = model.forward(x)
|
|
|
|
print(f"Output shape: {output.shape}")
|
|
print(f"Output sample (first 5 values): {output[0, 0, :5]}")
|
|
print()
|
|
|
|
total_params = model.count_parameters()
|
|
print(f"Total parameters: {total_params:,}")
|
|
print()
|
|
|
|
print("=== Attention Mask Examples ===")
|
|
padding_mask = create_padding_mask(x)
|
|
print(f"Padding mask shape: {padding_mask.shape}")
|
|
|
|
look_ahead_mask = create_look_ahead_mask(seq_len)
|
|
print(f"Look-ahead mask shape: {look_ahead_mask.shape}")
|
|
print(f"Look-ahead mask sample:\n{look_ahead_mask[:5, :5]}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |