11fromcollectionsimportCounter12fromtypingimportCallable1314importtorch15importtorchtext16fromtorchimportnn17fromtorch.utils.dataimportDataLoader18importtorchtext.vocab19fromtorchtext.vocabimportVocab2021fromlabmlimportlab,tracker,monit22fromlabml.configsimportoption23fromlabml\_helpers.deviceimportDeviceConfigs24fromlabml\_helpers.metrics.accuracyimportAccuracy25fromlabml\_helpers.moduleimportModule26fromlabml\_helpers.train\_validimportTrainValidConfigs,hook\_model\_outputs,BatchIndex27fromlabml\_nn.optimizers.configsimportOptimizerConfigs

30classNLPClassificationConfigs(TrainValidConfigs):

41optimizer:torch.optim.Adam

43device:torch.device=DeviceConfigs()

46model:Module

48batch\_size:int=16

50seq\_len:int=512

52vocab:Vocab='ag\_news'

54n\_tokens:int

56n\_classes:int='ag\_news'

58tokenizer:Callable='character'

61is\_save\_models=True

64loss\_func=nn.CrossEntropyLoss()

66accuracy=Accuracy()

68d\_model:int=512

70grad\_norm\_clip:float=1.0

73train\_loader:DataLoader='ag\_news'

75valid\_loader:DataLoader='ag\_news'

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

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

87definit(self):

92tracker.set\_scalar("accuracy.\*",True)93tracker.set\_scalar("loss.\*",True)

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

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

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

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

111ifself.mode.is\_train:112tracker.add\_global\_step(data.shape[1])

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

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

122loss=self.loss\_func(output,target)123tracker.add("loss.",loss)

126self.accuracy(output,target)127self.accuracy.track()

130ifself.mode.is\_train:

132loss.backward()

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

136self.optimizer.step()

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

141self.optimizer.zero\_grad()

144tracker.save()

147@option(NLPClassificationConfigs.optimizer)148def\_optimizer(c:NLPClassificationConfigs):

153optimizer=OptimizerConfigs()154optimizer.parameters=c.model.parameters()155optimizer.optimizer='Adam'156optimizer.d\_model=c.d\_model157158returnoptimizer

'tokenizer': 'basic_english',

161@option(NLPClassificationConfigs.tokenizer)162defbasic\_english():

176fromtorchtext.dataimportget\_tokenizer177returnget\_tokenizer('basic\_english')

180defcharacter\_tokenizer(x:str):

184returnlist(x)

187@option(NLPClassificationConfigs.tokenizer)188defcharacter():

192returncharacter\_tokenizer

195@option(NLPClassificationConfigs.n\_tokens)196def\_n\_tokens(c:NLPClassificationConfigs):

200returnlen(c.vocab)+2

203classCollateFunc:

208def\_\_init\_\_(self,tokenizer,vocab:Vocab,seq\_len:int,padding\_token:int,classifier\_token:int):

216self.classifier\_token=classifier\_token217self.padding\_token=padding\_token218self.seq\_len=seq\_len219self.vocab=vocab220self.tokenizer=tokenizer

222def\_\_call\_\_(self,batch):

228data=torch.full((self.seq\_len,len(batch)),self.padding\_token,dtype=torch.long)

230labels=torch.zeros(len(batch),dtype=torch.long)

233for(i,(\_label,\_text))inenumerate(batch):

235labels[i]=int(\_label)-1

237\_text=[self.vocab[token]fortokeninself.tokenizer(\_text)]

239\_text=\_text[:self.seq\_len]

241data[:len(\_text),i]=data.new\_tensor(\_text)

244data[-1,:]=self.classifier\_token

247returndata,labels

250@option([NLPClassificationConfigs.n\_classes,251NLPClassificationConfigs.vocab,252NLPClassificationConfigs.train\_loader,253NLPClassificationConfigs.valid\_loader])254defag\_news(c:NLPClassificationConfigs):

263train,valid=torchtext.datasets.AG\_NEWS(root=str(lab.get\_data\_path()/'ag\_news'),split=('train','test'))

266withmonit.section('Load data'):267fromlabml\_nn.utilsimportMapStyleDataset

270train,valid=MapStyleDataset(train),MapStyleDataset(valid)

273tokenizer=c.tokenizer

276counter=Counter()

278for(label,line)intrain:279counter.update(tokenizer(line))

281for(label,line)invalid:282counter.update(tokenizer(line))

284vocab=torchtext.vocab.vocab(counter,min\_freq=1)

287train\_loader=DataLoader(train,batch\_size=c.batch\_size,shuffle=True,288collate\_fn=CollateFunc(tokenizer,vocab,c.seq\_len,len(vocab),len(vocab)+1))

290valid\_loader=DataLoader(valid,batch\_size=c.batch\_size,shuffle=True,291collate\_fn=CollateFunc(tokenizer,vocab,c.seq\_len,len(vocab),len(vocab)+1))

294return4,vocab,train\_loader,valid\_loader