ධූරාවලිට්රාන්ස්ෆෝමර් වඩාත් කාර්යක්ෂම භාෂා ආකෘති වේ

30from typing import List
31
32import torch
33from torch import nn
34
35from labml_helpers.module import Module
36from labml_nn.transformers import MultiHeadAttention, TransformerLayer
37from labml_nn.transformers.feed_forward import FeedForward
38from labml_nn.transformers.utils import subsequent_mask

41class HourGlass(Module):

51    def __init__(self, n_heads: int, d_model: int, dropout: float, d_ff: int, shortening_factors: List[int]):

59        super().__init__()

62        self.pre = TransformerLayer(d_model=d_model,

64                                    self_attn=MultiHeadAttention(n_heads, d_model, dropout),

66                                    feed_forward=FeedForward(d_model, d_ff, dropout),

68                                    dropout_prob=dropout)

70        self.mask = AutoregressiveMask()

73        k = shortening_factors[0]

78        self.shift_right = ShiftRight(k - 1)

81        self.shortening = AvgPoolShortening(k)

84        if len(shortening_factors) == 1:

86            self.shortened = TransformerLayer(d_model=d_model,
87                                              self_attn=MultiHeadAttention(n_heads, d_model, dropout),
88                                              feed_forward=FeedForward(d_model, d_ff, dropout),
89                                              dropout_prob=dropout)

91            self.mask_short = AutoregressiveMask()
92            self.hour_glass = None
93        else:

95            self.hour_glass = HourGlass(n_heads, d_model, dropout, d_ff, shortening_factors[1:])

99        self.up_sampling = NaiveUpSampling(k)

102        self.post = TransformerLayer(d_model=d_model,
103                                     self_attn=MultiHeadAttention(n_heads, d_model, dropout),
104                                     feed_forward=FeedForward(d_model, d_ff, dropout),
105                                     dropout_prob=dropout)

107    def forward(self, x: torch.Tensor):

110        x = self.pre(x=x, mask=self.mask(x))

113        x_short = self.shortening(self.shift_right(x))

117        if self.hour_glass is None:

120            x_short = self.shortened(x=x_short, mask=self.mask_short(x_short))

122        else:

124            x_short = self.hour_glass(x_short)

128        x = x + self.up_sampling(x, x_short)

131        x = self.post(x=x, mask=self.mask(x))

134        return x

137class ShiftRight(Module):

144    def __init__(self, shift: int):

148        super().__init__()

150        assert shift >= 0

152        self.shift = shift

154    def forward(self, x: torch.Tensor):

159        if self.shift == 0:
160            return x

162        prefix = x.new_zeros([self.shift, *x.shape[1:]])

164        return torch.cat([prefix, x[:-self.shift]])

167class AvgPoolShortening(Module):

174    def __init__(self, k: int):

178        super().__init__()

180        self.pool = nn.AvgPool1d(k, ceil_mode=True)

182    def forward(self, x: torch.Tensor):

188        return self.pool(x.permute(1, 2, 0)).permute(2, 0, 1)

191class NaiveUpSampling(Module):

198    def __init__(self, k: int):

202        super().__init__()
203        self.k = k

205    def forward(self, x: torch.Tensor, x_short: torch.Tensor):

211        expanded = torch.repeat_interleave(x_short, self.k, dim=0)

213        expanded = expanded[:x.shape[0]]

216        return expanded

219class AutoregressiveMask(Module):

224    def __init__(self):
225        super().__init__()
226        self.mask = None

228    def forward(self, x: torch.Tensor):

230        if self.mask is None or self.mask.size(0) != len(x):

232            self.mask = subsequent_mask(len(x)).to(x.device)

235        return self.mask

238class LinearPoolingShortening(Module):

246    def __init__(self):
247        super().__init__()
248        raise NotImplementedError

251class AttentionBasedShortening(Module):

263    def __init__(self):
264        super().__init__()
265        raise NotImplementedError

268class LinearUpSampling(Module):

275    def __init__(self):
276        super().__init__()
277        raise NotImplementedError

280class AttentionBasedUpSampling(Module):

292    def __init__(self):
293        super().__init__()
294        raise NotImplementedError