15importmath16fromtypingimportList1718importtorch19importtorch.nn.functionalasF20fromlm\_evalimporttasks,evaluator,utils21fromlm\_eval.baseimportBaseLM22fromtokenizersimportTokenizer23fromtorchimportnn24fromtqdmimporttqdm2526fromlabmlimportmonit27fromlabml\_nn.neox.tokenizerimportget\_tokenizer

30classEvalHarnessAdapter(BaseLM):

37def\_\_init\_\_(self,tokenizer:Tokenizer,vocab\_size:int,batch\_size:int):

45super().\_\_init\_\_()46self.tokenizer=tokenizer47self.\_eot\_token\_id=self.tokenizer.token\_to\_id("\<|endoftext|\>")48self.\_vocab\_size=vocab\_size4950self.\_batch\_size=batch\_size

52@property53defdevice(self):54raiseRuntimeError()5556@property57defvocab\_size(self):

59returnself.\_vocab\_size

61@property62defeot\_token\_id(self):

64returnself.\_eot\_token\_id

66@property67defmax\_length(self):

69return2048

71@property72defmax\_gen\_toks(self):

74return128

76@property77defbatch\_size(self):

81returnself.\_batch\_size

83deftok\_encode(self,string:str):

87returnself.tokenizer.encode(string).ids

89deftok\_decode(self,tokens:List[int]):

93returnself.tokenizer.decode(tokens)

95def\_model\_call(self,inps:torch.Tensor):96raiseNotImplementedError

98def\_model\_generate(self,context,max\_length,eos\_token\_id):99raiseRuntimeError()

101defgreedy\_until(self,requests):102raiseRuntimeError()

[email protected]\_grad()105def\_loglikelihood\_tokens(self,requests,disable\_tqdm=False):

114res=[]

118def\_collate(x):119toks=x[1]+x[2]120return-len(toks),tuple(toks)121122reord=utils.Reorderer(requests,\_collate)

125forchunkinutils.chunks(tqdm(reord.get\_reordered(),disable=disable\_tqdm),self.batch\_size):

127inps=[]

129continuations=[]

131inplens=[]

133padded\_length=None

135for\_,context\_enc,continuation\_encinchunk:

137inp=context\_enc+continuation\_enc

139inp=inp[-(self.max\_length+1):]

141inp=inp[:-1]

143inp=torch.tensor(inp,dtype=torch.long)

145inplen=inp.shape[0]

149ifpadded\_lengthisNone:150padded\_length=int(math.ceil(inplen/32))\*32

154padding=torch.zeros(padded\_length-inplen,dtype=torch.long)

157inp=torch.cat([inp,padding],dim=0)158159inps.append(inp)160continuations.append(continuation\_enc)161inplens.append(inplen)

164logits=self.\_model\_call(torch.stack(inps))

167multi\_logits=F.log\_softmax(logits,dim=-1)

170forlogits,inplen,cont\_toksinzip(multi\_logits,inplens,continuations):

172contlen=len(cont\_toks)

174logits=logits[inplen-contlen:inplen]

176greedy\_tokens=logits.argmax(dim=-1)

178cont\_toks=torch.tensor(cont\_toks,dtype=torch.long).to(logits.device)

180max\_equal=(greedy\_tokens==cont\_toks).all()

182logits=torch.gather(logits,1,cont\_toks[:,None])

184res.append((float(logits.sum()),bool(max\_equal)))

187returnreord.get\_original(res)

[email protected]\_grad()190defrun\_eval(self,name:str,eval\_tasks:List[str]):

196results=evaluator.evaluate(lm=self,task\_dict=tasks.get\_task\_dict(eval\_tasks))

199results["config"]={200"name":name,201}

204returnresults

207classNoeXEvalHarnessAdapter(EvalHarnessAdapter):

214def\_\_init\_\_(self,model:nn.Module,tokenizer:Tokenizer,vocab\_size:int,batch\_size:int,device:torch.device):

224super().\_\_init\_\_(tokenizer,vocab\_size,batch\_size)225self.model=model226self.\_device=device

228def\_model\_call(self,inps:torch.Tensor):

232returnself.model(inps.to(self.\_device))

235defrun\_eval\_harness(model:nn.Module,name:str,eval\_tasks:List[str],device:torch.device,batch\_size:int=8):

241withmonit.section('Load tokenizer'):242tokenizer=get\_tokenizer()

245ifnoteval\_tasks:246eval\_tasks=[247"anli\_r1",248"anli\_r2",249"anli\_r3",250"hellaswag",251"lambada",252"piqa",253"winogrande",254"wsc",255"mathqa",256]

259adapter=NoeXEvalHarnessAdapter(model,tokenizer,50\_432,batch\_size,device)

262returnadapter.run\_eval(name,eval\_tasks)