fount-vlm-nell-02 0.3.11__py3-none-any.whl

mlx_vlm/convert.py

@@ -0,0 +1,284 @@
+import argparse
+import glob
+import shutil
+from pathlib import Path
+from typing import Callable, Optional, Union
+
+import mlx.core as mx
+import mlx.nn as nn
+from mlx.utils import tree_map_with_path
+from mlx_lm.utils import dequantize_model, quantize_model
+
+from .utils import (
+    MODEL_CONVERSION_DTYPES,
+    fetch_from_hub,
+    get_model_path,
+    save_config,
+    save_weights,
+    skip_multimodal_module,
+    upload_to_hub,
+)
+
+QUANT_RECIPES = [
+    "mixed_2_6",
+    "mixed_3_4",
+    "mixed_3_5",
+    "mixed_3_6",
+    "mixed_3_8",
+    "mixed_4_6",
+    "mixed_4_8",
+]
+
+
+def mixed_quant_predicate_builder(
+    recipe: str, model: nn.Module
+) -> Callable[[str, nn.Module], Union[bool, dict]]:
+    group_size = 64
+
+    recipe_config = {
+        "mixed_2_6": (2, 6),
+        "mixed_3_4": (3, 4),
+        "mixed_3_5": (3, 5),
+        "mixed_3_6": (3, 6),
+        "mixed_3_8": (3, 8),
+        "mixed_4_6": (4, 6),
+        "mixed_4_8": (4, 8),
+    }
+
+    if recipe not in recipe_config:
+        raise ValueError(f"Invalid quant recipe {recipe}")
+
+    low_bits, high_bits = recipe_config[recipe]
+
+    down_keys = [k for k, _ in model.named_modules() if "down_proj" in k]
+    if len(down_keys) == 0:
+        raise ValueError("Model does not have expected keys for mixed quant.")
+
+    # Look for the layer index location in the path:
+    for layer_location, k in enumerate(down_keys[0].split(".")):
+        if k.isdigit():
+            break
+    num_layers = len(model.layers)
+
+    def mixed_quant_predicate(
+        path: str,
+        module: nn.Module,
+    ) -> Union[bool, dict]:
+        """Implements mixed quantization predicates with similar choices to, for example, llama.cpp's Q4_K_M.
+        Ref: https://github.com/ggerganov/llama.cpp/blob/917786f43d0f29b7c77a0c56767c0fa4df68b1c5/src/llama.cpp#L5265
+        By Alex Barron: https://gist.github.com/barronalex/84addb8078be21969f1690c1454855f3
+        """
+
+        if skip_multimodal_module(path):
+            return False
+        if not hasattr(module, "to_quantized"):
+            return False
+        if module.weight.shape[1] % group_size != 0:
+            return False
+
+        path_parts = path.split(".")
+        index = 0
+
+        if len(path_parts) > layer_location:
+            element = path_parts[layer_location]
+            if element.isdigit():
+                index = int(element)
+
+        use_more_bits = (
+            index < num_layers // 8
+            or index >= 7 * num_layers // 8
+            or (index - num_layers // 8) % 3 == 2
+        )
+
+        if use_more_bits and ("v_proj" in path or "down_proj" in path):
+            return {"group_size": group_size, "bits": high_bits}
+
+        if "lm_head" in path or "embed_tokens" in path:
+            return {"group_size": group_size, "bits": high_bits}
+
+        return {"group_size": group_size, "bits": low_bits}
+
+    return mixed_quant_predicate
+
+
+def convert(
+    hf_path: str,
+    mlx_path: str = "mlx_model",
+    quantize: bool = False,
+    q_group_size: int = 64,
+    q_bits: int = 4,
+    q_mode: str = "affine",
+    dtype: Optional[str] = None,
+    upload_repo: str = None,
+    revision: Optional[str] = None,
+    dequantize: bool = False,
+    trust_remote_code: bool = True,
+    quant_predicate: Optional[str] = None,
+):
+    print("[INFO] Loading")
+    model_path = get_model_path(hf_path, revision=revision)
+    model, config, processor = fetch_from_hub(
+        model_path, lazy=True, trust_remote_code=trust_remote_code
+    )
+
+    def base_quant_predicate(path, module):
+        if skip_multimodal_module(path):
+            return False
+        return True
+
+    if isinstance(quant_predicate, str):
+        quant_predicate = mixed_quant_predicate_builder(quant_predicate, model)
+
+    quant_predicate = quant_predicate or base_quant_predicate
+
+    if dtype is None:
+        dtype = config.get("torch_dtype", None)
+    if dtype in MODEL_CONVERSION_DTYPES:
+        print("[INFO] Using dtype:", dtype)
+        dtype = getattr(mx, dtype)
+        cast_predicate = getattr(model, "cast_predicate", lambda _: True)
+
+        def set_dtype(k, v):
+            if cast_predicate(k) and mx.issubdtype(v.dtype, mx.floating):
+                return v.astype(dtype)
+            else:
+                return v
+
+        model.update(tree_map_with_path(set_dtype, model.parameters()))
+
+    if quantize and dequantize:
+        raise ValueError("Choose either quantize or dequantize, not both.")
+
+    if quantize:
+        print("[INFO] Quantizing")
+        config.setdefault("vision_config", {})
+        model, config = quantize_model(
+            model,
+            config,
+            q_group_size,
+            q_bits,
+            mode=q_mode,
+            quant_predicate=quant_predicate,
+        )
+
+    if dequantize:
+        print("[INFO] Dequantizing")
+        model = dequantize_model(model)
+
+    if isinstance(mlx_path, str):
+        mlx_path = Path(mlx_path)
+
+    save_weights(mlx_path, model, donate_weights=True)
+
+    # Copy Python and JSON files from the model path to the MLX path
+    for pattern in ["*.py", "*.json"]:
+        files = glob.glob(str(model_path / pattern))
+        for file in files:
+            # Skip the index file - save_weights() already generated the correct one
+            if Path(file).name == "model.safetensors.index.json":
+                continue
+            shutil.copy(file, mlx_path)
+
+    processor.save_pretrained(mlx_path)
+
+    save_config(config, config_path=mlx_path / "config.json")
+
+    if upload_repo is not None:
+        upload_to_hub(mlx_path, upload_repo, hf_path)
+
+
+def configure_parser() -> argparse.ArgumentParser:
+    """
+    Configures and returns the argument parser for the script.
+
+    Returns:
+        argparse.ArgumentParser: Configured argument parser.
+    """
+    parser = argparse.ArgumentParser(
+        description="Convert Hugging Face model to MLX format"
+    )
+    parser.add_argument(
+        "--hf-path",
+        "--model",
+        type=str,
+        help="Path to the model. This can be a local path or a Hugging Face Hub model identifier.",
+    )
+    parser.add_argument(
+        "--revision",
+        type=str,
+        help="Hugging Face revision (branch), when converting a model from the Hub.",
+        default=None,
+    )
+    parser.add_argument(
+        "--mlx-path", type=str, default="mlx_model", help="Path to save the MLX model."
+    )
+    parser.add_argument(
+        "-q", "--quantize", help="Generate a quantized model.", action="store_true"
+    )
+    parser.add_argument(
+        "--q-group-size",
+        help="Group size for quantization.",
+        type=int,
+        default=None,
+    )
+    parser.add_argument(
+        "--q-bits",
+        help="Bits per weight for quantization.",
+        type=int,
+        default=None,
+    )
+    parser.add_argument(
+        "--q-mode",
+        help="The quantization mode.",
+        type=str,
+        choices=["affine", "mxfp4", "nvfp4", "mxfp8"],
+        default="affine",
+    )
+    parser.add_argument(
+        "--dtype",
+        help="Type to save the parameter. Defaults to config.json's `torch_dtype` or the current model weights dtype",
+        type=str,
+        choices=MODEL_CONVERSION_DTYPES,
+        default=None,
+    )
+    parser.add_argument(
+        "--quant-predicate",
+        help=f"Mixed-bit quantization recipe.",
+        choices=QUANT_RECIPES,
+        type=str,
+        required=False,
+    )
+    parser.add_argument(
+        "--upload-repo",
+        help="The Hugging Face repo to upload the model to.",
+        type=str,
+        default=None,
+    )
+    parser.add_argument(
+        "-d",
+        "--dequantize",
+        help="Dequantize a quantized model.",
+        action="store_true",
+        default=False,
+    )
+    parser.add_argument(
+        "--trust-remote-code",
+        help="Trust remote code.",
+        action="store_true",
+        default=False,
+    )
+    return parser
+
+
+def main():
+    parser = configure_parser()
+    args = parser.parse_args()
+    convert(**vars(args))
+
+
+if __name__ == "__main__":
+    print(
+        "Calling `python -m mlx_vlm.convert ...` directly is deprecated."
+        " Use `mlx_vlm.convert ...` or `python -m mlx_vlm convert ...` instead."
+    )
+    main()

mlx_vlm/deprecation.py

@@ -0,0 +1,52 @@
+import functools
+import warnings
+from typing import Callable, Optional
+
+
+def deprecate(
+    remove_version: str,
+    message: str,
+    instead: Optional[str] = None,
+    since: Optional[str] = None,
+) -> Callable:
+    """
+    Mark a function or method as deprecated.
+
+    Args:
+        remove_version: Version when this will be removed
+        message: Deprecation message
+        instead: What to use instead
+        since: Version when this was deprecated
+
+    Example:
+        @deprecate(
+            remove_version="2.0.0",
+            message="Legacy API function",
+            instead="new_api()",
+            since="1.0.0"
+        )
+        def old_function():
+            pass
+    """
+
+    def decorator(func: Callable) -> Callable:
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            msg = f"`{func.__name__}` is deprecated"
+
+            if since:
+                msg += f" since v{since}"
+
+            msg += f". {message}"
+
+            if instead:
+                msg += f" Use `{instead}` instead."
+
+            msg += f" Will be removed in v{remove_version}."
+
+            warnings.warn(msg, DeprecationWarning, stacklevel=2)
+            return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator