Skip to content

Commit

Permalink
Remove profiling test remaining from Outlines
Browse files Browse the repository at this point in the history
  • Loading branch information
rlouf committed Oct 9, 2024
1 parent fbb7642 commit c3a5ecd
Showing 1 changed file with 0 additions and 52 deletions.
52 changes: 0 additions & 52 deletions tests/fsm/test_regex.py
Original file line number Diff line number Diff line change
Expand Up @@ -460,58 +460,6 @@ def test_regex_index_performance():
profiler.print_stats(output_unit=1e-3, summarize=True, stripzeros=True)


@pytest.mark.skip(reason="Only for local profiling")
def test_json_index_performance():
import json
from enum import Enum

from line_profiler import LineProfiler # type: ignore [import]
from pydantic import BaseModel, constr

import outlines_core

class Weapon(str, Enum):
sword = "sword"
axe = "axe"
mace = "mace"
spear = "spear"
bow = "bow"
crossbow = "crossbow"

class Armor(str, Enum):
leather = "leather"
chainmail = "chainmail"
plate = "plate"

class Character(BaseModel):
name: constr(max_length=10)
# TODO: Add support for conint
age: int # conint(int, ge=18, le=100)
armor: Armor
weapon: Weapon
# TODO: Add support for conint
strength: int # conint(int, ge=0, le=100)

model = outlines_core.models.transformers("gpt2", device="cuda")
json_schema = json.dumps(Character.model_json_schema())

def build_regex():
regex_str = outlines_core.index.json_schema.build_regex_from_object(json_schema)
outlines_core.generate.regex(model, regex_str)

profiler = LineProfiler(create_fsm_index_end_to_end)
profiler.add_function(create_fsm_index_tokenizer)
profiler.add_function(outlines_core.index.index.RegexFSM.__init__)

profiler.runctx(
"build_regex()",
globals(),
locals(),
)
profiler.dump_stats("line-profiler-build-json-regex.pkl")
profiler.print_stats(output_unit=1e-3, summarize=True, stripzeros=True)


def test_token_trans_keys_identical():
"""assert two tokens w/ identical behavior wrt FSM have same trans key seq"""

Expand Down

0 comments on commit c3a5ecd

Please sign in to comment.