Load the model and tokenizer from a given path or a huggingface repository. Args: path_or_hf_repo (Path): The path or the huggingfac

In-place extend this cache with the other cache.

Return only generated tokens ignoring batch generation responses. Returns: List of GenerationBatch.Response objects

Construct the model's cache for use in generation. This function will defer the cache construction to the model if it has a ``make_cache

A generator producing text based on the given prompt from the model. Args: model (nn.Module): The model to use for generation.

A generator producing token ids based on the given prompt from the model. Args: prompt (mx.array): The input prompt. model (

Load a prompt cache from a file. Args: file_name (str): The ``.safetensors`` file name. return_metadata (bool): Whether or n

Save a pre-computed prompt cache to a file. Args: file_name (str): The ``.safetensors`` file name. cache (List[Any]): The mo

Make logits processors for use with ``generate_step``. Args: repetition_penalty (float, optional): A (sign-aware) multiplicative

Trim the model's cache by the given number of tokens. This function will trim the cache if possible (in-place) and return the number of

Load a huggingface tokenizer and try to infer the type of streaming detokenizer to use.

Extend this batch with another generation batch.

Generate a complete response from the model. Args: model (nn.Module): The language model. tokenizer (PreTrainedTokenizer): The

Apply top-p (nucleus) sampling to logits. Args: logprobs: A vector of log probabilities. top_p: The cumulative probability t

Make a sampler function for use with ``generate_step``. Args: temp (float): The temperature for sampling, if 0 the argmax is used.

Functional RMSNorm with no learnable parameters.