අංක ගණිත දත්ත කට්ටලය

11importrandom12importstring13fromtypingimportList1415importtorch16fromlabml.loggerimportText17fromtorch.utils.dataimportDataLoader,Dataset1819fromlabmlimportmonit,logger,tracker20fromlabml.configsimportoption21fromlabml\_nn.experiments.nlp\_autoregressionimportNLPAutoRegressionConfigs,transpose\_batch

24classArithmeticDataset(Dataset):

34def\_\_init\_\_(self,seq\_len:int,max\_digits:int,n\_sequences:int):

41self.n\_sequences=n\_sequences42self.max\_digits=max\_digits43self.seq\_len=seq\_len

45self.itos=list(string.digits+'xe =\n?+;')

47self.stoi={c:ifori,cinenumerate(self.itos)}

49@staticmethod50defmake\_int(n\_digits:int):

54res=055foriinrange(n\_digits):56d=random.randrange(1,11)ifi==0elserandom.randrange(0,11)57res=res\*10+d5859returnres

61@staticmethod62defget\_add\_explanation(x:int,y:int):

69carry=070e=071explanation=[]72whilex\>0ory\>0orcarry\>0:73rx,ry=x%10,y%1074total=rx+ry+carry75explanation.append(f"{rx}e{e}+{ry}e{e}+{carry}e{e}=={total}e{e}")76x,y,carry=x//10,y//10,total//1077e+=17879return' '.join(explanation)

82defmake\_add\_problem(self):

86x=self.make\_int(n\_digits=random.randrange(1,self.max\_digits+1))87y=self.make\_int(n\_digits=random.randrange(1,self.max\_digits+1))8889explanation=self.get\_add\_explanation(x,y)90returnf"x={x}+{y}; {explanation} x=={x+y}\n"

92defget\_qa(self):

96x=self.make\_int(n\_digits=random.randrange(1,self.max\_digits+1))97y=self.make\_int(n\_digits=random.randrange(1,self.max\_digits+1))9899returnf'x={x}+{y};',f'{x+y}'

101defget\_packed\_math\_input(self):

105s\_enc=[]106whilelen(s\_enc)\<=self.seq\_len:107s\_part=self.make\_add\_problem()108s\_part\_enc=self.encode('?'+s\_part)109s\_enc=s\_enc+s\_part\_enc110returns\_enc

112defencode(self,s:str):

116return[self.stoi[c]forcins]

118defdecode(self,arr:List[int]):

122return''.join([self.itos[c]forcinarr])

124def\_\_getitem\_\_(self,idx:int):

128s=torch.tensor(self.get\_packed\_math\_input())129returns[:self.seq\_len],s[1:self.seq\_len+1]

131def\_\_len\_\_(self):

135returnself.n\_sequences

138classArithmeticAutoregression(NLPAutoRegressionConfigs):

143max\_digits:int=4

145train\_sequences\_per\_epoch:int=2\*\*12

147train\_loader:DataLoader='arithmetic\_train\_loader'

149n\_tests:int=64

151validator=None

153inner\_iterations=4

155n\_tokens=len(ArithmeticDataset(1,1,1).itos)

[email protected]\_grad()158defsample(self):

166ifself.training\_loop.idx\<1:167return

170dataset=ArithmeticDataset(self.seq\_len,self.max\_digits,1)

172qa=[dataset.get\_qa()for\_inrange(self.n\_tests)]

174questions=[p[0]forpinqa]

177data=torch.tensor([[dataset.stoi[p[0]]forpinquestions]])

179data=data.to(self.device)

182finished=torch.zeros((len(questions),)).bool().to(self.device)

184new\_line=dataset.stoi['\n']

187results=[p[0]forpinquestions]

190foriinmonit.iterate('Sample',self.seq\_len-1):

192iffinished.sum()==len(finished):193continue

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

198output=output[-1].argmax(dim=-1)

201finished=finished|(output==new\_line)

203iffinished.sum()==len(finished):204continue

207forj,pinenumerate(questions):208iflen(p)\>i+1:209output[j]=dataset.stoi[p[i+1]]

212data=torch.cat([data,output[None,:]],dim=0)

215forj,cinenumerate(output):216results[j]+=dataset.itos[c]

219results=[r.split('\n')[0]forrinresults]

222res\_sample=results[0].split(';')223logger.log([(res\_sample[0],Text.key),(';',Text.subtle),(';'.join(res\_sample[1:]),Text.none)])

226results=[r.split('x==')[-1]forrinresults]

229correct=0230forr,\_qainzip(results,qa):231ifr==\_qa[1]:232correct+=1

235tracker.save('score',correct/len(results))

238@option(ArithmeticAutoregression.train\_loader)239defarithmetic\_train\_loader(c:ArithmeticAutoregression):

243returnDataLoader(ArithmeticDataset(c.seq\_len,c.max\_digits,c.train\_sequences\_per\_epoch),244batch\_size=c.batch\_size,245collate\_fn=transpose\_batch,246num\_workers=4)

249def\_test():

253dataset=ArithmeticDataset(256,8,10)254255print(dataset.decode(dataset.get\_packed\_math\_input()))

259if\_\_name\_\_=='\_\_main\_\_':260\_test()