From f160fbd6d1f75b2397dbffa1eb85cde90e2b5554 Mon Sep 17 00:00:00 2001
From: Martin Evans <martindevans@gmail.com>
Date: Fri, 12 Jan 2024 15:43:58 +0000
Subject: [PATCH] Added a check for EOS token in LLamaStatelessExecutor

---
 LLama/LLamaStatelessExecutor.cs | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/LLama/LLamaStatelessExecutor.cs b/LLama/LLamaStatelessExecutor.cs
index 77c9dbe4f..e03fe7a1c 100644
--- a/LLama/LLamaStatelessExecutor.cs
+++ b/LLama/LLamaStatelessExecutor.cs
@@ -106,6 +106,10 @@ public async IAsyncEnumerable<string> InferAsync(string prompt, IInferenceParams
                     );
                 }
 
+                // Check if this is the EOS token
+                if (id == _weights.EndOfSentenceToken)
+                    break;
+
                 // Decode this token into text
                 decoder.Add(id);
                 var decoded = decoder.Read();