From 4d487eed8a36f5504842a527bd987c1ebad39333 Mon Sep 17 00:00:00 2001
From: kc611 <ckaustubhm06@gmail.com>
Date: Mon, 28 Oct 2024 13:06:12 +0530
Subject: [PATCH] Added EOS condition to next_state in Index struct

---
 src/index.rs | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/index.rs b/src/index.rs
index f74b3ce1..d299a0b4 100644
--- a/src/index.rs
+++ b/src/index.rs
@@ -113,6 +113,9 @@ impl Index {
     }
 
     pub(crate) fn next_state(&self, state: u32, token_id: u32) -> Option<u32> {
+        if token_id == self.eos_token_id {
+            return None;
+        }
         Some(*self.states_to_token_subsets.get(&state)?.get(&token_id)?)
     }