get_index_dict -> get_transitions

dottxt-ai · Oct 31, 2024 · 86b7d86 · 86b7d86
1 parent 9c0868f
commit 86b7d86
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 6 deletions.
diff --git a/python/outlines_core/fsm/guide.py b/python/outlines_core/fsm/guide.py
@@ -306,4 +306,4 @@ def copy(self):
 
     def get_index_dict(self):
         """Returns the Index as a Python Dict object."""
-        return self.states_to_token_maps.get_index_dict()
+        return self.states_to_token_maps.get_transitions()
diff --git a/src/index.rs b/src/index.rs
@@ -35,7 +35,7 @@ impl FSMInfo {
 
 #[derive(Debug)]
 pub struct Index {
-    pub(crate) initial: u32,
+    initial: u32,
     finals: HashSet<u32>,
     states_to_token_subsets: HashMap<u32, HashMap<u32, u32>>,
     eos_token_id: u32,
@@ -126,7 +126,7 @@ impl Index {
         self.finals.contains(&state)
     }
 
-    pub(crate) fn index(&self) -> &HashMap<u32, HashMap<u32, u32>> {
+    pub(crate) fn transitions(&self) -> &HashMap<u32, HashMap<u32, u32>> {
         &self.states_to_token_subsets
     }
 }
diff --git a/src/python_bindings/mod.rs b/src/python_bindings/mod.rs
@@ -101,8 +101,8 @@ impl PyIndex {
         self.0.is_final(state)
     }
 
-    fn get_index_dict(&self) -> HashMap<u32, HashMap<u32, u32>> {
-        self.0.index().clone()
+    fn get_transitions(&self) -> HashMap<u32, HashMap<u32, u32>> {
+        self.0.transitions().clone()
     }
 
     fn get_initial_state(&self) -> u32 {

diff --git a/tests/fsm/test_regex.py b/tests/fsm/test_regex.py
@@ -372,7 +372,7 @@ def test_create_fsm_index_tokenizer(hf_tokenizer_uri, revision):
     )
 
     assert not empty_token_ids
-    assert len(states_to_token_subsets.get_index_dict()) / num_fsm_states > 0.94
+    assert len(states_to_token_subsets.get_transitions()) / num_fsm_states > 0.94
 
 
 @pytest.mark.parametrize(