රේඛීය පක්ෂග්රාහී (අලිබී) අත්හදා බැලීම සමඟ අවධානය යොමු කරන්න

16import torch
17from torch.utils.data import DataLoader
18
19from labml import experiment, tracker
20from labml.configs import option, calculate
21from labml_helpers.datasets.text import SequentialUnBatchedDataset
22from labml_nn.transformers.alibi import AlibiMultiHeadAttention
23from labml_nn.experiments.nlp_autoregression import transpose_batch
24from labml_nn.transformers import TransformerConfigs
25from labml_nn.transformers.gpt import Configs as GPTConfigs

28class Configs(GPTConfigs):

36    transformer: TransformerConfigs = 'GPT_ALiBi'

38    valid_seq_len: int = 128
39    valid_loader = 'shuffled_longer_valid_loader'

41    def other_metrics(self, output: torch.Tensor, target: torch.Tensor):

46        if self.seq_len < output.shape[0]:

48            tracker.add(f'loss.{self.seq_len - 1}.', self.loss_func(output[self.seq_len - 1], target[self.seq_len - 1]))

50            tracker.add(f'loss.0.', self.loss_func(output[0], target[0]))

52        tracker.add(f'loss.{int(output.shape[0]) - 1}.', self.loss_func(output[-1], target[-1]))

55def _alibi_mha(c: TransformerConfigs):

59    return AlibiMultiHeadAttention(c.n_heads, c.d_model, dropout_prob=c.dropout)

63calculate(TransformerConfigs.encoder_attn, 'alibi_mha', _alibi_mha)
64calculate(TransformerConfigs.decoder_attn, 'alibi_mha', _alibi_mha)
65calculate(TransformerConfigs.decoder_mem_attn, 'alibi_mha', _alibi_mha)

68@option(Configs.valid_loader)
69def shuffled_longer_valid_loader(c: Configs):

73    return DataLoader(SequentialUnBatchedDataset(text=c.text.valid,
74                                                 dataset=c.text,
75                                                 seq_len=c.valid_seq_len),
76                      batch_size=c.batch_size,
77                      collate_fn=transpose_batch,
78                      shuffle=True)

81@option(Configs.transformer, 'GPT_ALiBi')
82def _transformer_configs(c: Configs):

89    conf = TransformerConfigs()

91    conf.n_src_vocab = c.n_tokens
92    conf.n_tgt_vocab = c.n_tokens

94    conf.ffn.activation = 'GELU'

97    conf.src_embed = 'no_pos'
98    conf.tgt_embed = 'no_pos'

101    conf.encoder_attn = 'alibi_mha'
102    conf.decoder_attn = 'alibi_mha'
103    conf.decoder_mem_attn = 'alibi_mha'

106    return conf

109def main():

111    experiment.create(name="gpt_alibi")

113    conf = Configs()

115    experiment.configs(conf, {

117        'tokenizer': 'character',

119        'prompt_separator': '',

121        'prompt': 'It is ',

123        'text': 'tiny_shakespeare',

127        'seq_len': 64,

129        'valid_seq_len': 80,

131        'epochs': 128,

133        'batch_size': 128,

136        'inner_iterations': 10,

139        'transformer.d_model': 128,
140        'transformer.ffn.d_ff': 512,
141        'transformer.n_heads': 8,
142        'transformer.n_layers': 4,
143        'transformer.dropout': 0.1,
144    })

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

150    with experiment.start():

152        conf.run()

156if __name__ == '__main__':
157    main()