Skip to content

Commit

Permalink
fix style
Browse files Browse the repository at this point in the history
  • Loading branch information
Eitan Turok committed May 30, 2024
1 parent a012690 commit 944df28
Show file tree
Hide file tree
Showing 3 changed files with 162 additions and 54 deletions.
72 changes: 54 additions & 18 deletions tests/generate/test_integration_llamacpp.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -248,14 +248,14 @@ def test_llamacpp_json_dict(model):
prompt = "<|im_start|>user\nOutput some JSON<|im_end|>\n<|im_start|>assistant\n"

schema_dict = {
"title": "spam",
"type": "object",
"properties": {
"foo" : {"type": "boolean"},
"bar": {"type": "string", "maxLength": 4}
"title": "spam",
"type": "object",
"properties": {
"foo": {"type": "boolean"},
"bar": {"type": "string", "maxLength": 4},
},
"required": ["foo", "bar"]
}
"required": ["foo", "bar"],
}

result = generate.json(model, schema_dict, whitespace_pattern="")(
prompt, max_tokens=100, temperature=0, seed=10
Expand All @@ -269,20 +269,36 @@ def test_json_equivalence(model):
"""Test that all methods of generating from json create the same fsm."""

# Different inputs for json generator
def user_callable(user_id: Optional[int], name: Dict[str, int], password: Union[str, int]):
def user_callable(
user_id: Optional[int], name: Dict[str, int], password: Union[str, int]
):
pass

class UserPydantic(BaseModel):
user_id: Optional[int]
name: Dict[str, int]
password: Union[str, int]

user_dict = {'properties': {'user_id': {'anyOf': [{'type': 'integer'}, {'type': 'null'}], 'title': 'User Id'},
'name': {'additionalProperties': {'type': 'integer'}, 'title': 'Name', 'type': 'object'},
'password': {'anyOf': [{'type': 'string'}, {'type': 'integer'}], 'title': 'Password'}},
'required': ['user_id', 'name', 'password'],
'title': 'User',
'type': 'object'}
user_dict = {
"properties": {
"user_id": {
"anyOf": [{"type": "integer"}, {"type": "null"}],
"title": "User Id",
},
"name": {
"additionalProperties": {"type": "integer"},
"title": "Name",
"type": "object",
},
"password": {
"anyOf": [{"type": "string"}, {"type": "integer"}],
"title": "Password",
},
},
"required": ["user_id", "name", "password"],
"title": "User",
"type": "object",
}

user_str = '{"properties": {"user_id": {"anyOf": [{"type": "integer"}, {"type": "null"}], "title": "User Id"}, "name": {"additionalProperties": {"type": "integer"}, "title": "Name", "type": "object"}, "password": {"anyOf": [{"type": "string"}, {"type": "integer"}], "title": "Password"}}, "required": ["user_id", "name", "password"], "title": "User", "type": "object"}'

Expand All @@ -293,10 +309,30 @@ class UserPydantic(BaseModel):
generator_str = generate.json(model, user_str)

# Check finite state machines are the same
assert generator_callable.fsm.states_to_token_maps == generator_pydantic.fsm.states_to_token_maps == generator_dict.fsm.states_to_token_maps == generator_str.fsm.states_to_token_maps
assert generator_callable.fsm.empty_token_ids == generator_pydantic.fsm.empty_token_ids == generator_dict.fsm.empty_token_ids == generator_str.fsm.empty_token_ids
assert generator_callable.fsm.eos_token_id == generator_pydantic.fsm.eos_token_id == generator_dict.fsm.eos_token_id == generator_str.fsm.eos_token_id
assert generator_callable.fsm.final_states == generator_pydantic.fsm.final_states == generator_dict.fsm.final_states == generator_str.fsm.final_states
assert (
generator_callable.fsm.states_to_token_maps
== generator_pydantic.fsm.states_to_token_maps
== generator_dict.fsm.states_to_token_maps
== generator_str.fsm.states_to_token_maps
)
assert (
generator_callable.fsm.empty_token_ids
== generator_pydantic.fsm.empty_token_ids
== generator_dict.fsm.empty_token_ids
== generator_str.fsm.empty_token_ids
)
assert (
generator_callable.fsm.eos_token_id
== generator_pydantic.fsm.eos_token_id
== generator_dict.fsm.eos_token_id
== generator_str.fsm.eos_token_id
)
assert (
generator_callable.fsm.final_states
== generator_pydantic.fsm.final_states
== generator_dict.fsm.final_states
== generator_str.fsm.final_states
)


def test_llamacpp_cfg(model):
Expand Down
72 changes: 54 additions & 18 deletions tests/generate/test_integration_transformers.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -448,14 +448,14 @@ def test_transformers_json_dict():
prompt = "Output some JSON "

schema_dict = {
"title": "spam",
"type": "object",
"properties": {
"foo" : {"type": "integer"},
"bar": {"type": "string", "maxLength": 4}
"title": "spam",
"type": "object",
"properties": {
"foo": {"type": "integer"},
"bar": {"type": "string", "maxLength": 4},
},
"required": ["foo", "bar"]
}
"required": ["foo", "bar"],
}

rng = torch.Generator()
rng.manual_seed(0) # make sure that `bar` is not an int
Expand Down Expand Up @@ -572,20 +572,36 @@ def test_json_equivalence():
"""Test that all methods of generating from json create the same fsm."""

# Different inputs for json generator
def user_callable(user_id: Optional[int], name: Dict[str, int], password: Union[str, int]):
def user_callable(
user_id: Optional[int], name: Dict[str, int], password: Union[str, int]
):
pass

class UserPydantic(BaseModel):
user_id: Optional[int]
name: Dict[str, int]
password: Union[str, int]

user_dict = {'properties': {'user_id': {'anyOf': [{'type': 'integer'}, {'type': 'null'}], 'title': 'User Id'},
'name': {'additionalProperties': {'type': 'integer'}, 'title': 'Name', 'type': 'object'},
'password': {'anyOf': [{'type': 'string'}, {'type': 'integer'}], 'title': 'Password'}},
'required': ['user_id', 'name', 'password'],
'title': 'User',
'type': 'object'}
user_dict = {
"properties": {
"user_id": {
"anyOf": [{"type": "integer"}, {"type": "null"}],
"title": "User Id",
},
"name": {
"additionalProperties": {"type": "integer"},
"title": "Name",
"type": "object",
},
"password": {
"anyOf": [{"type": "string"}, {"type": "integer"}],
"title": "Password",
},
},
"required": ["user_id", "name", "password"],
"title": "User",
"type": "object",
}

user_str = '{"properties": {"user_id": {"anyOf": [{"type": "integer"}, {"type": "null"}], "title": "User Id"}, "name": {"additionalProperties": {"type": "integer"}, "title": "Name", "type": "object"}, "password": {"anyOf": [{"type": "string"}, {"type": "integer"}], "title": "Password"}}, "required": ["user_id", "name", "password"], "title": "User", "type": "object"}'

Expand All @@ -598,10 +614,30 @@ class UserPydantic(BaseModel):
generator_str = generate.json(model, user_str)

# Check finite state machines are the same
assert generator_callable.fsm.states_to_token_maps == generator_pydantic.fsm.states_to_token_maps == generator_dict.fsm.states_to_token_maps == generator_str.fsm.states_to_token_maps
assert generator_callable.fsm.empty_token_ids == generator_pydantic.fsm.empty_token_ids == generator_dict.fsm.empty_token_ids == generator_str.fsm.empty_token_ids
assert generator_callable.fsm.eos_token_id == generator_pydantic.fsm.eos_token_id == generator_dict.fsm.eos_token_id == generator_str.fsm.eos_token_id
assert generator_callable.fsm.final_states == generator_pydantic.fsm.final_states == generator_dict.fsm.final_states == generator_str.fsm.final_states
assert (
generator_callable.fsm.states_to_token_maps
== generator_pydantic.fsm.states_to_token_maps
== generator_dict.fsm.states_to_token_maps
== generator_str.fsm.states_to_token_maps
)
assert (
generator_callable.fsm.empty_token_ids
== generator_pydantic.fsm.empty_token_ids
== generator_dict.fsm.empty_token_ids
== generator_str.fsm.empty_token_ids
)
assert (
generator_callable.fsm.eos_token_id
== generator_pydantic.fsm.eos_token_id
== generator_dict.fsm.eos_token_id
== generator_str.fsm.eos_token_id
)
assert (
generator_callable.fsm.final_states
== generator_pydantic.fsm.final_states
== generator_dict.fsm.final_states
== generator_str.fsm.final_states
)


def test_transformers_logits_vocab_size():
Expand Down
72 changes: 54 additions & 18 deletions tests/generate/test_integration_vllm.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -235,14 +235,14 @@ def test_vllm_json_dict(model):
prompt = "Output some JSON. "

schema_dict = {
"title": "spam",
"type": "object",
"properties": {
"foo" : {"type": "boolean"},
"bar": {"type": "string", "maxLength": 4}
"title": "spam",
"type": "object",
"properties": {
"foo": {"type": "boolean"},
"bar": {"type": "string", "maxLength": 4},
},
"required": ["foo", "bar"]
}
"required": ["foo", "bar"],
}

sampling_params = SamplingParams(temperature=0)
result = generate.json(model, schema_dict, whitespace_pattern="")(
Expand All @@ -257,20 +257,36 @@ def test_json_equivalence(model):
"""Test that all methods of generating from json create the same fsm."""

# Different inputs for json generator
def user_callable(user_id: Optional[int], name: Dict[str, int], password: Union[str, int]):
def user_callable(
user_id: Optional[int], name: Dict[str, int], password: Union[str, int]
):
pass

class UserPydantic(BaseModel):
user_id: Optional[int]
name: Dict[str, int]
password: Union[str, int]

user_dict = {'properties': {'user_id': {'anyOf': [{'type': 'integer'}, {'type': 'null'}], 'title': 'User Id'},
'name': {'additionalProperties': {'type': 'integer'}, 'title': 'Name', 'type': 'object'},
'password': {'anyOf': [{'type': 'string'}, {'type': 'integer'}], 'title': 'Password'}},
'required': ['user_id', 'name', 'password'],
'title': 'User',
'type': 'object'}
user_dict = {
"properties": {
"user_id": {
"anyOf": [{"type": "integer"}, {"type": "null"}],
"title": "User Id",
},
"name": {
"additionalProperties": {"type": "integer"},
"title": "Name",
"type": "object",
},
"password": {
"anyOf": [{"type": "string"}, {"type": "integer"}],
"title": "Password",
},
},
"required": ["user_id", "name", "password"],
"title": "User",
"type": "object",
}

user_str = '{"properties": {"user_id": {"anyOf": [{"type": "integer"}, {"type": "null"}], "title": "User Id"}, "name": {"additionalProperties": {"type": "integer"}, "title": "Name", "type": "object"}, "password": {"anyOf": [{"type": "string"}, {"type": "integer"}], "title": "Password"}}, "required": ["user_id", "name", "password"], "title": "User", "type": "object"}'

Expand All @@ -281,10 +297,30 @@ class UserPydantic(BaseModel):
generator_str = generate.json(model, user_str)

# Check finite state machines are the same
assert generator_callable.fsm.states_to_token_maps == generator_pydantic.fsm.states_to_token_maps == generator_dict.fsm.states_to_token_maps == generator_str.fsm.states_to_token_maps
assert generator_callable.fsm.empty_token_ids == generator_pydantic.fsm.empty_token_ids == generator_dict.fsm.empty_token_ids == generator_str.fsm.empty_token_ids
assert generator_callable.fsm.eos_token_id == generator_pydantic.fsm.eos_token_id == generator_dict.fsm.eos_token_id == generator_str.fsm.eos_token_id
assert generator_callable.fsm.final_states == generator_pydantic.fsm.final_states == generator_dict.fsm.final_states == generator_str.fsm.final_states
assert (
generator_callable.fsm.states_to_token_maps
== generator_pydantic.fsm.states_to_token_maps
== generator_dict.fsm.states_to_token_maps
== generator_str.fsm.states_to_token_maps
)
assert (
generator_callable.fsm.empty_token_ids
== generator_pydantic.fsm.empty_token_ids
== generator_dict.fsm.empty_token_ids
== generator_str.fsm.empty_token_ids
)
assert (
generator_callable.fsm.eos_token_id
== generator_pydantic.fsm.eos_token_id
== generator_dict.fsm.eos_token_id
== generator_str.fsm.eos_token_id
)
assert (
generator_callable.fsm.final_states
== generator_pydantic.fsm.final_states
== generator_dict.fsm.final_states
== generator_str.fsm.final_states
)


@pytest.mark.xfail(
Expand Down

0 comments on commit 944df28

Please sign in to comment.