Merge pull request SciSharp#920 from martindevans/fix_llama_get_logit…

…s_ith_null Fix `llama_get_logits_ith` Null Handling
SignalRT · Sep 4, 2024 · 0db26a4 · 0db26a4
2 parents 26a7114 + 846d2dc
commit 0db26a4
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 1 deletion.
diff --git a/LLama/Exceptions/RuntimeError.cs b/LLama/Exceptions/RuntimeError.cs
@@ -1,4 +1,4 @@
-using System;
+using System;
 using LLama.Native;
 
 namespace LLama.Exceptions;
@@ -56,4 +56,23 @@ public LLamaDecodeError(DecodeResult returnCode)
     {
         ReturnCode = returnCode;
     }
+}
+
+/// <summary>
+/// `llama_get_logits_ith` returned null, indicating that the index was invalid
+/// </summary>
+public class GetLogitsInvalidIndexException
+    : RuntimeError
+{
+    /// <summary>
+    /// The incorrect index passed to the `llama_get_logits_ith` call
+    /// </summary>
+    public int Index { get; }
+
+    /// <inheritdoc />
+    public GetLogitsInvalidIndexException(int index)
+        : base($"llama_get_logits_ith({index}) returned null")
+    {
+        Index = index;
+    }
 }
diff --git a/LLama/Native/SafeLLamaContextHandle.cs b/LLama/Native/SafeLLamaContextHandle.cs
@@ -472,6 +472,9 @@ public Span<float> GetLogitsIth(int i)
             unsafe
             {
                 var logits = llama_get_logits_ith(this, i);
+                if (logits == null)
+                    throw new GetLogitsInvalidIndexException(i);
+
                 return new Span<float>(logits, model.VocabCount);
             }
         }