zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-08-29 18:26:37 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

docs: update stream_generate return type annotation

Browse Source 
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
This commit is contained in:
madroid 2024-11-25 14:31:35 +08:00
parent 0f135396ae
commit 5cd7531c22
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
llms/mlx_lm/utils.py
View File

@ -300,7 +300,7 @@ def stream_generate(
prompt: Union[str, List[int]], prompt: Union[str, List[int]],
max_tokens: int = 100, max_tokens: int = 100,
**kwargs, **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. 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. See :func:`generate_step` for more details.
Yields: Yields:
Tuple[str, int, mx.array]: GenerationResponse: An instance containing the generated text segment and
The next text segment, token, and vector of log probabilities. associated metadata. See :class:`GenerationResponse` for details.
""" """
if not isinstance(tokenizer, TokenizerWrapper): if not isinstance(tokenizer, TokenizerWrapper):
tokenizer = TokenizerWrapper(tokenizer) tokenizer = TokenizerWrapper(tokenizer)