Added eos token entries to final states

dottxt-ai · Oct 28, 2024 · 6f3a798 · 6f3a798
1 parent 9281674
commit 6f3a798
Showing 1 changed file with 9 additions and 4 deletions.
diff --git a/src/index.rs b/src/index.rs
@@ -72,15 +72,20 @@ impl Index {
                 start_state,
             );
 
-            for (token_id, end_state) in token_ids_end_states {
+            for (token_id, end_state) in &token_ids_end_states {
                 let inner_map = states_to_token_subsets.entry(start_state).or_default();
-                inner_map.insert(token_id, end_state);
+                inner_map.insert(*token_id, *end_state);
 
-                if !seen.contains(&end_state) {
-                    next_states.insert(end_state);
+                if !seen.contains(end_state) {
+                    next_states.insert(*end_state);
                 }
             }
 
+            if fsm_info.finals.contains(&start_state) && !token_ids_end_states.is_empty() {
+                let inner_map = states_to_token_subsets.entry(start_state).or_default();
+                inner_map.insert(eos_token_id, start_state);
+            }
+
             seen.insert(start_state);
         }