From 1bdeb0eae01cd3be38c7ea17b96523c38a0a4838 Mon Sep 17 00:00:00 2001 From: kevinbu233 Date: Sat, 3 Aug 2024 22:06:59 +0000 Subject: [PATCH] fixed guided test --- tests/entrypoints/llm/test_guided_generate.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/entrypoints/llm/test_guided_generate.py b/tests/entrypoints/llm/test_guided_generate.py index 977aca0441bb6..873e115421257 100644 --- a/tests/entrypoints/llm/test_guided_generate.py +++ b/tests/entrypoints/llm/test_guided_generate.py @@ -111,6 +111,7 @@ def test_guided_grammar(sample_sql_statements, llm): sampling_params = SamplingParams( temperature=0.8, top_p=0.95, + max_tokens=1000, ) outputs = llm.generate( prompts=("Generate a sql state that select col_1 from " @@ -127,7 +128,6 @@ def test_guided_grammar(sample_sql_statements, llm): generated_text = output.outputs[0].text assert generated_text is not None - # use Lark to parse the output, and make sure it's a valid parse tree from lark import Lark parser = Lark(sample_sql_statements)