From df72c56b5b57adf51fb32a342406c45fda144947 Mon Sep 17 00:00:00 2001
From: Morgan Funtowicz <funtowiczmo@gmail.com>
Date: Thu, 28 Nov 2024 16:30:20 +0100
Subject: [PATCH] feat(backend): add guard in case top_k = 0

---
 backends/llamacpp/csrc/backend.cpp | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/backends/llamacpp/csrc/backend.cpp b/backends/llamacpp/csrc/backend.cpp
index 17709b72..46052435 100644
--- a/backends/llamacpp/csrc/backend.cpp
+++ b/backends/llamacpp/csrc/backend.cpp
@@ -28,7 +28,10 @@ namespace huggingface::tgi::backends::llamacpp {
                 false,
                 false
         ));
-        llama_sampler_chain_add(pSampler, llama_sampler_init_top_k(static_cast<int32_t>(top_k)));
+
+        if (top_k > 0) {
+            llama_sampler_chain_add(pSampler, llama_sampler_init_top_k(static_cast<int32_t>(top_k)));
+        }
 
         if (0 < top_p && top_p < 1) {
             llama_sampler_chain_add(pSampler, llama_sampler_init_top_p(top_p, 1));