#

中間アクティベーション用キャッシュ

推論中、モデルはトークンごとに出力します。このシンプルなキャッシュを使ってキーと値のアテンションレイヤーを格納するので、以前のトークンでそれらを再計算する必要がありません

。

15from typing import Any

#

キャッシュ

これにより、キーと値のキャッシュが維持され、値のプッシュとポップが同じ順序でキューに入れられます。キューは複数のアテンションレイヤーがあるので便利です

。

18class Cache:

#

26    def __init__(self):
27        self._cache = {}

#

キャッシュをクリア

29    def clear_all(self):

#

33        self._cache = {}

#

値をキューにプッシュ

name はキューの名前です
value プッシュする値です

35    def push(self, name: str, value: Any):

#

存在しない場合は空のキューを作成

44        if name not in self._cache:
45            self._cache[name] = []

#

キューにプッシュ

48        self._cache[name].append(value)

#

キューのサイズを返す

name はキューの名前です

キューが存在する場合はキューのサイズを返します。それ以外の場合は None

50    def q_size(self, name):

#

58        if name not in self._cache:
59            return None
60
61        if type(self._cache[name]) != list:
62            return None
63
64        return len(self._cache[name])

#

キューからポップする

name はキューの名前です

値を返す

66    def pop(self, name: str):

#

73        return self._cache[name].pop(0)

#

値をキャッシュする

key キャッシュされる値の名前
value は値です

75    def set(self, key: str, value: Any):

#

82        self._cache[key] = value

#

キャッシュから値を取得

key キャッシュ時に使用される名前です
default キャッシュが空の場合のデフォルト値です

キャッシュされた値を返します

84    def get(self, key: str, default: Any = None):

#

92        return self._cache.get(key, default)

#

キャッシュ値をクリアする

key キャッシュ時に使用される名前です

94    def clear(self, key: str):

#

100        del self._cache[key]

#

キャッシュ用シングルトン

104_INSTANCE = None

#

キャッシュインスタンスを取得

キャッシュインスタンスを返します

107def get_cache() -> Cache:

#

113    global _INSTANCE
114
115    if _INSTANCE is None:
116        _INSTANCE = Cache()
117
118    return _INSTANCE