phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-07-29 15:23:51 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

add back LogProbsConfig

Browse source
This commit is contained in:
Xi Yan 2024-10-07 17:36:39 -07:00
parent 20e78ff241
commit 8a67e7a2bd
4 changed files with 29 additions and 15 deletions
Download patch file Download diff file Expand all files Collapse all files

10
llama_stack/apis/inference/client.py
View file

@ -112,11 +112,19 @@ async def run_main(
content="hello world, write me a 2 sentence poem about the moon" content="hello world, write me a 2 sentence poem about the moon"
) )
cprint(f"User>{message.content}", "green") cprint(f"User>{message.content}", "green")
if logprobs:
logprobs_config = LogProbConfig(
top_k=1,
)
else:
logprobs_config = None
iterator = client.chat_completion( iterator = client.chat_completion(
model=model, model=model,
messages=[message], messages=[message],
stream=stream, stream=stream,
logprobs=logprobs, logprobs=logprobs_config,
) )
if logprobs: if logprobs:

20
llama_stack/apis/inference/inference.py
View file

@ -69,7 +69,7 @@ class ChatCompletionResponseEvent(BaseModel):
event_type: ChatCompletionResponseEventType event_type: ChatCompletionResponseEventType
delta: Union[str, ToolCallDelta] delta: Union[str, ToolCallDelta]
logprobs: Optional[List[float]] = None logprobs: Optional[List[TokenLogProbs]] = None
stop_reason: Optional[StopReason] = None stop_reason: Optional[StopReason] = None
@ -80,7 +80,7 @@ class CompletionRequest(BaseModel):
sampling_params: Optional[SamplingParams] = SamplingParams() sampling_params: Optional[SamplingParams] = SamplingParams()
stream: Optional[bool] = False stream: Optional[bool] = False
logprobs: Optional[bool] = False logprobs: Optional[LogProbConfig] = None
@json_schema_type @json_schema_type
@ -88,7 +88,7 @@ class CompletionResponse(BaseModel):
"""Completion response.""" """Completion response."""
completion_message: CompletionMessage completion_message: CompletionMessage
logprobs: Optional[List[float]] = None logprobs: Optional[List[TokenLogProbs]] = None
@json_schema_type @json_schema_type
@ -97,7 +97,7 @@ class CompletionResponseStreamChunk(BaseModel):
delta: str delta: str
stop_reason: Optional[StopReason] = None stop_reason: Optional[StopReason] = None
logprobs: Optional[List[float]] = None logprobs: Optional[List[TokenLogProbs]] = None
@json_schema_type @json_schema_type
@ -105,7 +105,7 @@ class BatchCompletionRequest(BaseModel):
model: str model: str
content_batch: List[InterleavedTextMedia] content_batch: List[InterleavedTextMedia]
sampling_params: Optional[SamplingParams] = SamplingParams() sampling_params: Optional[SamplingParams] = SamplingParams()
logprobs: Optional[bool] = False logprobs: Optional[LogProbConfig] = None
@json_schema_type @json_schema_type
@ -129,7 +129,7 @@ class ChatCompletionRequest(BaseModel):
) )
stream: Optional[bool] = False stream: Optional[bool] = False
logprobs: Optional[bool] = False logprobs: Optional[LogProbConfig] = None
@json_schema_type @json_schema_type
@ -144,7 +144,7 @@ class ChatCompletionResponse(BaseModel):
"""Chat completion response.""" """Chat completion response."""
completion_message: CompletionMessage completion_message: CompletionMessage
logprobs: Optional[List[float]] = None logprobs: Optional[List[TokenLogProbs]] = None
@json_schema_type @json_schema_type
@ -159,7 +159,7 @@ class BatchChatCompletionRequest(BaseModel):
tool_prompt_format: Optional[ToolPromptFormat] = Field( tool_prompt_format: Optional[ToolPromptFormat] = Field(
default=ToolPromptFormat.json default=ToolPromptFormat.json
) )
logprobs: Optional[bool] = False logprobs: Optional[LogProbConfig] = None
@json_schema_type @json_schema_type
@ -180,7 +180,7 @@ class Inference(Protocol):
content: InterleavedTextMedia, content: InterleavedTextMedia,
sampling_params: Optional[SamplingParams] = SamplingParams(), sampling_params: Optional[SamplingParams] = SamplingParams(),
stream: Optional[bool] = False, stream: Optional[bool] = False,
logprobs: Optional[bool] = None, logprobs: Optional[LogProbConfig] = None,
) -> Union[CompletionResponse, CompletionResponseStreamChunk]: ... ) -> Union[CompletionResponse, CompletionResponseStreamChunk]: ...
@webmethod(route="/inference/chat_completion") @webmethod(route="/inference/chat_completion")
@ -194,7 +194,7 @@ class Inference(Protocol):
tool_choice: Optional[ToolChoice] = ToolChoice.auto, tool_choice: Optional[ToolChoice] = ToolChoice.auto,
tool_prompt_format: Optional[ToolPromptFormat] = ToolPromptFormat.json, tool_prompt_format: Optional[ToolPromptFormat] = ToolPromptFormat.json,
stream: Optional[bool] = False, stream: Optional[bool] = False,
logprobs: Optional[bool] = None, logprobs: Optional[LogProbConfig] = None,
) -> Union[ChatCompletionResponse, ChatCompletionResponseStreamChunk]: ... ) -> Union[ChatCompletionResponse, ChatCompletionResponseStreamChunk]: ...
@webmethod(route="/inference/embeddings") @webmethod(route="/inference/embeddings")

