From 20e78ff2418b2947b006db1b00a40179d450c5c5 Mon Sep 17 00:00:00 2001 From: Xi Yan Date: Mon, 7 Oct 2024 16:45:06 -0700 Subject: [PATCH] fix log probs --- llama_stack/apis/inference/client.py | 18 ++++++++++++----- llama_stack/apis/inference/inference.py | 20 +++++++++---------- llama_stack/distribution/routers/routers.py | 4 ++-- .../meta_reference/inference/generation.py | 2 +- .../meta_reference/inference/inference.py | 7 +++++-- 5 files changed, 31 insertions(+), 20 deletions(-) diff --git a/llama_stack/apis/inference/client.py b/llama_stack/apis/inference/client.py index 5cfae633c..0a6c3e6e2 100644 --- a/llama_stack/apis/inference/client.py +++ b/llama_stack/apis/inference/client.py @@ -6,7 +6,6 @@ import asyncio import json -import sys from typing import Any, AsyncGenerator, List, Optional import fire @@ -101,7 +100,9 @@ class InferenceClient(Inference): print(f"Error with parsing or validation: {e}") -async def run_main(host: str, port: int, stream: bool, model: Optional[str]): +async def run_main( + host: str, port: int, stream: bool, model: Optional[str], logprobs: bool +): client = InferenceClient(f"http://{host}:{port}") if not model: @@ -115,9 +116,15 @@ async def run_main(host: str, port: int, stream: bool, model: Optional[str]): model=model, messages=[message], stream=stream, + logprobs=logprobs, ) - async for log in EventLogger().log(iterator): - log.print() + + if logprobs: + async for chunk in iterator: + cprint(f"Response: {chunk}", "red") + else: + async for log in EventLogger().log(iterator): + log.print() async def run_mm_main( @@ -149,13 +156,14 @@ def main( port: int, stream: bool = True, mm: bool = False, + logprobs: bool = False, file: Optional[str] = None, model: Optional[str] = None, ): if mm: asyncio.run(run_mm_main(host, port, stream, file, model)) else: - asyncio.run(run_main(host, port, stream, model)) + asyncio.run(run_main(host, port, stream, model, logprobs)) if __name__ == "__main__": diff --git a/llama_stack/apis/inference/inference.py b/llama_stack/apis/inference/inference.py index 428f29b88..abb2822e3 100644 --- a/llama_stack/apis/inference/inference.py +++ b/llama_stack/apis/inference/inference.py @@ -69,7 +69,7 @@ class ChatCompletionResponseEvent(BaseModel): event_type: ChatCompletionResponseEventType delta: Union[str, ToolCallDelta] - logprobs: Optional[List[TokenLogProbs]] = None + logprobs: Optional[List[float]] = None stop_reason: Optional[StopReason] = None @@ -80,7 +80,7 @@ class CompletionRequest(BaseModel): sampling_params: Optional[SamplingParams] = SamplingParams() stream: Optional[bool] = False - logprobs: Optional[LogProbConfig] = None + logprobs: Optional[bool] = False @json_schema_type @@ -88,7 +88,7 @@ class CompletionResponse(BaseModel): """Completion response.""" completion_message: CompletionMessage - logprobs: Optional[List[TokenLogProbs]] = None + logprobs: Optional[List[float]] = None @json_schema_type @@ -97,7 +97,7 @@ class CompletionResponseStreamChunk(BaseModel): delta: str stop_reason: Optional[StopReason] = None - logprobs: Optional[List[TokenLogProbs]] = None + logprobs: Optional[List[float]] = None @json_schema_type @@ -105,7 +105,7 @@ class BatchCompletionRequest(BaseModel): model: str content_batch: List[InterleavedTextMedia] sampling_params: Optional[SamplingParams] = SamplingParams() - logprobs: Optional[LogProbConfig] = None + logprobs: Optional[bool] = False @json_schema_type @@ -129,7 +129,7 @@ class ChatCompletionRequest(BaseModel): ) stream: Optional[bool] = False - logprobs: Optional[LogProbConfig] = None + logprobs: Optional[bool] = False @json_schema_type @@ -144,7 +144,7 @@ class ChatCompletionResponse(BaseModel): """Chat completion response.""" completion_message: CompletionMessage - logprobs: Optional[List[TokenLogProbs]] = None + logprobs: Optional[List[float]] = None @json_schema_type @@ -159,7 +159,7 @@ class BatchChatCompletionRequest(BaseModel): tool_prompt_format: Optional[ToolPromptFormat] = Field( default=ToolPromptFormat.json ) - logprobs: Optional[LogProbConfig] = None + logprobs: Optional[bool] = False @json_schema_type @@ -180,7 +180,7 @@ class Inference(Protocol): content: InterleavedTextMedia, sampling_params: Optional[SamplingParams] = SamplingParams(), stream: Optional[bool] = False, - logprobs: Optional[LogProbConfig] = None, + logprobs: Optional[bool] = None, ) -> Union[CompletionResponse, CompletionResponseStreamChunk]: ... @webmethod(route="/inference/chat_completion") @@ -194,7 +194,7 @@ class Inference(Protocol): tool_choice: Optional[ToolChoice] = ToolChoice.auto, tool_prompt_format: Optional[ToolPromptFormat] = ToolPromptFormat.json, stream: Optional[bool] = False, - logprobs: Optional[LogProbConfig] = None, + logprobs: Optional[bool] = None, ) -> Union[ChatCompletionResponse, ChatCompletionResponseStreamChunk]: ... @webmethod(route="/inference/embeddings") diff --git a/llama_stack/distribution/routers/routers.py b/llama_stack/distribution/routers/routers.py index c360bcfb0..fb562dc2f 100644 --- a/llama_stack/distribution/routers/routers.py +++ b/llama_stack/distribution/routers/routers.py @@ -101,7 +101,7 @@ class InferenceRouter(Inference): tool_choice: Optional[ToolChoice] = ToolChoice.auto, tool_prompt_format: Optional[ToolPromptFormat] = ToolPromptFormat.json, stream: Optional[bool] = False, - logprobs: Optional[LogProbConfig] = None, + logprobs: Optional[bool] = False, ) -> AsyncGenerator: params = dict( model=model, @@ -125,7 +125,7 @@ class InferenceRouter(Inference): content: InterleavedTextMedia, sampling_params: Optional[SamplingParams] = SamplingParams(), stream: Optional[bool] = False, - logprobs: Optional[LogProbConfig] = None, + logprobs: Optional[bool] = False, ) -> Union[CompletionResponse, CompletionResponseStreamChunk]: return await self.routing_table.get_provider_impl(model).completion( model=model, diff --git a/llama_stack/providers/impls/meta_reference/inference/generation.py b/llama_stack/providers/impls/meta_reference/inference/generation.py index 4351a3d56..27e086e0f 100644 --- a/llama_stack/providers/impls/meta_reference/inference/generation.py +++ b/llama_stack/providers/impls/meta_reference/inference/generation.py @@ -297,7 +297,7 @@ class Llama: token=next_token[0].item(), text=self.tokenizer.decode(next_token.tolist()), logprobs=( - token_logprobs[:, prev_pos + 1 : cur_pos + 1][0].tolist() + token_logprobs[:, cur_pos : cur_pos + 1][0].tolist() if logprobs else None ), diff --git a/llama_stack/providers/impls/meta_reference/inference/inference.py b/llama_stack/providers/impls/meta_reference/inference/inference.py index e89d8ec4c..b8cccf949 100644 --- a/llama_stack/providers/impls/meta_reference/inference/inference.py +++ b/llama_stack/providers/impls/meta_reference/inference/inference.py @@ -58,7 +58,7 @@ class MetaReferenceInferenceImpl(Inference, RoutableProvider): tool_choice: Optional[ToolChoice] = ToolChoice.auto, tool_prompt_format: Optional[ToolPromptFormat] = ToolPromptFormat.json, stream: Optional[bool] = False, - logprobs: Optional[LogProbConfig] = None, + logprobs: Optional[bool] = None, ) -> AsyncIterator[ Union[ChatCompletionResponseStreamChunk, ChatCompletionResponse] ]: @@ -132,7 +132,10 @@ class MetaReferenceInferenceImpl(Inference, RoutableProvider): if not request.stream: if request.logprobs: - logprobs.append(token_result.logprob) + assert ( + len(token_result.logprobs) == 1 + ), "Expected logprob to contain 1 result for the current token" + logprobs.append(token_result.logprobs[0]) continue