`.typings/mlx_lm/tuner/utils.pyi`

58 lines
"""
This type stub file was generated by pyright.
"""

import mlx.nn as nn
from typing import Dict

def build_schedule(schedule_config: Dict):  # -> Any:
    """
    Build a learning rate schedule from the given config.
    """
    ...

def linear_to_lora_layers(
    model: nn.Module, num_layers: int, config: Dict, use_dora: bool = ...
):  # -> None:
    """
    Convert some of the models linear layers to lora layers.

    Args:
        model (nn.Module): The neural network model.
        num_layers (int): The number of blocks to convert to lora layers
        starting from the last layer.
        config (dict): More configuration parameters for LoRA, including the
          rank, scale, and optional layer keys.
        use_dora (bool): If True, uses DoRA instead of LoRA.
          Default: ``False``
    """
    ...

def load_adapters(model: nn.Module, adapter_path: str) -> nn.Module:
    """
    Load any fine-tuned adapters / layers.

    Args:
        model (nn.Module): The neural network model.
        adapter_path (str): Path to the adapter configuration file.

    Returns:
        nn.Module: The updated model with LoRA layers applied.
    """
    ...

def remove_lora_layers(model: nn.Module) -> nn.Module:
    """
    Remove the LoRA layers from the model.

    Args:
        model (nn.Module): The model with LoRA layers.

    Returns:
        nn.Module: The model without LoRA layers.
    """
    ...

def print_trainable_parameters(model):  # -> None:
    ...