ධූරාවලි ට්රාන්ස්ෆෝමර් වඩාත් කාර්යක්ෂම භාෂා ආකෘති අත්හදා බැලීම

16import math
17from typing import List
18
19import torch
20from torch import nn
21
22from labml import experiment
23from labml.configs import option
24from labml_helpers.module import Module
25from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs
26from labml_nn.transformers.hour_glass import HourGlass
27from labml_nn.transformers.positional_encoding import PositionalEncoding

30class AutoregressiveTransformer(Module):

35    def __init__(self, n_tokens: int, d_model: int, dropout: float, hour_glass: HourGlass):

42        super().__init__()

44        self.embedding = nn.Embedding(n_tokens, d_model)

50        self.pos_embedding = PositionalEncoding(d_model, dropout)

52        self.hour_glass = hour_glass

54        self.norm = nn.LayerNorm([d_model])

56        self.d_model = d_model

58        self.output = nn.Linear(d_model, n_tokens)

60    def __call__(self, x: torch.Tensor):

65        x = self.embedding(x)

68        if self.pos_embedding is not None:
69            x = self.pos_embedding(x * math.sqrt(self.d_model))

72        x = self.hour_glass(x)

75        output = self.output(self.norm(x))

78        return output, None

81class Configs(NLPAutoRegressionConfigs):

89    model: AutoregressiveTransformer

91    n_heads: int = 8

93    dropout: float = 0.1

95    d_ff: int = 512

97    d_model: int = 256

99    shortening_factors: List[int] = [8, 4]

102@option(Configs.model)
103def _model(c: Configs):

109    hour_glass = HourGlass(c.n_heads, c.d_model, c.dropout, c.d_ff, c.shortening_factors)

111    m = AutoregressiveTransformer(c.n_tokens, c.d_model, c.dropout, hour_glass).to(c.device)

114    return m

117def main():

119    experiment.create(name="hour_glass")

121    conf = Configs()

123    experiment.configs(conf, {

125        'tokenizer': 'character',

127        'prompt_separator': '',

129        'prompt': 'It is ',

131        'text': 'tiny_shakespeare',

134        'seq_len': 256,

136        'epochs': 128,

138        'batch_size': 32,

141        'inner_iterations': 10,

144        'optimizer.optimizer': 'Noam',
145        'optimizer.learning_rate': 1.,

147    })

150    experiment.add_pytorch_models({'model': conf.model})

153    with experiment.start():

155        conf.run()

159if __name__ == '__main__':
160    main()