Added EOS condition to next_state in Index struct

dottxt-ai · Oct 28, 2024 · 4d487ee · 4d487ee
1 parent 4b13839
commit 4d487ee
Showing 1 changed file with 3 additions and 0 deletions.
diff --git 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)?)
     }