16importmath17fromtypingimportList1819importtorch20fromtorchimportnn2122fromlabmlimportexperiment23fromlabml.configsimportoption24fromlabml\_helpers.moduleimportModule25fromlabml\_nn.experiments.nlp\_autoregressionimportNLPAutoRegressionConfigs26fromlabml\_nn.transformers.hour\_glassimportHourGlass27fromlabml\_nn.transformers.positional\_encodingimportPositionalEncoding

30classAutoregressiveTransformer(Module):

35def\_\_init\_\_(self,n\_tokens:int,d\_model:int,dropout:float,hour\_glass:HourGlass):

42super().\_\_init\_\_()

44self.embedding=nn.Embedding(n\_tokens,d\_model)

50self.pos\_embedding=PositionalEncoding(d\_model,dropout)

52self.hour\_glass=hour\_glass

54self.norm=nn.LayerNorm([d\_model])

56self.d\_model=d\_model

58self.output=nn.Linear(d\_model,n\_tokens)

60def\_\_call\_\_(self,x:torch.Tensor):

65x=self.embedding(x)

68ifself.pos\_embeddingisnotNone:69x=self.pos\_embedding(x\*math.sqrt(self.d\_model))

72x=self.hour\_glass(x)

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

78returnoutput,None

81classConfigs(NLPAutoRegressionConfigs):

89model:AutoregressiveTransformer

91n\_heads:int=8

93dropout:float=0.1

95d\_ff:int=512

97d\_model:int=256

99shortening\_factors:List[int]=[8,4]

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

109hour\_glass=HourGlass(c.n\_heads,c.d\_model,c.dropout,c.d\_ff,c.shortening\_factors)

111m=AutoregressiveTransformer(c.n\_tokens,c.d\_model,c.dropout,hour\_glass).to(c.device)

114returnm

117defmain():

119experiment.create(name="hour\_glass")

121conf=Configs()

123experiment.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})

150experiment.add\_pytorch\_models({'model':conf.model})

153withexperiment.start():

155conf.run()

159if\_\_name\_\_=='\_\_main\_\_':160main()