ට්රාන්ස්ෆෝමර්40

36from typing import List, Optional
37
38import torch
39import torch.nn as nn
40
41from labml_helpers.module import Module
42from labml_nn.utils import clone_module_list
43from .relative_mha import RelativeMultiHeadAttention
44from ..feed_forward import FeedForward

47class TransformerXLLayer(Module):

53    def __init__(self, *,
54                 d_model: int,
55                 self_attn: RelativeMultiHeadAttention,
56                 feed_forward: FeedForward,
57                 dropout_prob: float):

64        super().__init__()
65        self.size = d_model
66        self.self_attn = self_attn
67        self.feed_forward = feed_forward
68        self.dropout = nn.Dropout(dropout_prob)
69        self.norm_self_attn = nn.LayerNorm([d_model])
70        self.norm_ff = nn.LayerNorm([d_model])

72    def forward(self, *,
73                x: torch.Tensor,
74                mem: Optional[torch.Tensor],
75                mask: torch.Tensor):

83        z = self.norm_self_attn(x)

85        if mem is not None:

87            mem = self.norm_self_attn(mem)

89            m_z = torch.cat((mem, z), dim=0)

91        else:
92            m_z = z

94        self_attn = self.self_attn(query=z, key=m_z, value=m_z, mask=mask)

96        x = x + self.dropout(self_attn)

99        z = self.norm_ff(x)

101        ff = self.feed_forward(z)

103        x = x + self.dropout(ff)

106        return x

109class TransformerXL(Module):

116    def __init__(self, layer: TransformerXLLayer, n_layers: int):
117        super().__init__()

119        self.layers = clone_module_list(layer, n_layers)

121        self.norm = nn.LayerNorm([layer.size])

123    def forward(self, x: torch.Tensor, mem: List[torch.Tensor], mask: torch.Tensor):

132        new_mem = []

134        for i, layer in enumerate(self.layers):

136            new_mem.append(x.detach())

138            m = mem[i] if mem else None

140            x = layer(x=x, mem=m, mask=mask)

142        return self.norm(x), new_mem