From e4bb709305945cd0b2890bc8fdf9e1c308d0dd26 Mon Sep 17 00:00:00 2001 From: kingbri Date: Sun, 12 May 2024 22:44:50 -0400 Subject: [PATCH] Model: Fix usage stats in non-streaming gens The wrong key was being returned from the model to the API. Signed-off-by: kingbri --- backends/exllamav2/model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/backends/exllamav2/model.py b/backends/exllamav2/model.py index 2246110..d42e2c0 100644 --- a/backends/exllamav2/model.py +++ b/backends/exllamav2/model.py @@ -656,7 +656,7 @@ class ExllamaV2Container: joined_generation["prompt_tokens"] = unwrap( generations[-1].get("prompt_tokens"), 0 ) - joined_generation["generation_tokens"] = unwrap( + joined_generation["generated_tokens"] = unwrap( generations[-1].get("generated_tokens"), 0 )