(self,
dim,
dim_attn,
dim_ffn,
num_heads,
num_buckets,
shared_pos=True,
dropout=0.1)
| 144 | class T5SelfAttention(nn.Module): |
| 145 | |
| 146 | def __init__(self, |
| 147 | dim, |
| 148 | dim_attn, |
| 149 | dim_ffn, |
| 150 | num_heads, |
| 151 | num_buckets, |
| 152 | shared_pos=True, |
| 153 | dropout=0.1): |
| 154 | super(T5SelfAttention, self).__init__() |
| 155 | self.dim = dim |
| 156 | self.dim_attn = dim_attn |
| 157 | self.dim_ffn = dim_ffn |
| 158 | self.num_heads = num_heads |
| 159 | self.num_buckets = num_buckets |
| 160 | self.shared_pos = shared_pos |
| 161 | |
| 162 | # layers |
| 163 | self.norm1 = T5LayerNorm(dim) |
| 164 | self.attn = T5Attention(dim, dim_attn, num_heads, dropout) |
| 165 | self.norm2 = T5LayerNorm(dim) |
| 166 | self.ffn = T5FeedForward(dim, dim_ffn, dropout) |
| 167 | self.pos_embedding = None if shared_pos else T5RelativeEmbedding( |
| 168 | num_buckets, num_heads, bidirectional=True) |
| 169 | |
| 170 | def forward(self, x, mask=None, pos_bias=None): |
| 171 | e = pos_bias if self.shared_pos else self.pos_embedding( |
nothing calls this directly
no test coverage detected