from __future__ import annotations

import importlib.util
import sys
from pathlib import Path


def load(path: str, **kwargs):
    source = Path(__file__).with_name("fuse2_mlx.py")
    module_name = "mlx_lm.models.fuse2"
    if module_name not in sys.modules:
        spec = importlib.util.spec_from_file_location(module_name, source)
        if spec is None or spec.loader is None:
            raise ImportError(f"cannot load {source}")
        module = importlib.util.module_from_spec(spec)
        sys.modules[module_name] = module
        spec.loader.exec_module(module)

    from mlx_lm import load as mlx_load
    from mlx_lm.utils import MODEL_REMAPPING

    MODEL_REMAPPING["fuse2"] = "fuse2"
    return mlx_load(path, **kwargs)