Spaces:
Running
Running
| import time | |
| from typing import Any, Dict, List, Optional, Tuple | |
| def clone_past_key_values(past_key_values: Any) -> Any: | |
| if past_key_values is None: | |
| return None | |
| from transformers.cache_utils import DynamicCache | |
| if isinstance(past_key_values, DynamicCache): | |
| cloned = DynamicCache() | |
| cloned.key_cache = [k.clone() for k in past_key_values.key_cache] | |
| cloned.value_cache = [v.clone() for v in past_key_values.value_cache] | |
| return cloned | |
| return tuple(tuple(t.clone() for t in layer) for layer in past_key_values) | |
| class RadixNode: | |
| def __init__(self, token_id: Optional[int] = None): | |
| self.token_id = token_id | |
| self.children: Dict[int, "RadixNode"] = {} | |
| self.past_key_values: Optional[Any] = None | |
| self.last_accessed: float = time.time() | |
| class RadixCacheManager: | |
| def __init__(self, max_nodes: int = 16, min_prefix_len: int = 10): | |
| self.root = RadixNode() | |
| self.max_nodes = max_nodes | |
| self.min_prefix_len = min_prefix_len | |
| self.node_list: List[RadixNode] = [] | |
| def query(self, token_ids: List[int]) -> Tuple[Optional[Any], int]: | |
| current = self.root | |
| last_pkv_node = None | |
| match_len = 0 | |
| current_depth = 0 | |
| for token_id in token_ids: | |
| if token_id in current.children: | |
| current = current.children[token_id] | |
| current_depth += 1 | |
| if current.past_key_values is not None: | |
| last_pkv_node = current | |
| match_len = current_depth | |
| else: | |
| break | |
| if last_pkv_node and match_len >= self.min_prefix_len: | |
| last_pkv_node.last_accessed = time.time() | |
| return clone_past_key_values(last_pkv_node.past_key_values), match_len | |
| return None, 0 | |
| def insert(self, token_ids: List[int], past_key_values: Any): | |
| if len(token_ids) < self.min_prefix_len or past_key_values is None: | |
| return | |
| current = self.root | |
| for token_id in token_ids: | |
| if token_id not in current.children: | |
| node = RadixNode(token_id) | |
| current.children[token_id] = node | |
| self.node_list.append(node) | |
| current = current.children[token_id] | |
| current.past_key_values = past_key_values | |
| current.last_accessed = time.time() | |
| self._evict_if_needed() | |
| def _evict_if_needed(self): | |
| nodes_with_cache = [n for n in self.node_list if n.past_key_values is not None] | |
| if len(nodes_with_cache) > self.max_nodes: | |
| nodes_with_cache.sort(key=lambda n: n.last_accessed) | |
| evicted = nodes_with_cache[0] | |
| evicted.past_key_values = None | |