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

add completion logprobs

Browse source
This commit is contained in:
Xi Yan 2025-01-15 16:48:39 -08:00
parent 965644ce68
commit 9c13a7b76b
2 changed files with 25 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

14
llama_stack/providers/remote/inference/fireworks/fireworks.py
View file

@ -168,7 +168,10 @@ class FireworksInferenceAdapter(
yield chunk
def _build_options(
self, sampling_params: Optional[SamplingParams], fmt: ResponseFormat
self,
sampling_params: Optional[SamplingParams],
fmt: ResponseFormat,
logprobs: Optional[LogProbConfig],
) -> dict:
options = get_sampling_options(sampling_params)
options.setdefault("max_tokens", 512)
@ -187,6 +190,11 @@ class FireworksInferenceAdapter(
else:
raise ValueError(f"Unknown response format {fmt.type}")
if logprobs and logprobs.top_k:
options["logprobs"] = logprobs.top_k
if options["logprobs"] <= 0 or options["logprobs"] >= 5:
raise ValueError("Required range: 0 < top_k < 5")
return options
async def chat_completion(
@ -280,7 +288,9 @@ class FireworksInferenceAdapter(
"model": request.model,
**input_dict,
"stream": request.stream,
**self._build_options(request.sampling_params, request.response_format),
**self._build_options(
request.sampling_params, request.response_format, request.logprobs
),
}
async def embeddings(