Skip to content

Commit

Permalink
Use TQDM to track index compilation progress
Browse files Browse the repository at this point in the history
  • Loading branch information
lapp0 authored and rlouf committed May 23, 2024
1 parent 6f655ca commit ffab2ac
Show file tree
Hide file tree
Showing 3 changed files with 14 additions and 1 deletion.
1 change: 1 addition & 0 deletions .pre-commit-config.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -30,3 +30,4 @@ repos:
- id: mypy
args: [--allow-redefinition]
exclude: ^examples/
additional_dependencies: [types-tqdm]
13 changes: 12 additions & 1 deletion outlines/fsm/regex.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
anything_else,
)
from numba.typed.typedobjectutils import _nonoptional
from tqdm import tqdm

if TYPE_CHECKING:
from outlines.models.tokenizer import Tokenizer
Expand Down Expand Up @@ -692,6 +693,12 @@ def create_fsm_index_end_to_end(
seen: Set[int] = set()
next_states = {fsm_info.initial}

pbar = tqdm(
total=len(set(fsm_info.transitions.values()))
+ 1, # all transitions plus initial
desc="Compiling FSM index for all state transitions",
)

while next_states:
start_state = next_states.pop()

Expand All @@ -713,7 +720,11 @@ def create_fsm_index_end_to_end(
if end_state not in seen:
next_states.add(end_state)

seen.add(start_state)
if start_state not in seen:
pbar.update(1)
seen.add(start_state)

pbar.close()

return states_to_token_subsets

Expand Down
1 change: 1 addition & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@ dependencies = [
"referencing",
"jsonschema",
"requests",
"tqdm"
]
dynamic = ["version"]

Expand Down

0 comments on commit ffab2ac

Please sign in to comment.