From f627485534cc8fcf68ddf72fddd9938efd29c6c8 Mon Sep 17 00:00:00 2001 From: kingbri Date: Sun, 11 Feb 2024 01:12:13 -0500 Subject: [PATCH] OAI: Fix completion token fetching The generator returns generated_tokens in the dict. Signed-off-by: kingbri --- OAI/utils/completion.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/OAI/utils/completion.py b/OAI/utils/completion.py index 81effe25..c768cfa7 100644 --- a/OAI/utils/completion.py +++ b/OAI/utils/completion.py @@ -43,7 +43,7 @@ def create_completion_response(generation: dict, model_name: Optional[str]): ) prompt_tokens = unwrap(generation.get("prompt_tokens"), 0) - completion_tokens = unwrap(generation.get("completion_tokens"), 0) + completion_tokens = unwrap(generation.get("generated_tokens"), 0) response = CompletionResponse( choices=[choice],