From b24a18cd3a88a926735936de57f00a3177ca2756 Mon Sep 17 00:00:00 2001 From: josc146 Date: Thu, 16 May 2024 13:48:06 +0800 Subject: [PATCH] fix a tps error --- backend-python/routes/completion.py | 5 ++++- backend-python/utils/rwkv.py | 5 ++++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/backend-python/routes/completion.py b/backend-python/routes/completion.py index 93254c0..3ffad8b 100644 --- a/backend-python/routes/completion.py +++ b/backend-python/routes/completion.py @@ -191,7 +191,10 @@ async def eval_rwkv( # torch_gc() requests_num = requests_num - 1 completion_end_time = time.time() - tps = completion_tokens / (completion_end_time - completion_start_time) + completion_interval = completion_end_time - completion_start_time + tps = 0 + if completion_interval > 0: + tps = completion_tokens / completion_interval print(f"Generation TPS: {tps:.2f}") if await request.is_disconnected(): diff --git a/backend-python/utils/rwkv.py b/backend-python/utils/rwkv.py index 1016c94..76438fc 100644 --- a/backend-python/utils/rwkv.py +++ b/backend-python/utils/rwkv.py @@ -257,7 +257,10 @@ class AbstractRWKV(ABC): self.fix_tokens(self.pipeline.encode(delta_prompt)) ) prompt_end_time = time.time() - tps = prompt_token_len / (prompt_end_time - prompt_start_time) + prompt_interval = prompt_end_time - prompt_start_time + tps = 0 + if prompt_interval > 0: + tps = prompt_token_len / prompt_interval print(f"Prompt Prefill TPS: {tps:.2f}", end=" ", flush=True) try: state_cache.add_state(