Added both LlamaChache classes Disk and RAM.

llama_cpp/llama.py

@@ -4,6 +4,7 @@ import uuid
 import time
 import math
 import multiprocessing
+from abc import ABC
 from typing import (
     List,
     Optional,
@@ -26,13 +27,86 @@ import numpy as np
 import numpy.typing as npt
 
 
+class LlamaCache(ABC):
+    """Base cache class for a llama.cpp model."""
 
-class LlamaCache:
+    def __init__(self, capacity_bytes: int = (2 << 30)):
-    """Cache for a llama.cpp model."""
+        pass
+
+    @property
+    def cache_size(self):
+        return 0
+
+    def _find_longest_prefix_key(
+            self,
+            key: Tuple[int, ...],
+    ) -> Optional[Tuple[int, ...]]:
+        pass
+
+    def __getitem__(self, key: Sequence[int]) -> "LlamaState":
+        pass
+
+    def __contains__(self, key: Sequence[int]) -> bool:
+        pass
+
+    def __setitem__(self, key: Sequence[int], value: "LlamaState"):
+        pass
+
+
+class LlamaRAMCache(LlamaCache):
+    """Cache for a llama.cpp model using RAM."""
+
+    def __init__(self, capacity_bytes: int = (2 << 30)):
+        super().__init__(capacity_bytes)
+        self.capacity_bytes = capacity_bytes
+        self.cache_state: OrderedDict[Tuple[int, ...], "LlamaState"] = OrderedDict()
+
+    @property
+    def cache_size(self):
+        return sum([state.llama_state_size for state in self.cache_state.values()])
+
+    def _find_longest_prefix_key(
+            self,
+            key: Tuple[int, ...],
+    ) -> Optional[Tuple[int, ...]]:
+        min_len = 0
+        min_key = None
+        keys = (
+            (k, Llama.longest_token_prefix(k, key)) for k in self.cache_state.keys()
+        )
+        for k, prefix_len in keys:
+            if prefix_len > min_len:
+                min_len = prefix_len
+                min_key = k
+        return min_key
+
+    def __getitem__(self, key: Sequence[int]) -> "LlamaState":
+        key = tuple(key)
+        _key = self._find_longest_prefix_key(key)
+        if _key is None:
+            raise KeyError("Key not found")
+        value = self.cache_state[_key]
+        self.cache_state.move_to_end(_key)
+        return value
+
+    def __contains__(self, key: Sequence[int]) -> bool:
+        return self._find_longest_prefix_key(tuple(key)) is not None
+
+    def __setitem__(self, key: Sequence[int], value: "LlamaState"):
+        key = tuple(key)
+        if key in self.cache_state:
+            del self.cache_state[key]
+        self.cache_state[key] = value
+        while self.cache_size > self.capacity_bytes:
+            self.cache_state.popitem(last=False)
+
+
+class LlamaDiskCache(LlamaCache):
+    """Cache for a llama.cpp model using disk."""
 
     def __init__(self, cache_dir="./llama_cache", capacity_bytes: int = (2 << 30)):
+        super().__init__(capacity_bytes)
         self.cache = diskcache.Cache(cache_dir)
-        self.capacity_bytes = capacity_bytes
 
     @property
     def cache_size(self):
@@ -60,9 +134,6 @@ class LlamaCache:
         self.cache.push(_key)
         return value
 
-    def __contains__(self, key: Sequence[int]) -> bool:
-        return self._find_longest_prefix_key(tuple(key)) is not None
-
     def __setitem__(self, key: Sequence[int], value: "LlamaState"):
         key = tuple(key)
         if key in self.cache:
@@ -73,7 +144,6 @@ class LlamaCache:
             del self.cache[key_to_remove]
 
 
-
 class LlamaState:
     def __init__(
             self,
@@ -317,7 +387,7 @@ class Llama:
         assert self.ctx is not None
         n_ctx = self._n_ctx
         for i in range(0, len(tokens), self.n_batch):
-            batch = tokens[i : min(len(tokens), i + self.n_batch)]
+            batch = tokens[i: min(len(tokens), i + self.n_batch)]
             n_past = min(n_ctx - len(batch), len(self._input_ids))
             n_tokens = len(batch)
             return_code = llama_cpp.llama_eval(
@@ -339,7 +409,7 @@ class Llama:
             n_vocab = self._n_vocab
             cols = n_vocab
             logits_view = llama_cpp.llama_get_logits(self.ctx)
-            logits = [logits_view[i * cols : (i + 1) * cols] for i in range(rows)]
+            logits = [logits_view[i * cols: (i + 1) * cols] for i in range(rows)]
             self.eval_logits.extend(logits)
             self._scores: npt.NDArray[np.single] = np.concatenate(
                 (self._scores, np.array(logits, dtype=np.single)), axis=0
@@ -508,7 +578,7 @@ class Llama:
         assert self.ctx is not None
         last_n_tokens_data = [llama_cpp.llama_token(0)] * max(
             0, self.last_n_tokens_size - len(self._input_ids)
-        ) + self._input_ids[-self.last_n_tokens_size :].tolist()
+        ) + self._input_ids[-self.last_n_tokens_size:].tolist()
         return self._sample(
             last_n_tokens_data=(llama_cpp.llama_token * self.last_n_tokens_size)(
                 *last_n_tokens_data