FilesExpand file tree

 batch-processing

/

llama_cache.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

150 lines (121 loc) · 4.85 KB

 batch-processing

/

llama_cache.py

Copy path

Top

File metadata and controls

• Code
• Blame

150 lines (121 loc) · 4.85 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

86

87

88

89

90

91

92

93

94

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

import sys

from abc import ABC, abstractmethod

from typing import (

Optional,

Sequence,

Tuple,

)

from collections import OrderedDict

import diskcache

import llama_cpp.llama

from .llama_types import *

class BaseLlamaCache(ABC):

"""Base cache class for a llama.cpp model."""

def __init__(self, capacity_bytes: int = (2 << 30)):

self.capacity_bytes = capacity_bytes

@property

@abstractmethod

def cache_size(self) -> int:

raise NotImplementedError

def _find_longest_prefix_key(

self,

key: Tuple[int, ...],

) -> Optional[Tuple[int, ...]]:

pass

@abstractmethod

def __getitem__(self, key: Sequence[int]) -> "llama_cpp.llama.LlamaState":

raise NotImplementedError

@abstractmethod

def __contains__(self, key: Sequence[int]) -> bool:

raise NotImplementedError

@abstractmethod

def __setitem__(self, key: Sequence[int], value: "llama_cpp.llama.LlamaState") -> None:

raise NotImplementedError

class LlamaRAMCache(BaseLlamaCache):

"""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, ...], "llama_cpp.llama.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_cpp.llama.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]) -> "llama_cpp.llama.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: "llama_cpp.llama.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 and len(self.cache_state) > 0:

self.cache_state.popitem(last=False)

# Alias for backwards compatibility

LlamaCache = LlamaRAMCache

class LlamaDiskCache(BaseLlamaCache):

"""Cache for a llama.cpp model using disk."""

def __init__(

self, cache_dir: str = ".cache/llama_cache", capacity_bytes: int = (2 << 30)

):

super().__init__(capacity_bytes)

self.cache = diskcache.Cache(cache_dir)

@property

def cache_size(self):

return int(self.cache.volume()) # type: ignore

def _find_longest_prefix_key(

self,

key: Tuple[int, ...],

) -> Optional[Tuple[int, ...]]:

min_len = 0

min_key: Optional[Tuple[int, ...]] = None

for k in self.cache.iterkeys(): # type: ignore

prefix_len = llama_cpp.llama.Llama.longest_token_prefix(k, key)

if prefix_len > min_len:

min_len = prefix_len

min_key = k # type: ignore

return min_key

def __getitem__(self, key: Sequence[int]) -> "llama_cpp.llama.LlamaState":

key = tuple(key)

_key = self._find_longest_prefix_key(key)

if _key is None:

raise KeyError("Key not found")

value: "llama_cpp.llama.LlamaState" = self.cache.pop(_key) # type: ignore

# NOTE: This puts an integer as key in cache, which breaks,

# Llama.longest_token_prefix(k, key) above since k is not a tuple of ints/tokens

# self.cache.push(_key, side="front") # type: ignore

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: "llama_cpp.llama.LlamaState"):

print("LlamaDiskCache.__setitem__: called", file=sys.stderr)

key = tuple(key)

if key in self.cache:

print("LlamaDiskCache.__setitem__: delete", file=sys.stderr)

del self.cache[key]

self.cache[key] = value

print("LlamaDiskCache.__setitem__: set", file=sys.stderr)

while self.cache_size > self.capacity_bytes and len(self.cache) > 0:

key_to_remove = next(iter(self.cache))

del self.cache[key_to_remove]

print("LlamaDiskCache.__setitem__: trim", file=sys.stderr)