15import math
16from typing import List
17
18import torch
19import torch.nn.functional as F
20from lm_eval import tasks, evaluator, utils
21from lm_eval.base import BaseLM
22from tokenizers import Tokenizer
23from torch import nn
24from tqdm import tqdm
25
26from labml import monit
27from labml_nn.neox.tokenizer import get_tokenizer
30class EvalHarnessAdapter(BaseLM):
tokenizer
ハギングフェイストークナイザーですvocab_size
はボキャブラリのサイズです (これはトークナイザーのボキャブサイズとは異なります。neoxは埋め込み層モデルを並列化するための追加機能を追加しているからです)。batch_size
はバッチサイズ37 def __init__(self, tokenizer: Tokenizer, vocab_size: int, batch_size: int):
45 super().__init__()
46 self.tokenizer = tokenizer
47 self._eot_token_id = self.tokenizer.token_to_id("<|endoftext|>")
48 self._vocab_size = vocab_size
49
50 self._batch_size = batch_size
ボキャブラリーのサイズ
52 @property
53 def device(self):
54 raise RuntimeError()
55
56 @property
57 def vocab_size(self):
59 return self._vocab_size
テキスト終了トークン
61 @property
62 def eot_token_id(self):
64 return self._eot_token_id
最大シーケンス長
66 @property
67 def max_length(self):
69 return 2048
生成するトークンの最大数
71 @property
72 def max_gen_toks(self):
74 return 128
バッチサイズ
76 @property
77 def batch_size(self):
81 return self._batch_size
与えられたテキストをエンコードする
83 def tok_encode(self, string: str):
87 return self.tokenizer.encode(string).ids
トークン ID からテキストをデコード
89 def tok_decode(self, tokens: List[int]):
93 return self.tokenizer.decode(tokens)
95 def _model_call(self, inps: torch.Tensor):
96 raise NotImplementedError
98 def _model_generate(self, context, max_length, eos_token_id):
99 raise RuntimeError()
101 def greedy_until(self, requests):
102 raise RuntimeError()
104 @torch.no_grad()
105 def _loglikelihood_tokens(self, requests, disable_tqdm=False):
結果について
114 res = []
同じ長さのシーケンスが近くなるように、リクエストを長さの降順に並べ替えます
118 def _collate(x):
119 toks = x[1] + x[2]
120 return -len(toks), tuple(toks)
121
122 reord = utils.Reorderer(requests, _collate)
batch_size
一度に複数のリクエストがあるリクエストをループスルーする
125 for chunk in utils.chunks(tqdm(reord.get_reordered(), disable=disable_tqdm), self.batch_size):
バッチの入力を保存するには
127 inps = []
バッチの継続
129 continuations = []
入力シーケンスの長さ
131 inplens = []
バッチ用のパッド入り長さ
133 padded_length = None
チャンク内の各リクエストをループ処理し、パディング付きの PyTorch テンソルにまとめます。
135 for _, context_enc, continuation_enc in chunk:
コンテキストと続きを連結する
137 inp = context_enc + continuation_enc
サイズが max_length
139 inp = inp[-(self.max_length + 1):]
最終トークンを削除
141 inp = inp[:-1]
テンソルの作成
143 inp = torch.tensor(inp, dtype=torch.long)
入力長さ
145 inplen = inp.shape[0]
パッドの長さを決めてください。短いシーケンスはパディングされます
。149 if padded_length is None:
150 padded_length = int(math.ceil(inplen / 32)) * 32
padded_length = padded_length が Padded_length でない場合はパディングされた_length、それ以外はプレンなし
パディング
154 padding = torch.zeros(padded_length - inplen, dtype=torch.long)
パディングを追加
157 inp = torch.cat([inp, padding], dim=0)
158
159 inps.append(inp)
160 continuations.append(continuation_enc)
161 inplens.append(inplen)
モデルロジットを取得
164 logits = self._model_call(torch.stack(inps))
ログソフトマックスを取得
167 multi_logits = F.log_softmax(logits, dim=-1)
バッチの入力と出力のペアをループ処理します
170 for logits, inplen, cont_toks in zip(multi_logits, inplens, continuations):
予測トークンの数を取得
172 contlen = len(cont_toks)
それらのログを取得
174 logits = logits[inplen - contlen: inplen]
最も確率の高いトークンを手に入れよう
176 greedy_tokens = logits.argmax(dim=-1)
対象トークンを取得
178 cont_toks = torch.tensor(cont_toks, dtype=torch.long).to(logits.device)
完全に一致するかどうか
180 max_equal = (greedy_tokens == cont_toks).all()
対象トークンの対数発生可能性
182 logits = torch.gather(logits, 1, cont_toks[:, None])
対数推定値の合計と、一致したかどうかを結果に加算します。
184 res.append((float(logits.sum()), bool(max_equal)))
並べ替えて結果を返す
187 return reord.get_original(res)
189 @torch.no_grad()
190 def run_eval(self, name: str, eval_tasks: List[str]):
196 results = evaluator.evaluate(lm=self, task_dict=tasks.get_task_dict(eval_tasks))
構成を追加
199 results["config"] = {
200 "name": name,
201 }
204 return results
207class NoeXEvalHarnessAdapter(EvalHarnessAdapter):
model
モデルですtokenizer
ハギングフェイストークナイザーですvocab_size
はボキャブラリのサイズです (これはトークナイザーのボキャブサイズとは異なります。neoxは埋め込み層モデルを並列化するための追加機能を追加しているからです)。batch_size
はバッチサイズdevice
モデルのデバイスです214 def __init__(self, model: nn.Module, tokenizer: Tokenizer, vocab_size: int, batch_size: int, device: torch.device):
224 super().__init__(tokenizer, vocab_size, batch_size)
225 self.model = model
226 self._device = device
モデルに電話する
228 def _model_call(self, inps: torch.Tensor):
232 return self.model(inps.to(self._device))
235def run_eval_harness(model: nn.Module, name: str, eval_tasks: List[str], device: torch.device, batch_size: int = 8):
トークナイザーをロード
241 with monit.section('Load tokenizer'):
242 tokenizer = get_tokenizer()
何も指定されていない場合はすべてのタスク
245 if not eval_tasks:
246 eval_tasks = [
247 "anli_r1",
248 "anli_r2",
249 "anli_r3",
250 "hellaswag",
251 "lambada",
252 "piqa",
253 "winogrande",
254 "wsc",
255 "mathqa",
256 ]
アダプターの作成
259 adapter = NoeXEvalHarnessAdapter(model, tokenizer, 50_432, batch_size, device)
実行
262 return adapter.run_eval(name, eval_tasks)