Hugging Face's logo

Spaces:
DebasishDhal99
/
Embedding-Visualization
Running on Zero

App Files Community
Embedding-Visualization
File size: 3,075 Bytes
8038560
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
# In simple terms, we are calculating this
### Sm,n = q_m^T * k_n + b_m-n ## here m and n are the positions of the query and key vectors respectively. The similarity score must depend on the relative position of the query and key vectors as well. b is some learned function of the relative position.

# Some minor lacuna pending here in the implementation, needs to be cleared.

import numpy as np

import numpy as np


def softmax(x, axis=-1):
    x = x - np.max(x, axis=axis, keepdims=True)
    exp_x = np.exp(x)
    return exp_x / np.sum(exp_x, axis=axis, keepdims=True)


def shaw_relative_attention(

    X,

    Wq,

    Wk,

    Wv,

    relative_key_embeddings,

    relative_value_embeddings,

    max_relative_position

):
    n_tokens, d_model = X.shape

    # --------------------------------------------------
    # 1. Relative positions
    # --------------------------------------------------

    positions = np.arange(n_tokens)

    relative_positions = (
        positions[:, None]
        - positions[None, :]
    )

    relative_positions = np.clip(
        relative_positions,
        -max_relative_position,
        max_relative_position
    )

    # Convert [-max, ..., +max] → [0, ..., 2*max]
    relative_indices = (
        relative_positions
        + max_relative_position
    )

    # --------------------------------------------------
    # 2. Query
    # --------------------------------------------------

    Q = X @ Wq

    # --------------------------------------------------
    # 3. Relative Key embeddings
    # --------------------------------------------------

    relative_key = (
        relative_key_embeddings[
            relative_indices
        ]
    )

    # Shape:
    # (n_tokens, n_tokens, d_model)

    # x_n + relative positional embedding
    K_input = (
        X[None, :, :]
        + relative_key
    )

    # Apply Wk
    K_relative = K_input @ Wk

    # --------------------------------------------------
    # 4. Attention scores
    # --------------------------------------------------

    scores = np.einsum(
        "md,mnd->mn",
        Q,
        K_relative
    )

    # --------------------------------------------------
    # 5. Relative Value embeddings
    # --------------------------------------------------

    relative_value = (
        relative_value_embeddings[
            relative_indices
        ]
    )

    V_input = (
        X[None, :, :]
        + relative_value
    )

    V_relative = V_input @ Wv

    # --------------------------------------------------
    # 6. Attention weights
    # --------------------------------------------------

    attention_weights = softmax(
        scores,
        axis=-1
    )

    # --------------------------------------------------
    # 7. Weighted Values
    # --------------------------------------------------

    output = np.einsum(
        "mn,mnd->md",
        attention_weights,
        V_relative
    )

    return output