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

fix nvidia sampling logic

Browse source
This commit is contained in:
Hardik Shah 2025-01-14 18:00:30 -08:00 • committed by Ashwin Bharambe
parent 0edd3ce78b
commit cb6c734460
1 changed files with 8 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

9
llama_stack/providers/remote/inference/nvidia/openai_utils.py
View file

@ -8,6 +8,11 @@ import json
import warnings import warnings
from typing import Any, AsyncGenerator, Dict, Generator, List, Optional from typing import Any, AsyncGenerator, Dict, Generator, List, Optional
from llama_models.datatypes import (
GreedySamplingStrategy,
TopKSamplingStrategy,
TopPSamplingStrategy,
)
from llama_models.llama3.api.datatypes import ( from llama_models.llama3.api.datatypes import (
BuiltinTool, BuiltinTool,
StopReason, StopReason,
@ -272,9 +277,11 @@ def convert_chat_completion_request(
if strategy.top_k != -1 and strategy.top_k < 1: if strategy.top_k != -1 and strategy.top_k < 1:
warnings.warn("top_k must be -1 or >= 1") warnings.warn("top_k must be -1 or >= 1")
nvext.update(top_k=strategy.top_k) nvext.update(top_k=strategy.top_k)
elif strategy.strategy == "greedy": elif isinstance(strategy, GreedySamplingStrategy):
nvext.update(top_k=-1) nvext.update(top_k=-1)
payload.update(temperature=strategy.temperature) payload.update(temperature=strategy.temperature)
else:
raise ValueError(f"Unsupported sampling strategy: {strategy}")
return payload return payload