From 16b636ef833eecf91dd29b8bbffd38dbf5e94f5d Mon Sep 17 00:00:00 2001
From: josc146 <josStorer@outlook.com>
Date: Wed, 31 Jan 2024 21:33:27 +0800
Subject: [PATCH] add EOS state cache point

---
 backend-python/utils/rwkv.py | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/backend-python/utils/rwkv.py b/backend-python/utils/rwkv.py
index df58448..6fa9438 100644
--- a/backend-python/utils/rwkv.py
+++ b/backend-python/utils/rwkv.py
@@ -272,6 +272,17 @@ class AbstractRWKV(ABC):
             )
 
             if token == self.EOS_ID:
+                try:
+                    state_cache.add_state(
+                        state_cache.AddStateBody(
+                            prompt=prompt + response,
+                            tokens=self.model_tokens,
+                            state=self.model_state,
+                            logits=logits,
+                        )
+                    )
+                except HTTPException:
+                    pass
                 yield response, "", prompt_token_len, completion_token_len
                 break