4
llama_stack/distribution/routers/routers.py
View file

@ -101,7 +101,7 @@ class InferenceRouter(Inference):
tool_choice: Optional[ToolChoice] = ToolChoice.auto, tool_choice: Optional[ToolChoice] = ToolChoice.auto,
tool_prompt_format: Optional[ToolPromptFormat] = ToolPromptFormat.json, tool_prompt_format: Optional[ToolPromptFormat] = ToolPromptFormat.json,
stream: Optional[bool] = False, stream: Optional[bool] = False,
logprobs: Optional[bool] = False, logprobs: Optional[LogProbConfig] = False,
) -> AsyncGenerator: ) -> AsyncGenerator:
params = dict( params = dict(
model=model, model=model,
@ -125,7 +125,7 @@ class InferenceRouter(Inference):
content: InterleavedTextMedia, content: InterleavedTextMedia,
sampling_params: Optional[SamplingParams] = SamplingParams(), sampling_params: Optional[SamplingParams] = SamplingParams(),
stream: Optional[bool] = False, stream: Optional[bool] = False,
logprobs: Optional[bool] = False, logprobs: Optional[LogProbConfig] = False,
) -> Union[CompletionResponse, CompletionResponseStreamChunk]: ) -> Union[CompletionResponse, CompletionResponseStreamChunk]:
return await self.routing_table.get_provider_impl(model).completion( return await self.routing_table.get_provider_impl(model).completion(
model=model, model=model,

10
llama_stack/providers/impls/meta_reference/inference/inference.py
View file

@ -58,7 +58,7 @@ class MetaReferenceInferenceImpl(Inference, RoutableProvider):
tool_choice: Optional[ToolChoice] = ToolChoice.auto, tool_choice: Optional[ToolChoice] = ToolChoice.auto,
tool_prompt_format: Optional[ToolPromptFormat] = ToolPromptFormat.json, tool_prompt_format: Optional[ToolPromptFormat] = ToolPromptFormat.json,
stream: Optional[bool] = False, stream: Optional[bool] = False,
logprobs: Optional[bool] = None, logprobs: Optional[LogProbConfig] = None,
) -> AsyncIterator[ ) -> AsyncIterator[
Union[ChatCompletionResponseStreamChunk, ChatCompletionResponse] Union[ChatCompletionResponseStreamChunk, ChatCompletionResponse]
]: ]:
@ -135,7 +135,13 @@ class MetaReferenceInferenceImpl(Inference, RoutableProvider):
assert ( assert (
len(token_result.logprobs) == 1 len(token_result.logprobs) == 1
), "Expected logprob to contain 1 result for the current token" ), "Expected logprob to contain 1 result for the current token"
logprobs.append(token_result.logprobs[0]) logprobs.append(
TokenLogProbs(
logprobs_by_token={
token_result.text: token_result.logprobs[0]
}
)
)
continue continue