A-Z-of-Tranformer-Architecture/h_self_attention.py at main · meetrais/A-Z-of-Tranformer-Architecture · GitHub

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
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
import torch

inputs = torch.tensor(
    [[0.43, 0.15, 0.89], # Your
     [0.55, 0.87, 0.66], # journey
     [0.57, 0.85, 0.64], # starts
     [0.22, 0.58, 0.33], # with
     [0.77, 0.25, 0.10], # one
     [0.05, 0.80, 0.55]] # step
)

query = inputs[1]
print("query: ", query)

attn_scores_2 = torch.empty(inputs.shape[0])
for i, x_i in enumerate(inputs):
    attn_scores_2[i] = torch.dot(x_i, query)
print("attn_scores_2: ", attn_scores_2)

attn_weights_2_tmp = attn_scores_2 / attn_scores_2.sum()
print("Attention weights:", attn_weights_2_tmp)
print("Sum:", attn_weights_2_tmp.sum())

print("\n################################################################\n")

def softmax_naive(x):
    return torch.exp(x) / torch.exp(x).sum(dim=0)

attn_weights_2_naive = softmax_naive(attn_scores_2)
print("Attention weights:", attn_weights_2_naive)
print("Sum:", attn_weights_2_naive.sum())

print("\n################################################################\n")

attn_weights_2 = torch.softmax(attn_scores_2, dim=0)
print("Attention weights:", attn_weights_2)
print("Sum:", attn_weights_2.sum())

print("\n################################################################\n")

query = inputs[1]
print("query: ", query)

context_vec_2 = torch.zeros(query.shape)
for i,x_i in enumerate(inputs):
    context_vec_2 += attn_weights_2[i]*x_i

print("context_vec_2: ",context_vec_2)

print("\n################################################################\n")

attn_scores = torch.empty(6, 6)
for i, x_i in enumerate(inputs):
    for j, x_j in enumerate(inputs):
        attn_scores[i, j] = torch.dot(x_i, x_j)
print(attn_scores)

print("\n################################################################\n")

attn_scores = inputs @ inputs.T
print(attn_scores)

print("\n################################################################\n")

attn_weights = torch.softmax(attn_scores, dim=-1)
print(attn_weights)

print("\n################################################################\n")

row_2_sum = sum([0.1385, 0.2379, 0.2333, 0.1240, 0.1082, 0.1581])
print("Row 2 sum:", row_2_sum)
print("All row sums:", attn_weights.sum(dim=-1))

print("\n################################################################\n")

all_context_vecs = attn_weights @ inputs
print(all_context_vecs)
print("\nPrevious 2nd context vector:", context_vec_2)

print("\n################################################################\n")

x_2 = inputs[1]
d_in = inputs.shape[1]
d_out = 2

torch.manual_seed(123)
W_query = torch.nn.Parameter(torch.rand(d_in, d_out), requires_grad=False)
W_key = torch.nn.Parameter(torch.rand(d_in, d_out), requires_grad=False)
W_value = torch.nn.Parameter(torch.rand(d_in, d_out), requires_grad=False)

query_2 = x_2 @ W_query
key_2 = x_2 @ W_key
value_2 = x_2 @ W_value

print("x_2.shape: ", x_2.shape)
print("W_query.shape: ", W_query.shape)
print("x_2: ", x_2)
print("W_query: ", W_query)
print("query_2: ", query_2)

print("\n################################################################\n")

keys = inputs @ W_key
values = inputs @ W_value
print("keys.shape:", keys.shape)
print("keys: ", keys)
print("values.shape:", values.shape)
print("values: ", values)

print("\n################################################################\n")

keys_2 = keys[1]
attn_score_22 = query_2.dot(keys_2)
print(attn_score_22)

print("\n################################################################\n")

attn_scores_2 = query_2 @ keys.T
print(attn_scores_2)

print("\n################################################################\n")

d_k = keys.shape[-1]
attn_weights_2 = torch.softmax(attn_scores_2 / d_k**0.5, dim=-1)
print(attn_weights_2)

print("\n################################################################\n")

context_vec_2 = attn_weights_2 @ values
print(context_vec_2)

print("\n################################################################\n")