import torch
import torch.nn as nn


class KVcache:
    def __init__(self, n_layer):
        self.n_layer = n_layer
        self.key_list = [[] for _ in range(n_layer)]
        self.value_list = [[] for _ in range(n_layer)]

    def reset(self):
        """Clear the cache."""
        self.key_list.clear()
        self.value_list.clear()

    def update(self ,layer, keys, values):
        """Append new keys and values to the cache."""
        self.key_list[layer] = keys
        self.value_list[layer] = values

    def get_key_values(self,layer):
        """Retrieve all cached keys."""
        return self.key_list[layer], self.value_list[layer]