diff --git a/tests/fsm/test_guide.py b/tests/fsm/test_guide.py index 510faf4b0..bf25c43c4 100644 --- a/tests/fsm/test_guide.py +++ b/tests/fsm/test_guide.py @@ -59,7 +59,7 @@ def convert_token_to_string(self, token): tokenizer = MockTokenizer() fsm = RegexGuide.from_regex(regex_str, tokenizer) - assert fsm.states_to_token_maps == {0: {1: 1}} + assert fsm.states_to_token_maps.get_transitions() == {0: {1: 1}} instruction = fsm.get_next_instruction(0) assert isinstance(instruction, Generate) @@ -70,9 +70,6 @@ def convert_token_to_string(self, token): assert fsm.is_final_state(0) is False - for state in fsm.final_states: - assert fsm.is_final_state(state) is True - def test_regex_multi_byte_llama_like(): class MockTokenizer: @@ -100,7 +97,7 @@ def convert_token_to_string(self, token): tokenizer = MockTokenizer() fsm = RegexGuide.from_regex(regex_str, tokenizer) - assert fsm.states_to_token_maps == { + assert fsm.states_to_token_maps.get_transitions() == { 0: {5: 1, 4: 2}, 1: {6: 3}, 3: {7: 4}, @@ -116,9 +113,6 @@ def convert_token_to_string(self, token): assert fsm.is_final_state(0) is False - for state in fsm.final_states: - assert fsm.is_final_state(state) is True - def test_regex_multi_byte_gpt2_like(): class MockTokenizer: @@ -147,7 +141,7 @@ def convert_token_to_string(self, token): tokenizer = MockTokenizer() fsm = RegexGuide.from_regex(regex_str, tokenizer) - assert fsm.states_to_token_maps == { + assert fsm.states_to_token_maps.get_transitions() == { 0: {5: 1, 10: 2}, 1: {8: 5, 4: 3}, 2: {11: 3}, @@ -163,9 +157,6 @@ def convert_token_to_string(self, token): assert fsm.is_final_state(0) is False - for state in fsm.final_states: - assert fsm.is_final_state(state) is True - def test_regex_final_state(): """Make sure that the FSM stays in the final state as we keep generating"""