jalr/tabbyAPI-ollama
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Tree: Fix extraneous bugs and update T/s print

Browse source 
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 <bdashore3@proton.me>
This commit is contained in:
kingbri 2023-11-18 22:34:40 -05:00
parent a51889bdb8
commit 6b9af58cc1
2 changed files with 12 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

4
auth.py
View file

@ -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")

11
model.py
View file

@ -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) + ")")