From f22e2fb5509fb639cdecd6ca98345789d8fe35eb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Adrien=20Gallou=C3=ABt?= <angt@huggingface.co>
Date: Wed, 5 Feb 2025 16:12:34 +0000
Subject: [PATCH] Cleanup
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Adrien Gallouët <angt@huggingface.co>
---
 backends/llamacpp/src/backend.rs | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/backends/llamacpp/src/backend.rs b/backends/llamacpp/src/backend.rs
index 1072dec4..bed7d2bd 100644
--- a/backends/llamacpp/src/backend.rs
+++ b/backends/llamacpp/src/backend.rs
@@ -290,6 +290,12 @@ impl Llamacpp {
         Ok(Llamacpp{model, ctx, vocab, logprobs, batch})
     }
 
+    fn decode(&mut self) -> i32 {
+        unsafe {
+            llamacpp::decode(self.ctx, self.batch)
+        }
+    }
+
     fn clear_kv_cache(&mut self, seq_id: llamacpp::llama_seq_id) {
         unsafe {
             llamacpp::kv_cache_seq_rm(self.ctx, seq_id, -1, -1);
@@ -543,14 +549,8 @@ impl LlamacppBackend {
                         running: true,
                     });
                 }
-                loop {
-                    if llamacpp.batch.n_tokens == 0 {
-                        break;
-                    }
-                    let decode = unsafe {
-                        llamacpp::decode(llamacpp.ctx, llamacpp.batch)
-                    };
-                    if decode != 0 {
+                while llamacpp.batch.n_tokens > 0 {
+                    if llamacpp.decode() != 0 {
                         warn!("llama_decode failed, clearing kv cache");
                         llamacpp.clear_kv_cache(-1);
                         for seq in seqs.iter_mut() {