add EOS state cache point

This commit is contained in:
josc146 2024-01-31 21:33:27 +08:00
parent 4339ce20d5
commit 16b636ef83

View File

@ -272,6 +272,17 @@ class AbstractRWKV(ABC):
) )
if token == self.EOS_ID: 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 yield response, "", prompt_token_len, completion_token_len
break break