160 lines
4.3 KiB
Python
160 lines
4.3 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
|
|
import sys
|
|
import types
|
|
from types import SimpleNamespace
|
|
|
|
|
|
class _DummyMetal:
|
|
@staticmethod
|
|
def is_available():
|
|
return False
|
|
|
|
|
|
class _DummyMX:
|
|
metal = _DummyMetal()
|
|
|
|
@staticmethod
|
|
def set_wired_limit(_limit):
|
|
return None
|
|
|
|
@staticmethod
|
|
def device_info():
|
|
return {"max_recommended_working_set_size": 1024}
|
|
|
|
|
|
class _DummyTokenizer:
|
|
pass
|
|
|
|
|
|
class _DummyProcessor:
|
|
tokenizer = _DummyTokenizer()
|
|
|
|
|
|
class _DummyModel:
|
|
pass
|
|
|
|
|
|
def _install_fake_mlx(monkeypatch):
|
|
mlx_pkg = types.ModuleType("mlx")
|
|
mlx_core = types.ModuleType("mlx.core")
|
|
mlx_core.metal = _DummyMetal()
|
|
mlx_core.set_wired_limit = _DummyMX.set_wired_limit
|
|
mlx_core.device_info = _DummyMX.device_info
|
|
mlx_pkg.core = mlx_core
|
|
monkeypatch.setitem(sys.modules, "mlx", mlx_pkg)
|
|
monkeypatch.setitem(sys.modules, "mlx.core", mlx_core)
|
|
|
|
|
|
def _install_fake_fast_mlx(monkeypatch, calls):
|
|
class _FastMLXModel:
|
|
@staticmethod
|
|
def from_pretrained(*args, **kwargs):
|
|
calls.append((args, kwargs))
|
|
if kwargs["text_only"] is False:
|
|
return _DummyModel(), _DummyProcessor()
|
|
return _DummyModel(), _DummyTokenizer()
|
|
|
|
unsloth_zoo_pkg = types.ModuleType("unsloth_zoo")
|
|
mlx_pkg = types.ModuleType("unsloth_zoo.mlx")
|
|
mlx_loader = types.ModuleType("unsloth_zoo.mlx.loader")
|
|
mlx_loader.FastMLXModel = _FastMLXModel
|
|
unsloth_zoo_pkg.mlx = mlx_pkg
|
|
mlx_pkg.loader = mlx_loader
|
|
monkeypatch.setitem(sys.modules, "unsloth_zoo", unsloth_zoo_pkg)
|
|
monkeypatch.setitem(sys.modules, "unsloth_zoo.mlx", mlx_pkg)
|
|
monkeypatch.setitem(sys.modules, "unsloth_zoo.mlx.loader", mlx_loader)
|
|
|
|
|
|
def test_mlx_inference_text_load_forwards_studio_settings(monkeypatch):
|
|
_install_fake_mlx(monkeypatch)
|
|
calls = []
|
|
_install_fake_fast_mlx(monkeypatch, calls)
|
|
|
|
from core.inference.mlx_inference import MLXInferenceBackend
|
|
|
|
backend = MLXInferenceBackend()
|
|
config = SimpleNamespace(identifier = "fake/text", is_vision = False, is_lora = False)
|
|
|
|
assert backend.load_model(
|
|
config,
|
|
max_seq_length = 4096,
|
|
load_in_4bit = False,
|
|
hf_token = "hf-token",
|
|
trust_remote_code = True,
|
|
dtype = "float16",
|
|
)
|
|
|
|
assert calls == [
|
|
(
|
|
("fake/text",),
|
|
{
|
|
"max_seq_length": 4096,
|
|
"dtype": "float16",
|
|
"load_in_4bit": False,
|
|
"token": "hf-token",
|
|
"trust_remote_code": True,
|
|
"text_only": True,
|
|
},
|
|
)
|
|
]
|
|
assert backend._is_vlm is False
|
|
assert isinstance(backend._tokenizer, _DummyTokenizer)
|
|
|
|
|
|
def test_mlx_inference_vlm_lora_uses_unsloth_loader_without_native_adapter_rewrite(
|
|
monkeypatch,
|
|
tmp_path,
|
|
):
|
|
_install_fake_mlx(monkeypatch)
|
|
calls = []
|
|
_install_fake_fast_mlx(monkeypatch, calls)
|
|
|
|
def _native_vlm_load(*_args, **_kwargs):
|
|
raise AssertionError("Studio MLX VLM inference must use FastMLXModel")
|
|
|
|
mlx_vlm = types.ModuleType("mlx_vlm")
|
|
mlx_vlm.load = _native_vlm_load
|
|
monkeypatch.setitem(sys.modules, "mlx_vlm", mlx_vlm)
|
|
|
|
adapter_dir = tmp_path / "adapter"
|
|
adapter_dir.mkdir()
|
|
cfg_path = adapter_dir / "adapter_config.json"
|
|
original_cfg = '{"base_model_name_or_path": "fake/base", "rank": 8}\n'
|
|
cfg_path.write_text(original_cfg)
|
|
|
|
from core.inference.mlx_inference import MLXInferenceBackend
|
|
|
|
backend = MLXInferenceBackend()
|
|
config = SimpleNamespace(
|
|
identifier = str(adapter_dir),
|
|
is_vision = True,
|
|
is_lora = True,
|
|
base_model = "fake/base",
|
|
)
|
|
|
|
assert backend.load_model(
|
|
config,
|
|
max_seq_length = 8192,
|
|
load_in_4bit = True,
|
|
hf_token = "hf-token",
|
|
trust_remote_code = True,
|
|
)
|
|
|
|
assert calls == [
|
|
(
|
|
(str(adapter_dir),),
|
|
{
|
|
"max_seq_length": 8192,
|
|
"dtype": None,
|
|
"load_in_4bit": True,
|
|
"token": "hf-token",
|
|
"trust_remote_code": True,
|
|
"text_only": False,
|
|
},
|
|
)
|
|
]
|
|
assert cfg_path.read_text() == original_cfg
|
|
assert backend._is_vlm is True
|
|
assert isinstance(backend._processor, _DummyProcessor)
|
|
assert isinstance(backend._tokenizer, _DummyTokenizer)
|