11fromtypingimportCallable1213importtorch14importtorch.nnasnn15fromtorch.utils.dataimportDataLoader,RandomSampler1617fromlabmlimportlab,monit,logger,tracker18fromlabml.configsimportoption19fromlabml.loggerimportText20fromlabml\_helpers.datasets.textimportTextDataset,SequentialDataLoader,SequentialUnBatchedDataset,TextFileDataset21fromlabml\_helpers.deviceimportDeviceConfigs22fromlabml\_helpers.metrics.accuracyimportAccuracy23fromlabml\_helpers.moduleimportModule24fromlabml\_helpers.train\_validimportTrainValidConfigs,hook\_model\_outputs,BatchIndex25fromlabml\_nn.optimizers.configsimportOptimizerConfigs

28classCrossEntropyLoss(Module):

33def\_\_init\_\_(self):34super().\_\_init\_\_()35self.loss=nn.CrossEntropyLoss()

37defforward(self,outputs,targets):38returnself.loss(outputs.view(-1,outputs.shape[-1]),targets.view(-1))

41classNLPAutoRegressionConfigs(TrainValidConfigs):

52optimizer:torch.optim.Adam

54device:torch.device=DeviceConfigs()

57model:Module

59text:TextDataset

61batch\_size:int=16

63seq\_len:int=512

65n\_tokens:int

67tokenizer:Callable='character'

70prompt:str

72prompt\_separator:str

75is\_save\_models=True

78loss\_func=CrossEntropyLoss()

80accuracy=Accuracy()

82d\_model:int=512

84grad\_norm\_clip:float=1.0

87train\_loader:DataLoader='shuffled\_train\_loader'

89valid\_loader:DataLoader='shuffled\_valid\_loader'

92dataloader\_shuffle\_with\_replacement:bool=False

97is\_log\_model\_params\_grads:bool=False

102is\_log\_model\_activations:bool=False

104definit(self):

109tracker.set\_scalar("accuracy.\*",True)110tracker.set\_scalar("loss.\*",True)111tracker.set\_text("sampled",False)

113hook\_model\_outputs(self.mode,self.model,'model')

118self.state\_modules=[self.accuracy]

120defother\_metrics(self,output:torch.Tensor,target:torch.Tensor):

122pass

124defstep(self,batch:any,batch\_idx:BatchIndex):

130self.model.train(self.mode.is\_train)

133data,target=batch[0].to(self.device),batch[1].to(self.device)

136ifself.mode.is\_train:137tracker.add\_global\_step(data.shape[0]\*data.shape[1])

140withself.mode.update(is\_log\_activations=batch\_idx.is\_lastandself.is\_log\_model\_activations):

144output,\*\_=self.model(data)

147loss=self.loss\_func(output,target)148tracker.add("loss.",loss)

151self.accuracy(output,target)152self.accuracy.track()153154self.other\_metrics(output,target)

157ifself.mode.is\_train:

159loss.backward()

161torch.nn.utils.clip\_grad\_norm\_(self.model.parameters(),max\_norm=self.grad\_norm\_clip)

163self.optimizer.step()

165ifbatch\_idx.is\_lastandself.is\_log\_model\_params\_grads:166tracker.add('model',self.model)

168self.optimizer.zero\_grad()

171tracker.save()

173defsample(self):

179prompt=self.prompt

181log=[(prompt,Text.subtle)]

183foriinmonit.iterate('Sample',25):

185data=self.text.text\_to\_i(prompt).unsqueeze(-1)186data=data.to(self.device)

188output,\*\_=self.model(data)

190output=output.argmax(dim=-1).squeeze()

192prompt+=self.prompt\_separator+self.text.itos[output[-1]]

194log+=[(self.prompt\_separator+self.text.itos[output[-1]],Text.value)]195196tracker.add({'sampled':prompt})

198logger.log(log)

201@option(NLPAutoRegressionConfigs.optimizer)202def\_optimizer(c:NLPAutoRegressionConfigs):

207optimizer=OptimizerConfigs()208optimizer.parameters=c.model.parameters()209optimizer.optimizer='Adam'210optimizer.d\_model=c.d\_model211212returnoptimizer

215@option(NLPAutoRegressionConfigs.n\_tokens)216def\_n\_tokens(c:NLPAutoRegressionConfigs):

220returnc.text.n\_tokens

'tokenizer': 'basic_english',

223@option(NLPAutoRegressionConfigs.tokenizer)224defbasic\_english():

238fromtorchtext.dataimportget\_tokenizer239returnget\_tokenizer('basic\_english')

242defcharacter\_tokenizer(x:str):

246returnlist(x)

249@option(NLPAutoRegressionConfigs.tokenizer)250defcharacter():

254returncharacter\_tokenizer

257@option(NLPAutoRegressionConfigs.text)258deftiny\_shakespeare(c:NLPAutoRegressionConfigs):

264returnTextFileDataset(265lab.get\_data\_path()/'tiny\_shakespeare.txt',266c.tokenizer,267url='https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt')

270@option(NLPAutoRegressionConfigs.train\_loader)271defsequential\_train\_loader(c:NLPAutoRegressionConfigs):

275returnSequentialDataLoader(text=c.text.train,276dataset=c.text,277batch\_size=c.batch\_size,278seq\_len=c.seq\_len)

281@option(NLPAutoRegressionConfigs.valid\_loader)282defsequential\_valid\_loader(c:NLPAutoRegressionConfigs):

286returnSequentialDataLoader(text=c.text.valid,287dataset=c.text,288batch\_size=c.batch\_size,289seq\_len=c.seq\_len)

292deftranspose\_batch(batch):

300transposed\_data=list(zip(\*batch))

302src=torch.stack(transposed\_data[0],dim=1)303tgt=torch.stack(transposed\_data[1],dim=1)304305returnsrc,tgt

308@option(NLPAutoRegressionConfigs.train\_loader)309defshuffled\_train\_loader(c:NLPAutoRegressionConfigs):

313dataset=SequentialUnBatchedDataset(text=c.text.train,314dataset=c.text,315seq\_len=c.seq\_len)316sampler=RandomSampler(dataset,replacement=c.dataloader\_shuffle\_with\_replacement)317318returnDataLoader(dataset,319batch\_size=c.batch\_size,320collate\_fn=transpose\_batch,321sampler=sampler)

324@option(NLPAutoRegressionConfigs.valid\_loader)325defshuffled\_valid\_loader(c:NLPAutoRegressionConfigs):

329dataset=SequentialUnBatchedDataset(text=c.text.valid,330dataset=c.text,331seq\_len=c.seq\_len)332sampler=RandomSampler(dataset,replacement=c.dataloader\_shuffle\_with\_replacement)333334returnDataLoader(dataset,335batch\_size=c.batch\_size,336collate\_fn=transpose\_batch,337sampler=sampler)