From a5e173802ea0da999923d240f65e94ec8ad3c415 Mon Sep 17 00:00:00 2001 From: madroid Date: Tue, 26 Nov 2024 00:10:14 +0800 Subject: [PATCH] docs: update stream_generate return type annotation (#1121) Improve documentation clarity by: 1. Fix return type annotation to correctly reflect GenerationResponse 2. Simplify docstring by referencing GenerationResponse class 3. Remove redundant field descriptions --- llms/mlx_lm/utils.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/llms/mlx_lm/utils.py b/llms/mlx_lm/utils.py index 496ae4fc..5abd396d 100644 --- a/llms/mlx_lm/utils.py +++ b/llms/mlx_lm/utils.py @@ -300,7 +300,7 @@ def stream_generate( prompt: Union[str, List[int]], max_tokens: int = 100, **kwargs, -) -> Generator[Tuple[str, int, mx.array], None, None]: +) -> Generator[GenerationResponse, None, None]: """ A generator producing text based on the given prompt from the model. @@ -313,8 +313,8 @@ def stream_generate( See :func:`generate_step` for more details. Yields: - Tuple[str, int, mx.array]: - The next text segment, token, and vector of log probabilities. + GenerationResponse: An instance containing the generated text segment and + associated metadata. See :class:`GenerationResponse` for details. """ if not isinstance(tokenizer, TokenizerWrapper): tokenizer = TokenizerWrapper(tokenizer)