ensure logprobs is being passed as a boolean throughout the call chain

NexaAI · Sep 5, 2024 · 49615ec · 49615ec
1 parent 9ba5684
commit 49615ec
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/nexa/gguf/llama/llama.py b/nexa/gguf/llama/llama.py
@@ -732,7 +732,7 @@ def sample(
             apply_grammar=grammar is not None,
         )
 
-        if logprobs:
+        if logprobs is not None and (top_logprobs is not None and top_logprobs > 0):
             sampled_logprobs = self.logits_to_logprobs(logits)
             token_logprob = float(sampled_logprobs[id])
 

diff --git a/nexa/gguf/nexa_inference_text.py b/nexa/gguf/nexa_inference_text.py
@@ -332,11 +332,13 @@ def run_streamlit(self, model_path: str):
         help="Run the inference in Streamlit UI",
     )
     parser.add_argument(
+        "-lps",
         "--logprobs",
         action="store_true",
         help="Whether to return log probabilities of the output tokens",
     )
     parser.add_argument(
+        "-tlps",
         "--top_logprobs",
         type=int,
         default=None,