Model: Skip empty token chunks

This helps make the generation loop more efficient by skipping past chunks that aren't providing any tokens anyways. The offset isn't affected. Signed-off-by: kingbri <[email protected]>
theroyallab · Jul 22, 2024 · 21516bd · 21516bd
1 parent 0eedc8c
commit 21516bd
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/backends/exllamav2/model.py b/backends/exllamav2/model.py
@@ -1185,13 +1185,15 @@ async def generate_gen(
                 result_id = result.get("identifier")
 
                 if stage == "streaming" and result_id == job_id:
-                    chunk = unwrap(result.get("text"), "")
-                    full_response += chunk
-
                     chunk_tokens = result.get("token_ids")
-                    if chunk_tokens is not None:
+                    if chunk_tokens is None:
+                        continue
+                    else:
                         generated_tokens += chunk_tokens.size(dim=0)
 
+                    chunk = unwrap(result.get("text"), "")
+                    full_response += chunk
+
                     generation = {
                         "text": chunk,
                         "prompt_tokens": context_len,