【覚え書き】MLXのgenerate, generate_stepの引数が増えていた件

mlx-lm 0.12.0 で確認しました。
pythonのインターラクティブモードで、下記コマンドでチェック。

from mlx_lm.utils import generate_step
from mlx_lm import generate
help(generate)
help(generate_step)

generateについての説明

Help on function generate in module mlx_lm.utils:

generate(model: mlx.nn.layers.base.Module, tokenizer: Union[transformers.tokenization_utils.PreTrainedTokenizer, mlx_lm.tokenizer_utils.TokenizerWrapper], prompt: str, temp: float = 0.0, max_tokens: int = 100, verbose: bool = False, formatter: Optional[Callable] = None, repetition_penalty: Optional[float] = None, repetition_context_size: Optional[int] = None, top_p: float = 1.0, logit_bias: Optional[Dict[int, float]] = None) -> str
Generate text from the model.

Args:
model (nn.Module): The language model.
tokenizer (PreTrainedTokenizer): The tokenizer.
prompt (str): The string prompt.
temp (float): The temperature for sampling (default 0).
max_tokens (int): The maximum number of tokens (default 100).
verbose (bool): If ``True``, print tokens and timing information
(default ``False``).
formatter (Optional[Callable]): A function which takes a token and a
probability and displays it.
repetition_penalty (float, optional): The penalty factor for repeating tokens.
repetition_context_size (int, optional): The number of tokens to consider for repetition penalty.

generate_stepについての説明

Help on function generate_step in module mlx_lm.utils:

generate_step(prompt: mlx.core.array, model: mlx.nn.layers.base.Module, temp: float = 0.0, repetition_penalty: Optional[float] = None, repetition_context_size: Optional[int] = 20, top_p: float = 1.0, logit_bias: Optional[Dict[int, float]] = None) -> Generator[Tuple[mlx.core.array, mlx.core.array], NoneType, NoneType]

A generator producing text based on the given prompt from the model.

Args:
prompt (mx.array): The input prompt.
model (nn.Module): The model to use for generation.
temp (float): The temperature for sampling, if 0 the argmax is used.
repetition_penalty (float, optional): The penalty factor for repeating tokens.
repetition_context_size (int, optional): The number of tokens to consider for repetition penalty (default 20).
top_p (float, optional): Nulceus sampling, higher means model considers more less likely words

Yields:
Generator[Tuple[mx.array, mx.array]]: A generator producing
one token and probability per call.

共通するのは、
repetition_penalty (float, optional) と
repetition_context_size (int, optional) が追加されて、
generation_stepの方は加えて、
top_p (float, optional) が追加されています。

top_p と temperature はどっちかを一定にして、片方だけを調整することがお勧めだということをどっかで読んで覚えがあります。
いずれも確率を指定しているので、0~1の間の値だと自分は理解しています。

あと、end of tokenの指定も覚え書きとしてサンプルをメモ

model, tokenizer = load(
"qwen/Qwen-7B",
tokenizer_config={"eos_token": "<|endoftext|>", "trust_remote_code": True},
)

もうひとつおまけ

NAME
mlx_lm - # Copyright © 2023-2024 Apple Inc.

PACKAGE CONTENTS
convert
fuse
generate
gguf
lora
merge
models (package)
sample_utils
server
tokenizer_utils
tuner (package)
utils
version

VERSION
0.12.0

下で確認

import mlx_lm
help(mlx_lm)

---

#AI #AIとやってみた #やってみた #大規模言語モデル #mlx #MacbookPro #ローカルLLM