hymenjj/llama-cpp-python-prebuilt

1import sys2from abc import ABC, abstractmethod3from typing import (4    Optional,5    Sequence,6    Tuple,7)8from collections import OrderedDict9 10import diskcache11 12import llama_cpp.llama13 14from .llama_types import *15 16 17class BaseLlamaCache(ABC):18    """Base cache class for a llama.cpp model."""19 20    def __init__(self, capacity_bytes: int = (2 << 30)):21        self.capacity_bytes = capacity_bytes22 23    @property24    @abstractmethod25    def cache_size(self) -> int:26        raise NotImplementedError27 28    def _find_longest_prefix_key(29        self,30        key: Tuple[int, ...],31    ) -> Optional[Tuple[int, ...]]:32        pass33 34    @abstractmethod35    def __getitem__(self, key: Sequence[int]) -> "llama_cpp.llama.LlamaState":36        raise NotImplementedError37 38    @abstractmethod39    def __contains__(self, key: Sequence[int]) -> bool:40        raise NotImplementedError41 42    @abstractmethod43    def __setitem__(44        self, key: Sequence[int], value: "llama_cpp.llama.LlamaState"45    ) -> None:46        raise NotImplementedError47 48 49class LlamaRAMCache(BaseLlamaCache):50    """Cache for a llama.cpp model using RAM."""51 52    def __init__(self, capacity_bytes: int = (2 << 30)):53        super().__init__(capacity_bytes)54        self.capacity_bytes = capacity_bytes55        self.cache_state: OrderedDict[56            Tuple[int, ...], "llama_cpp.llama.LlamaState"57        ] = OrderedDict()58 59    @property60    def cache_size(self):61        return sum([state.llama_state_size for state in self.cache_state.values()])62 63    def _find_longest_prefix_key(64        self,65        key: Tuple[int, ...],66    ) -> Optional[Tuple[int, ...]]:67        min_len = 068        min_key = None69        keys = (70            (k, llama_cpp.llama.Llama.longest_token_prefix(k, key))71            for k in self.cache_state.keys()72        )73        for k, prefix_len in keys:74            if prefix_len > min_len:75                min_len = prefix_len76                min_key = k77        return min_key78 79    def __getitem__(self, key: Sequence[int]) -> "llama_cpp.llama.LlamaState":80        key = tuple(key)81        _key = self._find_longest_prefix_key(key)82        if _key is None:83            raise KeyError("Key not found")84        value = self.cache_state[_key]85        self.cache_state.move_to_end(_key)86        return value87 88    def __contains__(self, key: Sequence[int]) -> bool:89        return self._find_longest_prefix_key(tuple(key)) is not None90 91    def __setitem__(self, key: Sequence[int], value: "llama_cpp.llama.LlamaState"):92        key = tuple(key)93        if key in self.cache_state:94            del self.cache_state[key]95        self.cache_state[key] = value96        while self.cache_size > self.capacity_bytes and len(self.cache_state) > 0:97            self.cache_state.popitem(last=False)98 99 100# Alias for backwards compatibility101LlamaCache = LlamaRAMCache102 103 104class LlamaDiskCache(BaseLlamaCache):105    """Cache for a llama.cpp model using disk."""106 107    def __init__(108        self, cache_dir: str = ".cache/llama_cache", capacity_bytes: int = (2 << 30)109    ):110        super().__init__(capacity_bytes)111        self.cache = diskcache.Cache(cache_dir)112 113    @property114    def cache_size(self):115        return int(self.cache.volume())  # type: ignore116 117    def _find_longest_prefix_key(118        self,119        key: Tuple[int, ...],120    ) -> Optional[Tuple[int, ...]]:121        min_len = 0122        min_key: Optional[Tuple[int, ...]] = None123        for k in self.cache.iterkeys():  # type: ignore124            prefix_len = llama_cpp.llama.Llama.longest_token_prefix(k, key)125            if prefix_len > min_len:126                min_len = prefix_len127                min_key = k  # type: ignore128        return min_key129 130    def __getitem__(self, key: Sequence[int]) -> "llama_cpp.llama.LlamaState":131        key = tuple(key)132        _key = self._find_longest_prefix_key(key)133        if _key is None:134            raise KeyError("Key not found")135        value: "llama_cpp.llama.LlamaState" = self.cache.pop(_key)  # type: ignore136        # NOTE: This puts an integer as key in cache, which breaks,137        # Llama.longest_token_prefix(k, key) above since k is not a tuple of ints/tokens138        # self.cache.push(_key, side="front")  # type: ignore139        return value140 141    def __contains__(self, key: Sequence[int]) -> bool:142        return self._find_longest_prefix_key(tuple(key)) is not None143 144    def __setitem__(self, key: Sequence[int], value: "llama_cpp.llama.LlamaState"):145        print("LlamaDiskCache.__setitem__: called", file=sys.stderr)146        key = tuple(key)147        if key in self.cache:148            print("LlamaDiskCache.__setitem__: delete", file=sys.stderr)149            del self.cache[key]150        self.cache[key] = value151        print("LlamaDiskCache.__setitem__: set", file=sys.stderr)152        while self.cache_size > self.capacity_bytes and len(self.cache) > 0:153            key_to_remove = next(iter(self.cache))154            del self.cache[key_to_remove]155        print("LlamaDiskCache.__setitem__: trim", file=sys.stderr)156