Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add test that verifies that the generation config passed in at model.predict() is used correctly. #3523

Merged
merged 3 commits into from
Aug 11, 2023
Merged

Add test that verifies that the generation config passed in at model.predict() is used correctly. #3523

Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
60a496e
Add test for generation config
justinxzhao Aug 11, 2023
071865c
Add check that the state of the model object is unchanged.
justinxzhao Aug 11, 2023
0c038e2
PR feedback
justinxzhao Aug 11, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 21 additions & 1 deletion tests/integration_tests/test_llm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@

LOCAL_BACKEND = {"type": "local"}
TEST_MODEL_NAME = "hf-internal-testing/tiny-random-GPTJForCausalLM"
MAX_NEW_TOKENS_TEST_DEFAULT = 5

RAY_BACKEND = {
"type": "ray",
Expand All @@ -46,6 +47,11 @@
}


def get_num_non_empty_tokens(iterable):
"""Returns the number of non-empty tokens."""
return len(list(filter(bool, iterable)))
arnavgarg1 marked this conversation as resolved.
Show resolved Hide resolved


@pytest.fixture(scope="module")
def local_backend():
return LOCAL_BACKEND
Expand Down Expand Up @@ -79,7 +85,7 @@ def get_generation_config():
"top_p": 0.75,
"top_k": 40,
"num_beams": 4,
"max_new_tokens": 5,
"max_new_tokens": MAX_NEW_TOKENS_TEST_DEFAULT,
}


Expand Down Expand Up @@ -134,6 +140,20 @@ def test_llm_text_to_text(tmpdir, backend, ray_cluster_4cpu):
assert preds["Answer_probabilities"]
assert preds["Answer_probability"]

# Check that in-line generation parameters are used. Original prediction uses max_new_tokens = 5.
assert get_num_non_empty_tokens(preds["Answer_predictions"][0]) > 3
justinxzhao marked this conversation as resolved.
Show resolved Hide resolved
original_max_new_tokens = model.model.generation.max_new_tokens

# This prediction uses max_new_tokens = 2.
preds, _ = model.predict(
dataset=dataset_filename, output_directory=str(tmpdir), split="test", generation_config={"max_new_tokens": 2}
)
preds = convert_preds(preds)
assert get_num_non_empty_tokens(preds["Answer_predictions"][0]) < 3

# Check that the state of the model is unchanged.
assert model.model.generation.max_new_tokens == original_max_new_tokens


@pytest.mark.llm
@pytest.mark.parametrize(
Expand Down
Loading