add EOS state cache point
This commit is contained in:
parent
4339ce20d5
commit
16b636ef83
@ -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
|
||||||
|
|
||||||
|
Loading…
Reference in New Issue
Block a user