mirror of
https://github.com/rasbt/LLMs-from-scratch.git
synced 2025-08-13 19:21:24 +00:00
22 lines
646 B
Python
22 lines
646 B
Python
|
# Copyright (c) Sebastian Raschka under Apache License 2.0 (see LICENSE.txt).
|
||
|
|
# Source for "Build a Large Language Model From Scratch"
|
||
|
|
# - https://www.manning.com/books/build-a-large-language-model-from-scratch
|
||
|
|
# Code: https://github.com/rasbt/LLMs-from-scratch
|
||
|
|
|
||
|
|
class KVCache:
|
||
|
|
def __init__(self, n_layers):
|
||
|
|
self.cache = [None] * n_layers
|
||
|
|
|
||
|
|
def get(self, layer_idx):
|
||
|
|
return self.cache[layer_idx]
|
||
|
|
|
||
|
|
def update(self, layer_idx, value):
|
||
|
|
self.cache[layer_idx] = value
|
||
|
|
|
||
|
|
def get_all(self):
|
||
|
|
return self.cache
|
||
|
|
|
||
|
|
def reset(self):
|
||
|
|
for i in range(len(self.cache)):
|
||
|
|
self.cache[i] = None
|