webgpu(python) state cache

2023-12-28 20:43:57 +08:00
parent e33858f110
commit e083f2c629
3 changed files with 13 additions and 3 deletions
--- a/backend-python/utils/rwkv.py
+++ b/backend-python/utils/rwkv.py
@@ -239,7 +239,12 @@ class AbstractRWKV(ABC):
            self.model_tokens = []
        else:
            delta_prompt = prompt[len(cache["prompt"]) :]
-            self.model_state = copy.deepcopy(cache["state"])
+            state = cache["state"]
+            self.model_state = (
+                copy.deepcopy(state)
+                if type(state) == list or type(state) == np.ndarray
+                else state
+            )
            self.model_tokens = copy.deepcopy(cache["tokens"])
            logits = copy.deepcopy(cache["logits"])