File size: 2,821 Bytes
c34ff1f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 | """Unit tests for the external-checkpoint auto-conversion dispatch."""
from unittest.mock import patch
import pytest
from speculators.convert.entrypoints import maybe_convert_external_checkpoint
class TestMaybeConvertExternalCheckpoint:
@patch("speculators.convert.entrypoints.PretrainedConfig.get_config_dict")
def test_speculators_checkpoint_passes_through(self, mock_cfg):
mock_cfg.return_value = ({"speculators_model_type": "dflash"}, None)
assert maybe_convert_external_checkpoint("some/speculators-model") == (
"some/speculators-model"
)
@patch("speculators.convert.entrypoints.convert_model")
@patch("speculators.convert.entrypoints.PretrainedConfig.get_config_dict")
def test_external_dflash_converts(self, mock_cfg, mock_convert):
mock_cfg.return_value = ({"dflash_config": {}, "model_type": "qwen3"}, None)
out = maybe_convert_external_checkpoint(
"z-lab/Qwen3-8B-DFlash-b16",
verifier="Qwen/Qwen3-8B",
output_path="/tmp/out",
)
assert out == "/tmp/out"
mock_convert.assert_called_once()
assert mock_convert.call_args.kwargs["algorithm"] == "dflash"
assert mock_convert.call_args.kwargs["verifier"] == "Qwen/Qwen3-8B"
@patch("speculators.convert.entrypoints.convert_model")
@patch("speculators.convert.entrypoints.PretrainedConfig.get_config_dict")
def test_external_domino_converts(self, mock_cfg, mock_convert):
mock_cfg.return_value = (
{
"architectures": ["DominoDraftModel"],
"dflash_config": {"projector_type": "domino"},
},
None,
)
out = maybe_convert_external_checkpoint(
"some/domino-model",
verifier="Qwen/Qwen3-8B",
output_path="/tmp/out",
)
assert out == "/tmp/out"
mock_convert.assert_called_once()
assert mock_convert.call_args.kwargs["algorithm"] == "domino"
assert mock_convert.call_args.kwargs["verifier"] == "Qwen/Qwen3-8B"
@patch("speculators.convert.entrypoints.PretrainedConfig.get_config_dict")
def test_external_without_verifier_raises(self, mock_cfg):
mock_cfg.return_value = ({"architectures": ["DFlashDraftModel"]}, None)
with pytest.raises(ValueError, match="requires a verifier"):
maybe_convert_external_checkpoint("z-lab/Qwen3-8B-DFlash-b16")
@patch("speculators.convert.entrypoints.PretrainedConfig.get_config_dict")
def test_unrecognized_format_raises(self, mock_cfg):
mock_cfg.return_value = ({"model_type": "qwen3"}, None)
with pytest.raises(NotImplementedError, match="unrecognized external"):
maybe_convert_external_checkpoint("some/model", verifier="Qwen/Qwen3-8B")
|