Fix index interface in tests

dottxt-ai · Nov 22, 2024 · f099f96 · f099f96
1 parent 2eab80c
commit f099f96
Showing 1 changed file with 3 additions and 12 deletions.
diff --git 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"""