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()

次のトークンの対数発生確率を取得

  • requests コンテキストと予想される継続を含むリクエストのリスト。
  • disable_tqdm True の場合、tqdm プログレスバーを無効にします。
  • 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)