baalajimaestro/llama.cpp
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

Diskcache implementation for llama state.

Browse source
This commit is contained in:
Maximilian-Winter 2023-05-28 15:56:38 +02:00
parent 80066f0b80
commit 719c3eae0a
1 changed files with 17 additions and 14 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

31
llama_cpp/llama.py
View file

@ -17,20 +17,23 @@ from typing import (
) )
from collections import deque, OrderedDict from collections import deque, OrderedDict
import diskcache
from . import llama_cpp from . import llama_cpp
from .llama_types import * from .llama_types import *
class LlamaCache: class LlamaCache:
"""Cache for a llama.cpp model.""" """Cache for a llama.cpp model."""
def __init__(self, capacity_bytes: int = (2 << 30)): def __init__(self, cache_dir="./llama_cache", capacity_bytes: int = (2 << 30)):
self.cache_state: OrderedDict[Tuple[int, ...], "LlamaState"] = OrderedDict() self.cache = diskcache.Cache(cache_dir)
self.capacity_bytes = capacity_bytes self.capacity_bytes = capacity_bytes
@property @property
def cache_size(self): def cache_size(self):
return sum([state.llama_state_size for state in self.cache_state.values()]) return self.cache.volume()
def _find_longest_prefix_key( def _find_longest_prefix_key(
self, self,
@ -38,10 +41,8 @@ class LlamaCache:
) -> Optional[Tuple[int, ...]]: ) -> Optional[Tuple[int, ...]]:
min_len = 0 min_len = 0
min_key = None min_key = None
keys = ( for k in self.cache.iterkeys():
(k, Llama.longest_token_prefix(k, key)) for k in self.cache_state.keys() prefix_len = Llama.longest_token_prefix(k, key)
)
for k, prefix_len in keys:
if prefix_len > min_len: if prefix_len > min_len:
min_len = prefix_len min_len = prefix_len
min_key = k min_key = k
@ -51,9 +52,9 @@ class LlamaCache:
key = tuple(key) key = tuple(key)
_key = self._find_longest_prefix_key(key) _key = self._find_longest_prefix_key(key)
if _key is None: if _key is None:
raise KeyError(f"Key not found") raise KeyError("Key not found")
value = self.cache_state[_key] value = self.cache.pop(_key)
self.cache_state.move_to_end(_key) self.cache.push(_key)
return value return value
def __contains__(self, key: Sequence[int]) -> bool: def __contains__(self, key: Sequence[int]) -> bool:
@ -61,11 +62,13 @@ class LlamaCache:
def __setitem__(self, key: Sequence[int], value: "LlamaState"): def __setitem__(self, key: Sequence[int], value: "LlamaState"):
key = tuple(key) key = tuple(key)
if key in self.cache_state: if key in self.cache:
del self.cache_state[key] del self.cache[key]
self.cache_state[key] = value self.cache[key] = value
while self.cache_size > self.capacity_bytes: while self.cache_size > self.capacity_bytes:
self.cache_state.popitem(last=False) key_to_remove = next(iter(self.cache))
del self.cache[key_to_remove]
class LlamaState: class LlamaState: