From 6b9af58cc12aa2be8638d47470d073f20edb6625 Mon Sep 17 00:00:00 2001 From: kingbri Date: Sat, 18 Nov 2023 22:34:40 -0500 Subject: [PATCH] Tree: Fix extraneous bugs and update T/s print Model: Add extra information to print and fix the divide by zero error. Auth: Fix validation of API and admin keys to look for the entire key. References #7 and #6 Signed-off-by: kingbri --- auth.py | 4 ++-- model.py | 11 ++++++++++- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/auth.py b/auth.py index 39c8250..d74b8f2 100644 --- a/auth.py +++ b/auth.py @@ -48,7 +48,7 @@ def check_api_key(x_api_key: str = Header(None), authorization: str = Header(Non if len(split_key) < 2: raise HTTPException(401, "Invalid API key") - elif split_key[0].lower() == "bearer" and split_key[1] in auth_keys.api_key: + elif split_key[0].lower() == "bearer" and split_key[1] == auth_keys.api_key: return authorization else: raise HTTPException(401, "Invalid API key") @@ -66,7 +66,7 @@ def check_admin_key(x_admin_key: str = Header(None), authorization: str = Header if len(split_key) < 2: raise HTTPException(401, "Invalid admin key") - elif split_key[0].lower() == "bearer" and split_key[1] in auth_keys.admin_key: + elif split_key[0].lower() == "bearer" and split_key[1] == auth_keys.admin_key: return authorization else: raise HTTPException(401, "Invalid admin key") diff --git a/model.py b/model.py index a10b69f..7274fa0 100644 --- a/model.py +++ b/model.py @@ -367,4 +367,13 @@ class ModelContainer: if eos or generated_tokens == max_tokens: break elapsed_time = last_chunk_time - start_time - print(f"Response: {round(generated_tokens, 2)} tokens generated in {round(elapsed_time, 2)} seconds ({round(generated_tokens / elapsed_time, 2)} T/s)") + + initial_response = f"Response: {round(generated_tokens, 2)} tokens generated in {round(elapsed_time, 2)} seconds" + extra_responses = [] + + # Add tokens per second + extra_responses.append(f"{'Indeterminate' if elapsed_time == 0 else round(generated_tokens / elapsed_time, 2)} T/s") + extra_responses.append(f"{generated_tokens} tokens") + + # Print output + print(initial_response + " (" + ", ".join(extra_responses) + ")")