Move index to fsm

dottxt-ai · Dec 1, 2023 · b57e149 · b57e149
1 parent a774693
commit b57e149
Show file tree

Hide file tree

Showing 17 changed files with 18 additions and 21 deletions.
diff --git a/examples/parsing.py b/examples/parsing.py
@@ -14,7 +14,7 @@
     set_seed,
 )
 
-from outlines.index.parsing import PartialLark, PartialPythonIndenter
+from outlines.fsm.parsing import PartialLark, PartialPythonIndenter
 
 revision = None
 checkpoint = "Salesforce/codegen-350M-mono"

diff --git a/outlines/index/__init__.py → outlines/fsm/__init__.py b/outlines/index/__init__.py → outlines/fsm/__init__.py
diff --git a/outlines/index/index.py → outlines/fsm/fsm.py b/outlines/index/index.py → outlines/fsm/fsm.py
@@ -2,7 +2,7 @@
 
 import interegular
 
-from outlines.index.fsm import create_fsm_index_tokenizer, make_deterministic_fsm
+from outlines.fsm.regex import create_fsm_index_tokenizer, make_deterministic_fsm
 
 if TYPE_CHECKING:
     from outlines.models.tokenizer import Tokenizer

diff --git a/outlines/index/json_schema.py → outlines/fsm/json_schema.py b/outlines/index/json_schema.py → outlines/fsm/json_schema.py
diff --git a/outlines/index/parsing.py → outlines/fsm/parsing.py b/outlines/index/parsing.py → outlines/fsm/parsing.py
@@ -35,7 +35,7 @@
 from lark.parsers.lalr_interactive_parser import InteractiveParser
 from lark.parsers.lalr_parser import LALR_Parser, ParseConf, ParserState, _Parser
 
-from outlines.index.fsm import (
+from outlines.fsm.regex import (
     fsm_union,
     get_sub_fsms_from_seq,
     make_deterministic_fsm,

diff --git a/outlines/index/fsm.py → outlines/fsm/regex.py b/outlines/index/fsm.py → outlines/fsm/regex.py
diff --git a/outlines/index/types.py → outlines/fsm/types.py b/outlines/index/types.py → outlines/fsm/types.py
diff --git a/outlines/generate/api.py b/outlines/generate/api.py
@@ -3,14 +3,11 @@
 
 from pydantic import BaseModel
 
+from outlines.fsm.fsm import RegexFSM, StopAtTokenFSM
+from outlines.fsm.json_schema import build_regex_from_object, get_schema_from_signature
+from outlines.fsm.types import python_types_to_regex
 from outlines.generate.generator import SequenceGenerator
 from outlines.generate.samplers import Sampler, multinomial
-from outlines.index.index import RegexFSM, StopAtTokenFSM
-from outlines.index.json_schema import (
-    build_regex_from_object,
-    get_schema_from_signature,
-)
-from outlines.index.types import python_types_to_regex
 
 
 def text(model, max_tokens: Optional[int] = None, *, sampler: Sampler = multinomial):

diff --git a/outlines/generate/generator.py b/outlines/generate/generator.py
@@ -4,11 +4,11 @@
 
 import torch
 
-from outlines.index.index import FSMState
+from outlines.fsm.fsm import FSMState
 
 if TYPE_CHECKING:
+    from outlines.fsm.fsm import FSM
     from outlines.generate.samplers import Sampler
-    from outlines.index.index import FSM
 
 
 @dataclasses.dataclass(frozen=True)

diff --git a/tests/index/partial_python.lark → tests/fsm/partial_python.lark b/tests/index/partial_python.lark → tests/fsm/partial_python.lark
diff --git a/tests/index/test_index.py → tests/fsm/test_fsm.py b/tests/index/test_index.py → tests/fsm/test_fsm.py
@@ -1,6 +1,6 @@
 import pytest
 
-from outlines.index.index import RegexFSM, StopAtTokenFSM
+from outlines.fsm.fsm import RegexFSM, StopAtTokenFSM
 
 
 def test_stop_at_token():

diff --git a/tests/index/test_json_schema.py → tests/fsm/test_json_schema.py b/tests/index/test_json_schema.py → tests/fsm/test_json_schema.py
@@ -5,7 +5,7 @@
 import pytest
 from pydantic import BaseModel, constr
 
-from outlines.index.json_schema import (
+from outlines.fsm.json_schema import (
     BOOLEAN,
     INTEGER,
     NULL,

diff --git a/tests/index/test_parsing.py → tests/fsm/test_parsing.py b/tests/index/test_parsing.py → tests/fsm/test_parsing.py
@@ -4,14 +4,14 @@
 from lark.indenter import DedentError
 from lark.lexer import UnexpectedCharacters, UnexpectedToken
 
-from outlines.index.parsing import PartialLark, PartialPythonIndenter
+from outlines.fsm.parsing import PartialLark, PartialPythonIndenter
 
 
 def test_partial_parsing():
     lp = PartialLark.open_from_package(
         "tests",
         "partial_python.lark",
-        ["index"],
+        ["fsm"],
         parser="lalr",
         postlex=PartialPythonIndenter(),
         start="file_input",
@@ -123,7 +123,7 @@ def test_partial_parsing():
     lp = PartialLark.open_from_package(
         "tests",
         "partial_python.lark",
-        ["index"],
+        ["fsm"],
         parser="lalr",
         postlex=PartialPythonIndenter(),
         start="file_input",
@@ -160,7 +160,7 @@ def test_sequential_parse_example():
     lp = PartialLark.open_from_package(
         "tests",
         "partial_python.lark",
-        ["index"],
+        ["fsm"],
         parser="lalr",
         postlex=PartialPythonIndenter(),
         start="file_input",

diff --git a/tests/index/test_fsm.py → tests/fsm/test_regex.py b/tests/index/test_fsm.py → tests/fsm/test_regex.py
@@ -2,7 +2,7 @@
 import numba
 import pytest
 
-from outlines.index.fsm import (
+from outlines.fsm.regex import (
     _walk_fsm,
     create_fsm_index,
     create_fsm_index_end_to_end,

diff --git a/tests/index/test_types.py → tests/fsm/test_types.py b/tests/index/test_types.py → tests/fsm/test_types.py
@@ -2,7 +2,7 @@
 
 import pytest
 
-from outlines.index.types import (
+from outlines.fsm.types import (
     BOOLEAN,
     DATE,
     DATETIME,

diff --git a/tests/generate/test_generator.py b/tests/generate/test_generator.py
@@ -4,6 +4,7 @@
 import pytest
 import torch
 
+from outlines.fsm.fsm import FSMState
 from outlines.generate.generator import (
     SequenceGenerator,
     bias_logits,
@@ -15,7 +16,6 @@
     token_generator,
     update_token_ids,
 )
-from outlines.index.index import FSMState
 
 
 def test_sequence_generator_class():

diff --git a/tests/generate/test_integration_transfomers.py b/tests/generate/test_integration_transfomers.py
@@ -9,7 +9,7 @@
 
 import outlines.generate as generate
 import outlines.models as models
-from outlines.index.fsm import reduced_vocabulary
+from outlines.fsm.regex import reduced_vocabulary
 from outlines.models.transformers import TransformerTokenizer