Merge pull request #77 from NexaAI/david/bugfix

fix text_inference bug
NexaAI · Sep 10, 2024 · bd9629f · bd9629f
2 parents 479e30c + df10d1c
commit bd9629f
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/nexa/gguf/nexa_inference_text.py b/nexa/gguf/nexa_inference_text.py
@@ -170,7 +170,12 @@ def run(self):
                             decoding_start_time = time.time()
                             prefill_time = decoding_start_time - generation_start_time
                             first_token = False
-                        delta = chunk["choices"][0]["text"]
+                        choice = chunk["choices"][0]
+                        if "text" in choice:
+                            delta = choice["text"]
+                        elif "delta" in choice:
+                            delta = choice["delta"]["content"]
+
                         print(delta, end="", flush=True)
                         generated_text += delta