fixup some bug when run api.py

Tlntin · Oct 22, 2024 · d56c319 · d56c319
1 parent 8a27cbb
commit d56c319
Showing 1 changed file with 4 additions and 6 deletions.
diff --git a/utils/inference.py b/utils/inference.py
@@ -263,14 +263,13 @@ def predict(
                 if show_progress:
                     prefill_show_progress = True
                 # reset counter
-                self.session.run_times = 0
-                self.session.kv_cache.real_kv_size = 0
+                self.session.reset()
             else:
                 prefill_show_progress = False
             logits = self.session.run(
                 input_ids,
                 show_progress=prefill_show_progress
-            )[0]
+            )
             input_ids = self.sample_logits(
                 logits[0][-1:],
                 self.sampling_method,
@@ -320,14 +319,13 @@ def generate(
                 if show_progress:
                     prefill_show_progress = True
                 # reset counter
-                self.session.run_times = 0
-                self.session.kv_cache.real_kv_size = 0
+                self.session.reset()
             else:
                 prefill_show_progress = False
             logits = self.session.run(
                 input_ids,
                 show_progress=prefill_show_progress
-            )[0]
+            )
             input_ids = self.sample_logits(
                 logits[0][-1:],
                 self.sampling_method,