phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-08-07 02:58:21 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

top_k=1

Browse source
This commit is contained in:
Sixian Yi 2025-01-28 14:18:18 -08:00
parent 0052089ab8
commit 3816355e1d
1 changed files with 14 additions and 16 deletions
Download patch file Download diff file Expand all files Collapse all files

30
tests/client-sdk/inference/test_inference.py
View file

@ -16,11 +16,13 @@ PROVIDER_TOOL_PROMPT_FORMAT = {
"remote::fireworks": "json", "remote::fireworks": "json",
} }
PROVIDER_LOGPROBS_TOP_K = { PROVIDER_LOGPROBS_TOP_K = set(
"remote::together": 1, {
"remote::fireworks": 3, "remote::together",
# "remote:vllm" "remote::fireworks",
} # "remote:vllm"
}
)
@pytest.fixture(scope="session") @pytest.fixture(scope="session")
@ -95,7 +97,6 @@ def test_completion_log_probs_non_streaming(
if inference_provider_type not in PROVIDER_LOGPROBS_TOP_K: if inference_provider_type not in PROVIDER_LOGPROBS_TOP_K:
pytest.xfail(f"{inference_provider_type} doesn't support log probs yet") pytest.xfail(f"{inference_provider_type} doesn't support log probs yet")
logprobs_top_k = PROVIDER_LOGPROBS_TOP_K[inference_provider_type]
response = llama_stack_client.inference.completion( response = llama_stack_client.inference.completion(
content="Complete the sentence: Micheael Jordan is born in ", content="Complete the sentence: Micheael Jordan is born in ",
stream=False, stream=False,
@ -104,15 +105,14 @@ def test_completion_log_probs_non_streaming(
"max_tokens": 5, "max_tokens": 5,
}, },
logprobs={ logprobs={
"top_k": logprobs_top_k, "top_k": 1,
}, },
) )
assert response.logprobs, "Logprobs should not be empty" assert response.logprobs, "Logprobs should not be empty"
assert 1 <= len(response.logprobs) <= 5 assert (
assert all( 1 <= len(response.logprobs) <= 5
len(logprob.logprobs_by_token) == logprobs_top_k ) # each token has 1 logprob and here max_tokens=5
for logprob in response.logprobs assert all(len(logprob.logprobs_by_token) == 1 for logprob in response.logprobs)
)
def test_completion_log_probs_streaming( def test_completion_log_probs_streaming(
@ -121,7 +121,6 @@ def test_completion_log_probs_streaming(
if inference_provider_type not in PROVIDER_LOGPROBS_TOP_K: if inference_provider_type not in PROVIDER_LOGPROBS_TOP_K:
pytest.xfail(f"{inference_provider_type} doesn't support log probs yet") pytest.xfail(f"{inference_provider_type} doesn't support log probs yet")
logprobs_top_k = PROVIDER_LOGPROBS_TOP_K[inference_provider_type]
response = llama_stack_client.inference.completion( response = llama_stack_client.inference.completion(
content="Complete the sentence: Micheael Jordan is born in ", content="Complete the sentence: Micheael Jordan is born in ",
stream=True, stream=True,
@ -130,7 +129,7 @@ def test_completion_log_probs_streaming(
"max_tokens": 5, "max_tokens": 5,
}, },
logprobs={ logprobs={
"top_k": logprobs_top_k, "top_k": 1,
}, },
) )
streamed_content = [chunk for chunk in response] streamed_content = [chunk for chunk in response]
@ -138,8 +137,7 @@ def test_completion_log_probs_streaming(
if chunk.delta: # if there's a token, we expect logprobs if chunk.delta: # if there's a token, we expect logprobs
assert chunk.logprobs, "Logprobs should not be empty" assert chunk.logprobs, "Logprobs should not be empty"
assert all( assert all(
len(logprob.logprobs_by_token) == logprobs_top_k len(logprob.logprobs_by_token) == 1 for logprob in chunk.logprobs
for logprob in chunk.logprobs
) )
else: # no token, no logprobs else: # no token, no logprobs
assert not chunk.logprobs, "Logprobs should be empty" assert not chunk.logprobs, "Logprobs should be empty"