推論中、モデルはトークンごとに出力します。このシンプルなキャッシュを使ってキーと値のアテンションレイヤーを格納するので、以前のトークンでそれらを再計算する必要がありません
。15from typing import Any
18class Cache:
26 def __init__(self):
27 self._cache = {}
29 def clear_all(self):
33 self._cache = {}
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)
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])
66 def pop(self, name: str):
73 return self._cache[name].pop(0)
75 def set(self, key: str, value: Any):
82 self._cache[key] = value
84 def get(self, key: str, default: Any = None):
92 return self._cache.get(key, default)
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