From 6f3a79899e183b4c3ebb1e3a2632e12e086bfd63 Mon Sep 17 00:00:00 2001 From: kc611 Date: Mon, 28 Oct 2024 10:40:41 +0530 Subject: [PATCH] Added eos token entries to final states --- src/index.rs | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/src/index.rs b/src/index.rs index d38d8fcc..f74b3ce1 100644 --- 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); }