diff --git a/python/outlines_core/fsm/guide.py b/python/outlines_core/fsm/guide.py index 26deb00e..ce92bfa3 100644 --- a/python/outlines_core/fsm/guide.py +++ b/python/outlines_core/fsm/guide.py @@ -9,7 +9,7 @@ make_deterministic_fsm, ) -from .outlines_core_rs import PyVocabIndex +from .outlines_core_rs import Index @dataclass(frozen=True) @@ -109,7 +109,7 @@ def create_states_mapping( tokenizer, regex_parser: Callable[[str], interegular.Pattern] = interegular.parse_pattern, frozen_tokens: List[str] = [], -) -> Tuple[PyVocabIndex, Set[int], Set[int]]: +) -> Tuple[Index, Set[int], Set[int]]: """Create the variables related to the mapping between states and tokens from a regex string. The parameters of the function are used for caching purpose. @@ -145,7 +145,7 @@ def create_states_mapping_from_fsm( fsm: interegular.fsm.FSM, tokenizer, frozen_tokens: List[str] = [], -) -> Tuple[PyVocabIndex, Set[int], Set[int]]: +) -> Tuple[Index, Set[int], Set[int]]: """Create the variables related to the mapping between states and tokens from an FSM. The parameters of the function are used for caching purpose. diff --git a/python/outlines_core/fsm/outlines_core_rs.pyi b/python/outlines_core/fsm/outlines_core_rs.pyi index 3dfe41e6..2cfedb68 100644 --- a/python/outlines_core/fsm/outlines_core_rs.pyi +++ b/python/outlines_core/fsm/outlines_core_rs.pyi @@ -87,7 +87,7 @@ class Vocabulary: """ ... -class PyVocabIndex: +class Index: def get_allowed_tokens(self, state: int): """ Return the next instruction for guided generation. diff --git a/python/outlines_core/fsm/regex.py b/python/outlines_core/fsm/regex.py index 3ab8af3d..af337e34 100644 --- a/python/outlines_core/fsm/regex.py +++ b/python/outlines_core/fsm/regex.py @@ -24,7 +24,7 @@ from .outlines_core_rs import ( # noqa: F401 FSMInfo, - PyVocabIndex, + Index, Vocabulary, _walk_fsm, create_fsm_index_end_to_end, @@ -439,7 +439,7 @@ def create_fsm_index_tokenizer( fsm: BetterFSM, tokenizer, frozen_tokens: Optional[Iterable[str]] = None, -) -> Tuple[PyVocabIndex, Set[int]]: +) -> Tuple[Index, Set[int]]: """Construct an FMS index from a tokenizer. This uses the end-to-end approach of `create_fsm_index_end_to_end`. @@ -470,7 +470,7 @@ def create_fsm_index_tokenizer( """ tokens_to_token_ids, empty_token_ids = reduced_vocabulary(tokenizer) - states_to_token_subsets = PyVocabIndex( # type: ignore + states_to_token_subsets = Index( # type: ignore fsm.fsm_info, Vocabulary.from_dict(tokens_to_token_ids), tokenizer.eos_token_id, diff --git a/src/python_bindings/mod.rs b/src/python_bindings/mod.rs index f50f51dd..01ccc35b 100644 --- a/src/python_bindings/mod.rs +++ b/src/python_bindings/mod.rs @@ -46,7 +46,7 @@ impl FSMInfo { } #[pyclass] -pub struct PyVocabIndex { +pub struct Index { initial: u32, finals: HashSet, states_to_token_subsets: HashMap>, @@ -55,7 +55,7 @@ pub struct PyVocabIndex { } #[pymethods] -impl PyVocabIndex { +impl Index { #[new] fn new( fsm_info: &FSMInfo, @@ -332,7 +332,7 @@ fn outlines_core_rs(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(to_regex_py, m)?)?; m.add_class::()?; - m.add_class::()?; + m.add_class::()?; Ok(()) }